add in-transaction for upgrade

This commit is contained in:
Bogdan 2023-04-03 20:55:12 +04:00
parent 8cf50c58d7
commit 241b30a710
2 changed files with 28 additions and 18 deletions

View File

@ -36,7 +36,18 @@ class Command:
async def init(self):
await Migrate.init(self.tortoise_config, self.app, self.location)
async def upgrade(self):
async def _upgrade(self, conn, version_file):
file_path = Path(Migrate.migrate_location, version_file)
m = import_py_file(file_path)
upgrade = getattr(m, "upgrade")
await conn.execute_script(await upgrade(conn))
await Aerich.create(
version=version_file,
app=self.app,
content=get_models_describe(self.app),
)
async def upgrade(self, run_in_transaction: bool):
migrated = []
for version_file in Migrate.get_all_version_files():
try:
@ -44,18 +55,13 @@ class Command:
except OperationalError:
exists = False
if not exists:
async with in_transaction(
get_app_connection_name(self.tortoise_config, self.app)
) as conn:
file_path = Path(Migrate.migrate_location, version_file)
m = import_py_file(file_path)
upgrade = getattr(m, "upgrade")
await conn.execute_script(await upgrade(conn))
await Aerich.create(
version=version_file,
app=self.app,
content=get_models_describe(self.app),
)
app_conn_name = get_app_connection_name(self.tortoise_config, self.app)
if run_in_transaction:
async with in_transaction(app_conn_name) as conn:
await self._upgrade(conn, version_file)
else:
app_conn = get_app_connection(self.tortoise_config, self.app)
await self._upgrade(app_conn, version_file)
migrated.append(version_file)
return migrated

View File

@ -90,11 +90,17 @@ async def migrate(ctx: Context, name):
@cli.command(help="Upgrade to specified version.")
@click.option(
"--in-transaction",
default=True,
type=bool,
help="Make migrations in transaction or not. Can be helpful for large migrations or creating concurrent indexes.",
)
@click.pass_context
@coro
async def upgrade(ctx: Context):
async def upgrade(ctx: Context, in_transaction: bool):
command = ctx.obj["command"]
migrated = await command.upgrade()
migrated = await command.upgrade(run_in_transaction=in_transaction)
if not migrated:
click.secho("No upgrade items found", fg=Color.yellow)
else:
@ -233,9 +239,7 @@ async def init_db(ctx: Context, safe: bool):
click.secho(f"Success create app migrate location {dirname}", fg=Color.green)
click.secho(f'Success generate schema for app "{app}"', fg=Color.green)
except FileExistsError:
return click.secho(
f"Inited {app} already, or delete {dirname} and try again.", fg=Color.yellow
)
return click.secho(f"Inited {app} already, or delete {dirname} and try again.", fg=Color.yellow)
@cli.command(help="Introspects the database tables to standard output as TortoiseORM model.")