From 2a7a246f1677199e83185f4bc0b2b5597e03126c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sara=20Hincapi=C3=A9=20M?=
 <43832784+shincap8@users.noreply.github.com>
Date: Fri, 13 Sep 2024 17:34:24 -0500
Subject: [PATCH] add refresher middleware to close connections after a call
 and add pool_pre_ping when creating the connection engine (#304)

---
 backend/app/infrastructure/connection.py |  5 ++++-
 backend/app/main.py                      | 20 ++++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/backend/app/infrastructure/connection.py b/backend/app/infrastructure/connection.py
index d0c127588..7d8b30bd1 100644
--- a/backend/app/infrastructure/connection.py
+++ b/backend/app/infrastructure/connection.py
@@ -8,10 +8,13 @@
 
 import os
 
+import uuid
 from dotenv import load_dotenv
 from sqlalchemy import MetaData, create_engine
 from sqlalchemy.orm import sessionmaker
 
+from app.infrastructure.utils.singleton import Singleton
+
 
 load_dotenv()
 
@@ -28,7 +31,7 @@
 
 class Connection:
     def __init__(self) -> None:
-        self.engine = create_engine(CONNECTION_URI, echo=False)
+        self.engine = create_engine(CONNECTION_URI, echo=False, pool_pre_ping=True)
         self.Session = sessionmaker(bind=self.engine, expire_on_commit=False)
         self.session = self.Session()
         self.metadata = MetaData()
diff --git a/backend/app/main.py b/backend/app/main.py
index 4c5e6043e..4ce5d98ce 100644
--- a/backend/app/main.py
+++ b/backend/app/main.py
@@ -25,6 +25,7 @@
     user,
 )
 from app.api.endpoints.builder_and_evaluation import evaluation
+from app.infrastructure.connection import Connection
 
 
 load_dotenv()
@@ -72,3 +73,22 @@ def read_root():
     historical_data.router, prefix="/historical_data", tags=["historical_data"]
 )
 app.include_router(round.router, prefix="/round", tags=["round"])
+
+
+class RefreshSessionMiddleware:
+    """
+    Refreshes SQLAlchemy session on each call.
+    There may be issues since close_session is blocking and it is ran in an async function.
+    Possibly blocking the main asyncio loop.
+    Keep in mind that middlewares run on the main asyncio thread.
+    """
+
+    def __init__(self, app):
+        self.app = app
+
+    async def __call__(self, scope, receive, send):
+        await self.app(scope, receive, send)
+        Connection().close_session()
+
+
+app.add_middleware(RefreshSessionMiddleware)