From 241b30a7102a31dbe7977d38bc79d6fbf3bdca96 Mon Sep 17 00:00:00 2001
From: Bogdan <evstrat.bg@gmail.com>
Date: Mon, 3 Apr 2023 20:55:12 +0400
Subject: [PATCH] add in-transaction for upgrade

---
 aerich/__init__.py | 32 +++++++++++++++++++-------------
 aerich/cli.py      | 14 +++++++++-----
 2 files changed, 28 insertions(+), 18 deletions(-)

diff --git a/aerich/__init__.py b/aerich/__init__.py
index 75f4e05..2109b94 100644
--- a/aerich/__init__.py
+++ b/aerich/__init__.py
@@ -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
 
diff --git a/aerich/cli.py b/aerich/cli.py
index 072e547..e09dabe 100644
--- a/aerich/cli.py
+++ b/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.")