Newer
Older
cortex-hub / ai-hub / app / db / database.py
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()