add in-transaction for upgrade

aerich/__init__.py

@@ -36,17 +36,7 @@ class Command:
     async def init(self):
         await Migrate.init(self.tortoise_config, self.app, self.location)
 
-    async def upgrade(self):
-        migrated = []
-        for version_file in Migrate.get_all_version_files():
-            try:
-                exists = await Aerich.exists(version=version_file, app=self.app)
-            except OperationalError:
-                exists = False
-            if not exists:
-                async with in_transaction(
-                    get_app_connection_name(self.tortoise_config, self.app)
-                ) as conn:
+    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")
@@ -56,6 +46,22 @@ class Command:
             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:
+                exists = await Aerich.exists(version=version_file, app=self.app)
+            except OperationalError:
+                exists = False
+            if not exists:
+                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
 

aerich/cli.py

@@ -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.")