refactoring db calls

database.py

@@ -3,6 +3,7 @@ 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 contextlib import contextmanager
 from utils.utils import get_secret
 import random
 from better_profanity import profanity
@@ -82,30 +83,32 @@ def generate_chatbot_id():
 
     return f"{random.choice(adjectives)}-{random.choice(nouns)}-{random.randint(100, 999)}"
 
-def create_chatbot(name, custom_instruction):
+@contextmanager
+def get_db_session():
     db = SessionLocal()
     try:
+        yield db
+    except SQLAlchemyError as e:
+        db.rollback()
+        raise e
+    finally:
+        db.close()
+
+def create_chatbot(name, custom_instruction):
+    with get_db_session() as db:
         chatbot_id = generate_chatbot_id()
         new_chatbot = Chatbot(chatbot_id=chatbot_id, name=name, custom_instruction=custom_instruction)
         db.add(new_chatbot)
         db.commit()
         db.refresh(new_chatbot)
         return new_chatbot
-    except SQLAlchemyError as e:
-        db.rollback()
-        raise e
-    finally:
-        db.close()
 
 def get_chatbot(chatbot_id):
-    db = SessionLocal()
-    try:
+    with get_db_session() as db:
         return db.query(Chatbot).filter(Chatbot.chatbot_id == chatbot_id, Chatbot.is_active == True).first()
-    except SQLAlchemyError as e:
-        raise e
-    finally:
-        db.close()
 
+
+### 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: