pub/aerich
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add in-transaction for upgrade

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

32
aerich/__init__.py
View File

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

14
aerich/cli.py
View File

@@ -90,11 +90,17 @@ async def migrate(ctx: Context, name):
@cli.command(help="Upgrade to specified version.") @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 @click.pass_context
@coro @coro
async def upgrade(ctx: Context): async def upgrade(ctx: Context, in_transaction: bool):
command = ctx.obj["command"] command = ctx.obj["command"]
migrated = await command.upgrade() migrated = await command.upgrade(run_in_transaction=in_transaction)
if not migrated: if not migrated:
click.secho("No upgrade items found", fg=Color.yellow) click.secho("No upgrade items found", fg=Color.yellow)
else: 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 create app migrate location {dirname}", fg=Color.green)
click.secho(f'Success generate schema for app "{app}"', fg=Color.green) click.secho(f'Success generate schema for app "{app}"', fg=Color.green)
except FileExistsError: except FileExistsError:
return click.secho( return click.secho(f"Inited {app} already, or delete {dirname} and try again.", fg=Color.yellow)
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.") @cli.command(help="Introspects the database tables to standard output as TortoiseORM model.")