Refactor database engine initialization to include connection pooling…

… settings for improved performance
satishsurath · Feb 4, 2025 · 754b89c · 754b89c
1 parent 70245f7
commit 754b89c
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 4 deletions.
diff --git a/app.py b/app.py
@@ -27,7 +27,12 @@
 
 
 DB_URL = os.getenv("DATABASE_URL", "postgresql://user:pass@localhost:5432/mydb")
-engine = create_engine(DB_URL, echo=False)
+#engine = create_engine(DB_URL, echo=False)
+engine = create_engine(
+    DB_URL, 
+    echo=False,
+    pool_pre_ping=True,
+    pool_recycle=1800)  # 30 minutes
 SessionLocal = sessionmaker(bind=engine)
 
 app = Flask(__name__)
@@ -48,6 +53,7 @@
 download_statuses = {}
 summarize_v2_statuses = {}
 
+
 # Define a decorator to require a specific role
 def require_role(allowed_roles):
     """
@@ -156,11 +162,14 @@ def api_channel_start():
     }
 
     def run_download():
+        session = SessionLocal()
         try:
             # This function now ensures that for every video,
             # a row in video_folders(folder_name=<channel_id>, video_id=...) is inserted
             download_channel_transcripts(channel_url, download_statuses[task_id])
             download_statuses[task_id]["status"] = "completed"
+        finally:
+            session.close()
         except Exception as e:
             logger.error(f"Error in channel download: {e}")
             download_statuses[task_id]["status"] = "failed"

diff --git a/auth_utils.py b/auth_utils.py
@@ -21,7 +21,11 @@
 CLOUDFLARE_AUD_TAG = os.getenv("CLOUDFLARE_AUD_TAG")
 
 DB_URL = os.getenv("DATABASE_URL", "postgresql://user:pass@localhost:5432/mydb")
-engine = create_engine(DB_URL, echo=False)
+engine = create_engine(
+    DB_URL, 
+    echo=False,
+    pool_pre_ping=True,
+    pool_recycle=1800)
 SessionLocal = sessionmaker(bind=engine)
 
 

diff --git a/init_db.py b/init_db.py
@@ -30,7 +30,12 @@ def main():
     logger.info(f"Using DB_URL={DB_URL}")
 
     # Create engine & session
-    engine = create_engine(DB_URL, echo=True)
+    #engine = create_engine(DB_URL, echo=True)
+    engine = create_engine(
+    DB_URL, 
+    echo=False,
+    pool_pre_ping=True,
+    pool_recycle=1800)  # 30 minutes
     SessionLocal = sessionmaker(bind=engine)
 
     # Create all tables that do not exist

diff --git a/youtube_utils.py b/youtube_utils.py
@@ -16,7 +16,12 @@
 logger = logging.getLogger(__name__)
 
 DB_URL = os.getenv("DATABASE_URL", "postgresql://user:pass@localhost:5432/mydb")
-engine = create_engine(DB_URL)
+#engine = create_engine(DB_URL)
+engine = create_engine(
+    DB_URL, 
+    echo=False,
+    pool_pre_ping=True,
+    pool_recycle=1800)  # 30 minutes
 SessionLocal = sessionmaker(bind=engine)
 
 def download_channel_transcripts(channel_url, status_dict):