From 5bdaa32a9e66f09fbbcadf3d4d9a475a67361492 Mon Sep 17 00:00:00 2001
From: Patrick Schneeweis <pschneeweis@protonmail.com>
Date: Fri, 24 Jul 2020 09:17:19 -0400
Subject: [PATCH] Decorator for closing db connection after func run

---
 aerich/cli.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/aerich/cli.py b/aerich/cli.py
index 4feb041..4da07c1 100644
--- a/aerich/cli.py
+++ b/aerich/cli.py
@@ -1,3 +1,4 @@
+import functools
 import json
 import os
 import sys
@@ -27,6 +28,16 @@ class Color(str, Enum):
 parser = ConfigParser()
 
 
+def close_db(func):
+    @functools.wraps(func)
+    async def close_db_inner(*args, **kwargs):
+        result = await func(*args, **kwargs)
+        await Tortoise.close_connections()
+        return result
+
+    return close_db_inner
+
+
 @click.group(context_settings={"help_option_names": ["-h", "--help"]})
 @click.version_option(__version__, "-V", "--version")
 @click.option(