"""Database configuration and session management.""" import os from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker, AsyncSession from sqlalchemy.orm import DeclarativeBase DATABASE_URL = os.getenv("DATABASE_URL", "sqlite+aiosqlite:///./xdl.db") engine = create_async_engine(DATABASE_URL, echo=False) async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) class Base(DeclarativeBase): pass async def get_db(): async with async_session() as session: yield session async def init_db(): async with engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) # Ensure indexes exist on already-created tables (idempotent) from sqlalchemy import text await conn.execute(text( "CREATE INDEX IF NOT EXISTS ix_video_url_format_id ON videos (url, format_id)" )) await conn.execute(text( "CREATE INDEX IF NOT EXISTS ix_download_logs_video_id ON download_logs (video_id)" )) await conn.execute(text( "CREATE INDEX IF NOT EXISTS ix_download_logs_downloaded_at ON download_logs (downloaded_at)" )) # Migrate: add geo columns to existing download_logs table (idempotent) for col_def in [ "ALTER TABLE download_logs ADD COLUMN country_code VARCHAR(8) DEFAULT ''", "ALTER TABLE download_logs ADD COLUMN country VARCHAR(128) DEFAULT ''", "ALTER TABLE download_logs ADD COLUMN city VARCHAR(128) DEFAULT ''", ]: try: await conn.execute(text(col_def)) except Exception: pass # Column already exists # Seed default cleanup settings (only if not already set) defaults = { "cleanup_enabled": "true", "cleanup_retention_minutes": "10080", # 7 days "cleanup_storage_limit_pct": "80", "cleanup_last_run": "", "cleanup_last_result": "", } for k, v in defaults.items(): await conn.execute(text( "INSERT OR IGNORE INTO app_settings (key, value, updated_at) VALUES (:k, :v, datetime('now'))" ), {"k": k, "v": v})