database.py

import os
import time
import sqlalchemy
from sqlalchemy import create_engine
from sqlalchemy.orm import declarative_base, sessionmaker
import streamlit as st
from contextlib import contextmanager
from dotenv import load_dotenv

load_dotenv()

# Load database connection details from environment variables
DB_HOST = os.environ.get("PGHOST", "localhost")
DB_PORT = os.environ.get("PGPORT", "5432")
DB_NAME = os.environ.get("PGDATABASE", "privacychatbox")
DB_USER = os.environ.get("PGUSER", "postgres")
DB_PASSWORD = os.environ.get("PGPASSWORD", "postgres")
DATABASE_URL = os.environ.get("DATABASE_URL")

# Create SQLAlchemy base class
Base = declarative_base()

# Engine and session factory
engine = None
SessionLocal = None

# Connection retry settings
MAX_RETRIES = 3
RETRY_DELAY = 1  # seconds

def init_db():
    """Initialize the database connection"""
    global engine, SessionLocal
    
    retry_count = 0
    last_error = None
    
    while retry_count < MAX_RETRIES:
        try:
            # Try to use DATABASE_URL if available
            if DATABASE_URL:
                engine = create_engine(
                    DATABASE_URL,
                    pool_pre_ping=True,  # Enable connection health checks
                    pool_recycle=3600    # Recycle connections after 1 hour
                )
            else:
                # Otherwise construct the connection string from individual components
                conn_str = f"postgresql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
                engine = create_engine(
                    conn_str,
                    pool_pre_ping=True,
                    pool_recycle=3600
                )
            
            # Create session factory
            SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
            
            # Test the connection first
            with engine.connect() as conn:
                conn.execute(sqlalchemy.text("SELECT 1"))
            
            # Create tables if they don't exist
            # We'll import the models here to avoid circular imports
            from models import User, Settings, DetectionEvent, Conversation, Message, File
            
            # Create any missing tables in the schema, then add any columns that were
            # introduced after the tables were first created.
            Base.metadata.create_all(engine)
            sync_schema(engine)
            
            return True
        except Exception as e:
            last_error = e
            retry_count += 1
            
            if retry_count < MAX_RETRIES:
                time.sleep(RETRY_DELAY)
            else:
                st.error(f"Database connection error after {MAX_RETRIES} attempts.")
                return False
    
    return False

def _literal_default(column, dialect) -> str:
    """SQL DEFAULT clause for a column's scalar Python default, or an empty string."""
    default = column.default
    if default is None or not getattr(default, "is_scalar", False):
        return ""
    literal = sqlalchemy.literal(default.arg, type_=column.type)
    return " DEFAULT " + str(literal.compile(dialect=dialect, compile_kwargs={"literal_binds": True}))


def sync_schema(target_engine) -> list:
    """
    Add columns that exist in the ORM models but not in the database.

    ``create_all`` only creates missing tables; upgrading an existing deployment would
    otherwise fail with "column does not exist" until someone ran a migration script by
    hand. Changes are strictly additive (nullable columns, scalar defaults). Returns the
    list of ``table.column`` names that were added.
    """
    added = []
    inspector = sqlalchemy.inspect(target_engine)
    for table in Base.metadata.sorted_tables:
        if not inspector.has_table(table.name):
            continue
        existing = {column["name"] for column in inspector.get_columns(table.name)}
        for column in table.columns:
            if column.name in existing:
                continue
            preparer = target_engine.dialect.identifier_preparer
            column_type = column.type.compile(dialect=target_engine.dialect)
            ddl = (
                f"ALTER TABLE {preparer.quote(table.name)} "
                f"ADD COLUMN {preparer.quote(column.name)} {column_type}"
                f"{_literal_default(column, target_engine.dialect)}"
            )
            try:
                with target_engine.begin() as conn:
                    conn.execute(sqlalchemy.text(ddl))
                added.append(f"{table.name}.{column.name}")
            except Exception as exc:
                print(f"Schema sync: unable to add {table.name}.{column.name}: {exc}")
    if added:
        print(f"Schema sync: added columns {', '.join(added)}")
    return added


def get_session():
    """Get a new database session"""
    if SessionLocal is None:
        init_db()
    
    if SessionLocal is None:
        st.error("Failed to initialize database connection")
        return None
    
    session = SessionLocal()
    try:
        # Test the session with a simple query
        session.execute(sqlalchemy.text("SELECT 1"))
        return session
    except Exception:
        session.close()
        raise

@contextmanager
def session_scope():
    """Provide a transactional scope around a series of operations."""
    if SessionLocal is None:
        init_db()
        
    if SessionLocal is None:
        st.error("Failed to initialize database connection")
        yield None
        return
        
    session = SessionLocal()
    try:
        yield session
        session.commit()
    except Exception as e:
        session.rollback()
        raise
    finally:
        session.close()