from sqlalchemy import create_engine, event, text from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker import logging from backend.config import settings logger = logging.getLogger(__name__) IS_SQLITE = "sqlite" in settings.DATABASE_URL engine_kwargs = {"pool_pre_ping": True, "echo": settings.DEBUG} if IS_SQLITE: engine_kwargs["connect_args"] = {"check_same_thread": False} else: engine_kwargs["pool_size"] = 5 engine_kwargs["max_overflow"] = 10 engine_kwargs["pool_timeout"] = 30 engine = create_engine(settings.DATABASE_URL, **engine_kwargs) if IS_SQLITE: @event.listens_for(engine, "connect") def set_sqlite_pragma(dbapi_connection, connection_record): cursor = dbapi_connection.cursor() cursor.execute("PRAGMA journal_mode=WAL") cursor.execute("PRAGMA foreign_keys=ON") cursor.close() SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) Base = declarative_base() def get_db(): db = SessionLocal() try: yield db finally: db.close() def _table_has_column(conn, table_name: str, column_name: str) -> bool: """Check if a column exists in a PostgreSQL table.""" result = conn.execute(text(""" SELECT COUNT(*) FROM information_schema.columns WHERE table_name = :t AND column_name = :c """), {"t": table_name, "c": column_name}) return result.scalar() > 0 def create_tables(): """ Create all tables safely. If an existing table is missing columns (broken schema from old run), drop it and recreate cleanly. """ from backend import models # noqa: F401 if IS_SQLITE: Base.metadata.create_all(bind=engine, checkfirst=True) logger.info("✅ SQLite tables ready") return # PostgreSQL: check if tables exist with correct schema with engine.connect() as conn: # Check if users table exists but is broken (missing 'username' column) result = conn.execute(text( "SELECT COUNT(*) FROM information_schema.tables WHERE table_name='users'" )) users_exists = result.scalar() > 0 if users_exists and not _table_has_column(conn, "users", "username"): logger.warning("⚠️ Broken schema detected — dropping all tables and recreating...") conn.execute(text("DROP TABLE IF EXISTS messages CASCADE")) conn.execute(text("DROP TABLE IF EXISTS conversations CASCADE")) conn.execute(text("DROP TABLE IF EXISTS user_api_keys CASCADE")) conn.execute(text("DROP TABLE IF EXISTS users CASCADE")) conn.commit() logger.info("🗑 Old tables dropped") # Now create fresh Base.metadata.create_all(bind=engine, checkfirst=True) logger.info("✅ PostgreSQL tables ready")