Create database.py

database.py

@@ -0,0 +1,64 @@
+"""
+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()