KillerKing93/Transformers-TextEngine-OpenAPI

1#!/usr/bin/env python2# -*- coding: utf-8 -*-3"""4Database connection and session management for AI Marketplace Platform5 6Usage:7    from database import get_db, init_db8 9    # Initialize database10    init_db()11 12    # Use in FastAPI endpoints13    @app.get("/example")14    def example(db: Session = Depends(get_db)):15        ...16"""17 18import os19from sqlalchemy import create_engine20from sqlalchemy.orm import sessionmaker, Session21from sqlalchemy.pool import StaticPool22from models import Base23 24# Database configuration from environment25DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./marketplace.db")26 27# Create engine28# For SQLite: use StaticPool and check_same_thread=False for FastAPI compatibility29# For PostgreSQL/MySQL: remove connect_args30if DATABASE_URL.startswith("sqlite"):31    engine = create_engine(32        DATABASE_URL,33        connect_args={"check_same_thread": False},34        poolclass=StaticPool35    )36else:37    engine = create_engine(DATABASE_URL, pool_pre_ping=True)38 39# Create session factory40SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)41 42 43def init_db():44    """45    Initialize database by creating all tables.46    Should be called at application startup.47    """48    Base.metadata.create_all(bind=engine)49    print(f"[Database] Tables created/verified. Database: {DATABASE_URL}")50 51 52def get_db() -> Session:53    """54    Dependency for FastAPI endpoints to get database session.55 56    Usage:57        @app.get("/items")58        def get_items(db: Session = Depends(get_db)):59            items = db.query(Product).all()60            return items61    """62    db = SessionLocal()63    try:64        yield db65    finally:66        db.close()67