Merge pull request #296 from evstratbg/migrate-transaction

add in-transaction for upgrade

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## 0.7
 
+### 0.7.3
+
+- Added `-i` and `--in-transaction` options to `aerich migrate` command. (#296)
+
 ### 0.7.2
 
 - Support virtual fields.

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(
+                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)
         return migrated
 

aerich/cli.py

@@ -90,11 +90,18 @@ async def migrate(ctx: Context, name):
 
 
 @cli.command(help="Upgrade to specified version.")
+@click.option(
+    "--in-transaction",
+    "-i",
+    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:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "aerich"
-version = "0.7.2"
+version = "0.7.3"
 description = "A database migrations tool for Tortoise ORM."
 authors = ["long2ice <long2ice@gmail.com>"]
 license = "Apache-2.0"