Newer
Older
cortex-hub / ai-hub / app / db / session.py
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from app.config import settings
from app.db.database import Base

# Determine engine arguments based on the database mode from the central config
engine_args = {}
if settings.DB_MODE == "sqlite":
    # This argument is required for SQLite to allow it to be used by multiple threads,
    # which is the case in a web application like FastAPI.
    engine_args["connect_args"] = {"check_same_thread": False}
else:
    # 'pool_pre_ping' checks if a database connection is still alive before using it.
    # This prevents errors from connections that have been timed out by the DB server.
    engine_args["pool_pre_ping"] = True

# Create the SQLAlchemy engine using the centralized URL and determined arguments
engine = create_engine(settings.DATABASE_URL, **engine_args)

# SessionLocal is a factory for creating new database session objects.
# It's the standard way to interact with the database in SQLAlchemy.
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def create_db_and_tables():
    """
    Creates all database tables defined by models inheriting from Base.
    This is typically called once on application startup.
    """
    print("Creating database tables...")
    # Base.metadata contains all the schema information from your models.
    Base.metadata.create_all(bind=engine)

def get_db():
    """
    FastAPI dependency that provides a database session for a single API request.
    
    This pattern ensures that the database session is always closed after the
    request is finished, even if an error occurs.
    """
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()