database.py

# Database connection and session management
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from contextlib import contextmanager
from typing import Generator
import os

from models import Base
from config import config


class Database:
    """Database connection manager."""
    
    def __init__(self, database_url: str = None):
        """Initialize database connection."""
        self.database_url = database_url or config.get_database_url()
        self.engine = create_engine(
            self.database_url, 
            echo=False,
            connect_args={"check_same_thread": False}  # SQLite specific
        )
        self.SessionLocal = sessionmaker(
            autocommit=False, 
            autoflush=False, 
            bind=self.engine
        )
    
    def create_tables(self):
        """Create all database tables."""
        Base.metadata.create_all(bind=self.engine)
    
    def drop_tables(self):
        """Drop all database tables."""
        Base.metadata.drop_all(bind=self.engine)
    
    @contextmanager
    def get_session(self) -> Generator[Session, None, None]:
        """Get a database session with automatic cleanup."""
        session = self.SessionLocal()
        try:
            yield session
            session.commit()
        except Exception:
            session.rollback()
            raise
        finally:
            session.close()
    
    def get_session_direct(self) -> Session:
        """Get a database session (caller must manage lifecycle)."""
        return self.SessionLocal()


# Global database instance
db = Database()

# Auto-initialize tables on import
db.create_tables()


def init_database():
    """Initialize the database and create tables."""
    db.create_tables()
    return db


def get_db_session() -> Session:
    """Get a database session for use in Streamlit."""
    return db.get_session_direct()