import os from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker, declarative_base # <-- CORRECTED IMPORT # --- Configuration --- # Determines the database mode. Can be "postgres" or "sqlite". # Defaults to "postgres" if not set. DB_MODE = os.getenv("DB_MODE", "postgres").lower() # Default database URLs POSTGRES_DEFAULT_URL = "postgresql://user:password@localhost/ai_hub_db" SQLITE_DEFAULT_URL = "sqlite:///./data/ai_hub.db" DATABASE_URL = "" engine_args = {} # --- Database Initialization --- if DB_MODE == "sqlite": print("✅ Initializing with SQLite in-file database.") DATABASE_URL = SQLITE_DEFAULT_URL # SQLite requires a specific argument to allow access from multiple threads, # which is common in web applications. engine_args = {"connect_args": {"check_same_thread": False}} else: # Default to postgres # Use the provided DATABASE_URL or fall back to the default. DATABASE_URL = os.getenv("DATABASE_URL", POSTGRES_DEFAULT_URL) DB_MODE = "postgres" print(f"✅ Initializing with PostgreSQL database. URL: {DATABASE_URL}") # pool_pre_ping checks if a connection is still alive before using it from the pool. engine_args = {"pool_pre_ping": True} # Create the SQLAlchemy engine with the determined settings engine = create_engine(DATABASE_URL, **engine_args) # SessionLocal is a factory for creating new database session objects SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) # Base is a class that our database model classes will inherit from. Base = declarative_base() # --- Dependency for FastAPI --- def get_db(): """ FastAPI dependency that provides a database session for a single API request. It ensures the session is always closed after the request is finished. """ db = SessionLocal() try: yield db finally: db.close()