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:///./data/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()