From b6d209926ff70865e79fbe2538302339dc92223c Mon Sep 17 00:00:00 2001
From: Brian Wiborg <me@brianwib.org>
Date: Sat, 27 Sep 2025 17:19:37 +0200
Subject: [PATCH] =?UTF-8?q?=F0=9F=A9=B9=20Ensure=20ORM=20initialization=20?=
 =?UTF-8?q?and=20tear-down?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/ohmyapi/cli.py | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/ohmyapi/cli.py b/src/ohmyapi/cli.py
index 38d8fcc..58b4f11 100644
--- a/src/ohmyapi/cli.py
+++ b/src/ohmyapi/cli.py
@@ -1,4 +1,5 @@
 import asyncio
+import atexit
 import importlib
 import sys
 import typer
@@ -45,6 +46,28 @@ def shell(root: str = "."):
     project_path = Path(root).resolve()
     project = runtime.Project(project_path)
 
+    # Ensure the ORM is shutdown
+    async def close_project():
+        try:
+            await project.close_orm()
+            print("Tortoise ORM closed successfully.")
+        except Exception as e:
+            print(f"Error closing ORM: {e}")
+
+    def cleanup():
+        loop = None
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            pass
+        if loop and loop.is_running():
+            asyncio.create_task(close_project())
+        else:
+            asyncio.run(close_project())
+
+    # Ensure the ORM is initialized
+    asyncio.run(project.init_orm())
+
     try:
         from IPython import start_ipython
         shell_vars = {
@@ -56,11 +79,13 @@ def shell(root: str = "."):
         c.TerminalInteractiveShell.banner2 = banner.format(**{
             "project_name": f"{f'{project.settings.PROJECT_NAME} ' if getattr(project.settings, 'PROJECT_NAME', '') else ''}[{Path(project_path).resolve()}]",
         })
+        atexit.register(cleanup)
         start_ipython(argv=[], user_ns=shell_vars, config=c)
     except ImportError:
         typer.echo("IPython is not installed. Falling back to built-in Python shell.")
         import code
-        code.interact(local={"settings": project.settings})
+        atexit.register(cleanup)
+        code.interact(local={"p": project})
 
 
 @app.command()