From 9c3ba7e2732cc7ba8267034307b96ce17de6d4c5 Mon Sep 17 00:00:00 2001
From: Waket Zheng <waketzheng@gmail.com>
Date: Thu, 6 Mar 2025 13:39:56 +0800
Subject: [PATCH] fix: `aerich init-db` process is suspended (#435)

---
 CHANGELOG.md                 |  7 +++++++
 aerich/cli.py                | 16 ++++++++++++++++
 tests/test_sqlite_migrate.py | 23 +++++++++++++++++++----
 3 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e85bac7..eab5ca8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,13 @@
 
 ## 0.8
 
+### [0.8.3]**(Unreleased)**
+
+#### Fixed
+- fix: `aerich init-db` process is suspended. ([#435])
+
+[#435]: https://github.com/tortoise/aerich/pull/435
+
 ### [0.8.2](../../releases/tag/v0.8.2) - 2025-02-28
 
 #### Added
diff --git a/aerich/cli.py b/aerich/cli.py
index feafb6a..54b554a 100644
--- a/aerich/cli.py
+++ b/aerich/cli.py
@@ -19,6 +19,22 @@ CONFIG_DEFAULT_VALUES = {
 }
 
 
+def _patch_context_to_close_tortoise_connections_when_exit() -> None:
+    from tortoise import Tortoise, connections
+
+    origin_aexit = Context.__aexit__
+
+    async def aexit(*args, **kw) -> None:
+        await origin_aexit(*args, **kw)
+        if Tortoise._inited:
+            await connections.close_all()
+
+    Context.__aexit__ = aexit  # type:ignore[method-assign]
+
+
+_patch_context_to_close_tortoise_connections_when_exit()
+
+
 @click.group(context_settings={"help_option_names": ["-h", "--help"]})
 @click.version_option(__version__, "-V", "--version")
 @click.option(
diff --git a/tests/test_sqlite_migrate.py b/tests/test_sqlite_migrate.py
index 91fb79f..bbe6fb7 100644
--- a/tests/test_sqlite_migrate.py
+++ b/tests/test_sqlite_migrate.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import contextlib
 import os
+import platform
 import shlex
 import shutil
 import subprocess
@@ -12,11 +13,16 @@ from pathlib import Path
 from tests._utils import Dialect, chdir, copy_files
 
 
-def run_aerich(cmd: str) -> None:
-    with contextlib.suppress(subprocess.TimeoutExpired):
-        if not cmd.startswith("aerich") and not cmd.startswith("poetry"):
+def run_aerich(cmd: str) -> subprocess.CompletedProcess | None:
+    if not cmd.startswith("poetry") and not cmd.startswith("python"):
+        if not cmd.startswith("aerich"):
             cmd = "aerich " + cmd
-        subprocess.run(shlex.split(cmd), timeout=2)
+        if platform.system() == "Windows":
+            cmd = "python -m " + cmd
+    r = None
+    with contextlib.suppress(subprocess.TimeoutExpired):
+        r = subprocess.run(shlex.split(cmd), timeout=2)
+    return r
 
 
 def run_shell(cmd: str) -> subprocess.CompletedProcess:
@@ -43,6 +49,15 @@ def prepare_sqlite_project(tmp_path: Path) -> Generator[tuple[Path, str]]:
         yield models_py, models_py.read_text("utf-8")
 
 
+def test_close_tortoise_connections_patch(tmp_path: Path) -> None:
+    if not Dialect.is_sqlite():
+        return
+    with prepare_sqlite_project(tmp_path) as (models_py, models_text):
+        run_aerich("aerich init -t settings.TORTOISE_ORM")
+        r = run_aerich("aerich init-db")
+        assert r is not None
+
+
 def test_sqlite_migrate_alter_indexed_unique(tmp_path: Path) -> None:
     if not Dialect.is_sqlite():
         return