db threadsafe session handling

database.py

@@ -1,9 +1,9 @@
-# database.py
 import os
 from sqlalchemy import create_engine, Column, Integer, String, Text, Boolean, DateTime, func
 from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import sessionmaker
+from sqlalchemy.orm import sessionmaker, scoped_session
 from contextlib import contextmanager
+from sqlalchemy.exc import SQLAlchemyError
 from utils.utils import get_secret
 import random
 from better_profanity import profanity
@@ -31,7 +31,9 @@ except Exception as e:
     logger.error(f"Error connecting to the database: {str(e)}")
     raise e
 
-SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+# Use scoped session to handle thread-local sessions
+session_factory = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+SessionLocal = scoped_session(session_factory)
 
 Base = declarative_base()
 
@@ -92,7 +94,7 @@ def get_db_session():
         db.rollback()
         raise e
     finally:
-        db.close()
+        SessionLocal.remove()
 
 def create_chatbot(name, custom_instruction):
     with get_db_session() as db:
@@ -107,11 +109,8 @@ def get_chatbot(chatbot_id):
     with get_db_session() as db:
         return db.query(Chatbot).filter(Chatbot.chatbot_id == chatbot_id, Chatbot.is_active == True).first()
 
-
-### These functions need to be rewritten to use the get_db_session context manager
 def update_chatbot(chatbot_id, name=None, custom_instruction=None, is_active=None):
-    db = SessionLocal()
-    try:
+    with get_db_session() as db:
         chatbot = db.query(Chatbot).filter(Chatbot.chatbot_id == chatbot_id).first()
         if chatbot:
             if name:
@@ -123,32 +122,16 @@ def update_chatbot(chatbot_id, name=None, custom_instruction=None, is_active=Non
             db.commit()
             db.refresh(chatbot)
         return chatbot
-    except SQLAlchemyError as e:
-        db.rollback()
-        raise e
-    finally:
-        db.close()
 
 def delete_chatbot(chatbot_id):
-    db = SessionLocal()
-    try:
+    with get_db_session() as db:
         chatbot = db.query(Chatbot).filter(Chatbot.chatbot_id == chatbot_id).first()
         if chatbot:
             db.delete(chatbot)
             db.commit()
             return True
         return False
-    except SQLAlchemyError as e:
-        db.rollback()
-        raise e
-    finally:
-        db.close()
 
 def get_all_chatbots():
-    db = SessionLocal()
-    try:
+    with get_db_session() as db:
         return db.query(Chatbot).all()
-    except SQLAlchemyError as e:
-        raise e
-    finally:
-        db.close()