database.py

# 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 contextlib import contextmanager
from utils.utils import get_secret
import random
from better_profanity import profanity
import logging

# Database connection
DB_USER = get_secret("DB_USER")
DB_PASSWORD = get_secret("DB_PASSWORD")
DB_HOST = get_secret("DB_HOST")
DB_NAME = get_secret("DB_NAME")
DB_PORT = get_secret("DB_PORT", "3306")  

DATABASE_URL = f"mysql+pymysql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

try:
    logger.info("Connecting to database at %s", DB_HOST)
    engine = create_engine(DATABASE_URL, connect_args={"connect_timeout": 10})
    # Test the connection
    with engine.connect() as connection:
        logger.info("Successfully connected to the database")
except Exception as e:
    logger.error(f"Error connecting to the database: {str(e)}")
    raise e

SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

Base = declarative_base()

class Chatbot(Base):
    __tablename__ = "chatbots"

    id = Column(Integer, primary_key=True, index=True)
    chatbot_id = Column(String(50), unique=True, index=True, nullable=False)
    name = Column(String(100), nullable=False)
    custom_instruction = Column(Text, nullable=False)
    is_active = Column(Boolean, default=True) 
    created_at = Column(DateTime(timezone=True), server_default=func.now())
    updated_at = Column(DateTime(timezone=True), onupdate=func.now())

# Create tables
Base.metadata.create_all(bind=engine)

# Load German profanity words
german_profanity = set()
with open('german_profanity.txt', 'r', encoding='utf-8') as f:
    german_profanity = set(word.strip().lower() for word in f)

profanity.add_censor_words(german_profanity)

def filter_profanity(text):
    return profanity.censor(text)

def generate_chatbot_id():
    adjectives = [
        'happy', 'clever', 'bright', 'shiny', 'fluffy', 'gentle', 'brave', 'calm',
        'kind', 'joyful', 'radiant', 'sparkling', 'cheerful', 'gracious', 'elegant', 
        'vivacious', 'serene', 'vibrant', 'splendid', 'charismatic', 'delightful', 
        'blissful', 'generous', 'charming', 'dazzling', 'glowing', 'harmonious', 
        'jovial', 'luminous', 'majestic', 'mellow', 'noble', 'optimistic', 
        'passionate', 'playful', 'resilient', 'spirited', 'tranquil', 'upbeat', 
        'valiant', 'whimsical', 'witty', 'zealous', 'admirable', 'affectionate', 
        'brilliant', 'courteous', 'devoted', 'ecstatic', 'faithful', 'gleeful'
    ]

    nouns = [
        'elephant', 'penguin', 'lion', 'dolphin', 'koala', 'panda', 'tiger', 'whale',
        'butterfly', 'eagle', 'peacock', 'unicorn', 'phoenix', 'otter', 'swan', 
        'chameleon', 'ladybug', 'flamingo', 'puppy', 'kitten', 'fawn', 'hummingbird', 
        'koala', 'firefly', 'bunny', 'goldfish', 'puffin', 'orca', 'red_panda', 'turtle', 
        'parrot', 'owl', 'seahorse', 'hedgehog', 'sloth', 'duckling', 'starfish', 
        'gazelle', 'panther', 'robin', 'seal', 'lynx', 'jellyfish', 'gecko', 
        'kangaroo', 'lemur', 'meerkat', 'platypus', 'quokka', 'squirrel', 'toucan'
    ]

    return f"{random.choice(adjectives)}-{random.choice(nouns)}-{random.randint(100, 999)}"

@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

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:
        chatbot = db.query(Chatbot).filter(Chatbot.chatbot_id == chatbot_id).first()
        if chatbot:
            if name:
                chatbot.name = name
            if custom_instruction:
                chatbot.custom_instruction = custom_instruction
            if is_active is not None:
                chatbot.is_active = is_active
            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:
        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:
        return db.query(Chatbot).all()
    except SQLAlchemyError as e:
        raise e
    finally:
        db.close()