From c27a18cc174bc50ff195aceb208c47da08daf1a0 Mon Sep 17 00:00:00 2001
From: magowiz <magowiz@gmail.com>
Date: Tue, 7 Dec 2021 16:58:12 +0100
Subject: [PATCH] completed chain, handling exception

---
 mobilizon_reshare/cli/__init__.py | 19 ++++++++++++++--
 mobilizon_reshare/storage/db.py   | 38 +++++++++++++++++++++++++++----
 2 files changed, 50 insertions(+), 7 deletions(-)

diff --git a/mobilizon_reshare/cli/__init__.py b/mobilizon_reshare/cli/__init__.py
index 2b6bc02..dc17e76 100644
--- a/mobilizon_reshare/cli/__init__.py
+++ b/mobilizon_reshare/cli/__init__.py
@@ -3,6 +3,7 @@ import logging
 import traceback
 from logging.config import dictConfig
 from pathlib import Path
+import sys
 
 from mobilizon_reshare.config.config import get_settings
 from mobilizon_reshare.storage.db import tear_down, MoReDB
@@ -20,11 +21,25 @@ async def init(settings_file):
     dictConfig(settings["logging"])
     db_path = Path(settings.db_path)
     db = MoReDB(db_path)
-    await db.setup()
+    db_setup = asyncio.create_task(db.setup())
+    _, _ = await asyncio.wait({db_setup},
+                              return_when=asyncio.FIRST_EXCEPTION)
+    if db_setup.exception():
+        logging.critical("exception during db setup")
+        raise db_setup.exception()
 
 
 async def _safe_execution(f, settings_file):
-    await init(settings_file)
+    init_task = asyncio.create_task(init(settings_file))
+    _, _ = await asyncio.wait({init_task},
+                              return_when=asyncio.FIRST_EXCEPTION)
+    if init_task.exception():
+        logging.critical("exception during init")
+        # raise init_task.exception()
+        # sys.exit(1)
+        loop = asyncio.get_event_loop()
+        loop.stop()
+
     return_code = 1
     try:
         return_code = await f()
diff --git a/mobilizon_reshare/storage/db.py b/mobilizon_reshare/storage/db.py
index b5ce121..6c048df 100644
--- a/mobilizon_reshare/storage/db.py
+++ b/mobilizon_reshare/storage/db.py
@@ -1,5 +1,5 @@
 import logging
-import sys
+import asyncio
 from pathlib import Path
 from tortoise import Tortoise
 from aerich import Command
@@ -37,8 +37,33 @@ TORTOISE_ORM = {
 }
 
 
+async def shutdown(loop):
+    """shutdown method"""
+    logging.critical("SHUTDOWN CALLED")
+    logging.info("closing database connections")
+    tasks = [t for t in asyncio.all_tasks() if t is asyncio.current_task()]
+    _ = [task.cancel() for task in tasks]
+    logging.info("Cancelling %i tasks", len(tasks))
+    await asyncio.gather(*tasks, return_exceptions=True)
+    logging.info("flushing metrics")
+    loop.stop()
+
+
+async def handle_exception(loop, context):
+    """exception handler"""
+    logging.critical("HANDLER CALLED")
+    msg = context.get("exception", context["message"])
+    logging.critical("Caught exception: %s", msg)
+    logging.info("shutting down")
+    await shutdown(loop)
+    await tear_down()
+
+
 class MoReDB:
     def __init__(self, path: Path):
+        loop = asyncio.get_event_loop()
+        loop.set_exception_handler(handle_exception)
+
         self.path = path
         # TODO: Check if DB is openable/"queriable"
         self.is_init = self.path.exists() and (not self.path.is_dir())
@@ -55,12 +80,15 @@ class MoReDB:
             await command.upgrade()
         except FileNotFoundError:
             logging.critical("aerich configuration not found, fatal error")
-            # raise
-            sys.exit(1)
-
+            raise
 
     async def setup(self):
-        await self._implement_db_changes()
+        implement_db_changes = asyncio.create_task(self._implement_db_changes())
+        _, _ = await asyncio.wait({implement_db_changes},
+                                  return_when=asyncio.FIRST_EXCEPTION)
+        if implement_db_changes.exception():
+            logging.critical("exception during aerich init")
+            raise implement_db_changes.exception()
         await Tortoise.init(
             db_url=get_db_url(),
             modules={