Newer
Older
cortex-hub / ai-hub / app / db_setup.py
import os
from dotenv import load_dotenv
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, Session
from app.db.database import Base # Assuming `Base` is in this file

# Load environment variables from a .env file
load_dotenv()

# --- Database Connection Setup ---
# This configuration allows for easy switching between SQLite and PostgreSQL.
DB_MODE = os.getenv("DB_MODE", "sqlite")
if DB_MODE == "sqlite":
    DATABASE_URL = "sqlite:///./ai_hub.db"
    # The connect_args are needed for SQLite to work with FastAPI's multiple threads
    engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
else:
    DATABASE_URL = os.getenv("DATABASE_URL", "postgresql://user:password@localhost/ai_hub_db")
    engine = create_engine(DATABASE_URL)

# Create a database session class
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def create_db_tables():
    """Create all database tables based on the models."""
    print("Creating database tables...")
    Base.metadata.create_all(bind=engine)

# The dependency to get a database session
def get_db():
    """Dependency that provides a database session."""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()