""" Database initialization and session management """ import os from urllib.parse import urlparse, parse_qs from sqlalchemy import create_engine, event from sqlalchemy.orm import sessionmaker, Session from sqlalchemy.pool import StaticPool from models import Base # Database URL from environment or use SQLite DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./hf_uploader.db") # Parse and fix the database URL for TiDB Cloud if DATABASE_URL and DATABASE_URL.startswith("mysql://"): # Parse the URL parsed = urlparse(DATABASE_URL) # Build a clean URL without SSL parameters clean_url = f"mysql+pymysql://{parsed.username}:{parsed.password}@{parsed.hostname}:{parsed.port or 3306}{parsed.path}" # Add SSL parameters properly formatted for PyMySQL clean_url += "?ssl_verify_cert=false&ssl_verify_identity=false" DATABASE_URL = clean_url # Create engine with proper configuration if "sqlite" in DATABASE_URL: engine = create_engine( DATABASE_URL, connect_args={"check_same_thread": False}, poolclass=StaticPool, echo=False ) # Enable foreign keys for SQLite @event.listens_for(engine, "connect") def set_sqlite_pragma(dbapi_conn, connection_record): cursor = dbapi_conn.cursor() cursor.execute("PRAGMA foreign_keys=ON") cursor.close() else: engine = create_engine( DATABASE_URL, echo=False, pool_pre_ping=True, pool_recycle=3600 ) # Create session factory SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) def init_db(): """Initialize database tables""" Base.metadata.create_all(bind=engine) def get_db() -> Session: """Get database session""" db = SessionLocal() try: yield db finally: db.close()