diff --git a/aerich/cli.py b/aerich/cli.py index 26a0164..e4d3cf7 100644 --- a/aerich/cli.py +++ b/aerich/cli.py @@ -8,7 +8,7 @@ from typing import List import click from click import Context, UsageError -from tortoise import Tortoise, generate_schema_for_client, run_async +from tortoise import Tortoise, generate_schema_for_client from tortoise.exceptions import OperationalError from tortoise.transactions import in_transaction from tortoise.utils import get_schema_sql @@ -34,7 +34,14 @@ parser = ConfigParser() def coro(f): @wraps(f) def wrapper(*args, **kwargs): - run_async(f(*args, **kwargs)) + loop = asyncio.get_event_loop() + + # Close db connections at the end of all all but the cli group function + try: + loop.run_until_complete(f(*args, **kwargs)) + finally: + if f.__name__ != "cli": + loop.run_until_complete(Tortoise.close_connections()) return wrapper