From 8e5ae800d5071fd6fa9178a9d4ba0e14e3567337 Mon Sep 17 00:00:00 2001
From: Taha Yassine Kraiem <tahayk2@gmail.com>
Date: Tue, 14 Mar 2023 15:15:48 +0100
Subject: [PATCH] feat(chalice): upgraded startup/shutdown logic

---
 api/app.py                        | 77 ++++++++++++++++---------------
 api/requirements-alerts.txt       |  2 +-
 api/requirements.txt              |  2 +-
 ee/api/chalicelib/core/tenants.py |  6 +--
 ee/api/requirements-alerts.txt    |  2 +-
 ee/api/requirements.txt           |  2 +-
 6 files changed, 48 insertions(+), 43 deletions(-)

diff --git a/api/app.py b/api/app.py
index 883cf6704..50cd7342f 100644
--- a/api/app.py
+++ b/api/app.py
@@ -1,4 +1,5 @@
 import logging
+from contextlib import asynccontextmanager
 
 from apscheduler.schedulers.asyncio import AsyncIOScheduler
 from decouple import config
@@ -14,7 +15,40 @@ from routers.crons import core_crons
 from routers.crons import core_dynamic_crons
 from routers.subs import insights, metrics, v1_api, health
 
-app = FastAPI(root_path="/api", docs_url=config("docs_url", default=""), redoc_url=config("redoc_url", default=""))
+loglevel = config("LOGLEVEL", default=logging.INFO)
+print(f">Loglevel set to: {loglevel}")
+logging.basicConfig(level=loglevel)
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup
+    logging.info(">>>>> starting up <<<<<")
+    ap_logger = logging.getLogger('apscheduler')
+    ap_logger.setLevel(loglevel)
+
+    app.schedule = AsyncIOScheduler()
+    await pg_client.init()
+    app.schedule.start()
+
+    for job in core_crons.cron_jobs + core_dynamic_crons.cron_jobs:
+        app.schedule.add_job(id=job["func"].__name__, **job)
+
+    ap_logger.info(">Scheduled jobs:")
+    for job in app.schedule.get_jobs():
+        ap_logger.info({"Name": str(job.id), "Run Frequency": str(job.trigger), "Next Run": str(job.next_run_time)})
+
+    # App listening
+    yield
+
+    # Shutdown
+    logging.info(">>>>> shutting down <<<<<")
+    app.schedule.shutdown(wait=False)
+    await pg_client.terminate()
+
+
+app = FastAPI(root_path="/api", docs_url=config("docs_url", default=""), redoc_url=config("redoc_url", default=""),
+              lifespan=lifespan)
 app.add_middleware(GZipMiddleware, minimum_size=1000)
 
 
@@ -55,38 +89,9 @@ app.include_router(health.public_app)
 app.include_router(health.app)
 app.include_router(health.app_apikey)
 
-loglevel = config("LOGLEVEL", default=logging.INFO)
-print(f">Loglevel set to: {loglevel}")
-logging.basicConfig(level=loglevel)
-ap_logger = logging.getLogger('apscheduler')
-ap_logger.setLevel(loglevel)
-app.schedule = AsyncIOScheduler()
-
-
-@app.on_event("startup")
-async def startup():
-    logging.info(">>>>> starting up <<<<<")
-    await pg_client.init()
-    app.schedule.start()
-
-    for job in core_crons.cron_jobs + core_dynamic_crons.cron_jobs:
-        app.schedule.add_job(id=job["func"].__name__, **job)
-
-    ap_logger.info(">Scheduled jobs:")
-    for job in app.schedule.get_jobs():
-        ap_logger.info({"Name": str(job.id), "Run Frequency": str(job.trigger), "Next Run": str(job.next_run_time)})
-
-
-@app.on_event("shutdown")
-async def shutdown():
-    logging.info(">>>>> shutting down <<<<<")
-    app.schedule.shutdown(wait=False)
-    await pg_client.terminate()
-
-
-@app.get('/private/shutdown', tags=["private"])
-async def stop_server():
-    logging.info("Requested shutdown")
-    await shutdown()
-    import os, signal
-    os.kill(1, signal.SIGTERM)
+# @app.get('/private/shutdown', tags=["private"])
+# async def stop_server():
+#     logging.info("Requested shutdown")
+#     await shutdown()
+#     import os, signal
+#     os.kill(1, signal.SIGTERM)
diff --git a/api/requirements-alerts.txt b/api/requirements-alerts.txt
index b208d28c2..edb644c87 100644
--- a/api/requirements-alerts.txt
+++ b/api/requirements-alerts.txt
@@ -8,7 +8,7 @@ jira==3.4.1
 
 
 
-fastapi==0.92.0
+fastapi==0.94.1
 uvicorn[standard]==0.20.0
 python-decouple==3.7
 pydantic[email]==1.10.4
diff --git a/api/requirements.txt b/api/requirements.txt
index 4a8d35090..27b95f17e 100644
--- a/api/requirements.txt
+++ b/api/requirements.txt
@@ -8,7 +8,7 @@ jira==3.4.1
 
 
 
-fastapi==0.92.0
+fastapi==0.94.1
 uvicorn[standard]==0.20.0
 python-decouple==3.7
 pydantic[email]==1.10.4
diff --git a/ee/api/chalicelib/core/tenants.py b/ee/api/chalicelib/core/tenants.py
index 30a87bd29..7ea621007 100644
--- a/ee/api/chalicelib/core/tenants.py
+++ b/ee/api/chalicelib/core/tenants.py
@@ -51,7 +51,7 @@ def get_by_api_key(api_key):
                                 WHERE tenants.api_key = %(api_key)s 
                                     AND tenants.deleted_at ISNULL
                                 LIMIT 1;""",
-            {"api_key": api_key})
+                            {"api_key": api_key})
         cur.execute(query=query)
         return helper.dict_to_camel_case(cur.fetchone())
 
@@ -94,7 +94,7 @@ def update(tenant_id, user_id, data: schemas.UpdateTenantSchema):
     return edit_client(tenant_id=tenant_id, changes=changes)
 
 
-def tenants_exists():
-    with pg_client.PostgresClient() as cur:
+def tenants_exists(use_pool=True):
+    with pg_client.PostgresClient(use_pool=use_pool) as cur:
         cur.execute(f"SELECT EXISTS(SELECT 1 FROM public.tenants)")
         return cur.fetchone()["exists"]
diff --git a/ee/api/requirements-alerts.txt b/ee/api/requirements-alerts.txt
index 250882623..6b6901ca5 100644
--- a/ee/api/requirements-alerts.txt
+++ b/ee/api/requirements-alerts.txt
@@ -8,7 +8,7 @@ jira==3.4.1
 
 
 
-fastapi==0.92.0
+fastapi==0.94.1
 uvicorn[standard]==0.20.0
 python-decouple==3.7
 pydantic[email]==1.10.4
diff --git a/ee/api/requirements.txt b/ee/api/requirements.txt
index 9ce06fe06..cad05e873 100644
--- a/ee/api/requirements.txt
+++ b/ee/api/requirements.txt
@@ -8,7 +8,7 @@ jira==3.4.1
 
 
 
-fastapi==0.92.0
+fastapi==0.94.1
 uvicorn[standard]==0.20.0
 python-decouple==3.7
 pydantic[email]==1.10.4