"""Turso / libSQL async connection management with fallback replay.""" from __future__ import annotations import json import os import re from pathlib import Path from typing import Any, Optional from loguru import logger from config import settings from utils.file_utils import ensure_dir _FALLBACK_DIR = ensure_dir(settings.upload_dir / "_db_fallback") # Patterns whose values should be masked in fallback logs _SENSITIVE_PARAM_PATTERNS: list[re.Pattern[str]] = [ re.compile(r"token", re.IGNORECASE), re.compile(r"secret", re.IGNORECASE), re.compile(r"password", re.IGNORECASE), re.compile(r"passwd", re.IGNORECASE), re.compile(r"auth", re.IGNORECASE), re.compile(r"api_key", re.IGNORECASE), ] class _HttpTursoClient: """HTTP-based Turso client (fallback when libsql client is unavailable). Talks to the Turso HTTP API (v2/pipeline) directly using httpx. """ def __init__(self, base_url: str, auth_token: str) -> None: self._base_url = base_url.rstrip("/") self._auth_token = auth_token def execute(self, sql: str, params: tuple = ()) -> Any: """Execute SQL synchronously via HTTP API (runs in executor).""" import httpx # Build typed args for Turso HTTP API def _to_turso_arg(v: Any) -> dict[str, Any]: """Convert Python value to Turso HTTP API arg format. NOTE: Turso HTTP API expects ALL values as strings, even integers! """ if v is None: return {"type": "null", "value": None} if isinstance(v, bool): return {"type": "integer", "value": "1" if v else "0"} if isinstance(v, int): return {"type": "integer", "value": str(v)} return {"type": "text", "value": str(v)} payload = { "requests": [ { "type": "execute", "stmt": { "sql": sql, "args": [_to_turso_arg(p) for p in params] if params else [], }, } ] } resp = httpx.post( f"{self._base_url}/v2/pipeline", json=payload, headers={ "Authorization": f"Bearer {self._auth_token}", "Content-Type": "application/json", }, timeout=15, ) resp.raise_for_status() result = resp.json() # Check for errors in the response for r in result.get("results", []): if r.get("type") == "error": raise RuntimeError(r.get("error", {}).get("message", "Unknown Turso error")) return _HttpTursoResult(result) def fetchall(self) -> list[tuple]: """Not used directly; _HttpTursoResult handles this.""" class _HttpTursoResult: """Simulates the fetchall() return of libsql results.""" def __init__(self, raw: dict) -> None: self._raw = raw def fetchall(self) -> list[tuple]: rows: list[tuple] = [] for r in self._raw.get("results", []): resp = r.get("response", {}) result = resp.get("result", {}) vals = result.get("rows", []) for v in vals: row: list[Any] = [] for c in v: raw_val = c.get("value") col_type = c.get("type", "text") # Convert Turso types to Python types if col_type == "integer": row.append(int(raw_val) if raw_val is not None else 0) elif col_type == "real": row.append(float(raw_val) if raw_val is not None else 0.0) elif col_type == "null": row.append(None) else: row.append(raw_val) rows.append(tuple(row)) return rows class DatabaseClient: """Thin wrapper around libsql-client (sync with run_in_executor). Falls back to local JSON files when Turso is unreachable. When the connection is restored pending fallback entries are replayed automatically. """ def __init__(self) -> None: self._conn: Any = None self._fallback_path: Path = _FALLBACK_DIR / "fallback.json" # ------------------------------------------------------------------ # Public helpers # ------------------------------------------------------------------ async def ensure_tables(self): """Run CREATE TABLE IF NOT EXISTS statements with correct constraints. The UNIQUE constraint is on (tmdb_id_hash, quality), not just tmdb_id_hash, so the same movie can be uploaded in multiple qualities. """ success = await self._execute(""" CREATE TABLE IF NOT EXISTS movies ( id INTEGER PRIMARY KEY AUTOINCREMENT, title TEXT NOT NULL, quality TEXT NOT NULL, tmdb_id INTEGER NOT NULL, tmdb_id_hash TEXT NOT NULL, streamtape_urls TEXT NOT NULL, telegram_file_ids TEXT NOT NULL, file_size INTEGER NOT NULL, original_filename TEXT, bandwidth_used INTEGER DEFAULT 0, processed_at DATETIME DEFAULT CURRENT_TIMESTAMP, admin_id INTEGER NOT NULL ) """) if success: # Migration: remove old UNIQUE constraint on tmdb_id_hash # if it exists (old schema). We create without UNIQUE now. await self._execute( "CREATE UNIQUE INDEX IF NOT EXISTS idx_movies_hash_quality " "ON movies(tmdb_id_hash, quality)" ) await self._execute(""" CREATE TABLE IF NOT EXISTS series ( id INTEGER PRIMARY KEY AUTOINCREMENT, title TEXT NOT NULL, quality TEXT NOT NULL, tmdb_id INTEGER NOT NULL, tmdb_id_hash TEXT NOT NULL, season_number INTEGER NOT NULL, episode_number INTEGER NOT NULL, tmdb_episode_id INTEGER, streamtape_urls TEXT NOT NULL, telegram_file_ids TEXT NOT NULL, file_size INTEGER NOT NULL, original_filename TEXT, bandwidth_used INTEGER DEFAULT 0, processed_at DATETIME DEFAULT CURRENT_TIMESTAMP, admin_id INTEGER NOT NULL ) """) await self._execute( "CREATE UNIQUE INDEX IF NOT EXISTS idx_series_hash_season_ep_quality " "ON series(tmdb_id_hash, season_number, episode_number, quality)" ) await self._execute(""" CREATE TABLE IF NOT EXISTS bandwidth_stats ( id INTEGER PRIMARY KEY AUTOINCREMENT, total_bytes_used INTEGER DEFAULT 0, last_updated DATETIME DEFAULT CURRENT_TIMESTAMP ) """) # Replay any queries that queued while DB was down await self._replay_fallback() async def insert_movie(self, data: dict) -> bool: sql = """INSERT INTO movies (title, quality, tmdb_id, tmdb_id_hash, streamtape_urls, telegram_file_ids, file_size, original_filename, bandwidth_used, admin_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""" params = ( data["title"], data["quality"], data["tmdb_id"], data["tmdb_id_hash"], json.dumps(data["streamtape_urls"]), json.dumps(data["telegram_file_ids"]), data["file_size"], data["original_filename"], data.get("bandwidth_used", 0), data["admin_id"], ) return await self._execute(sql, params) async def insert_series(self, data: dict) -> bool: sql = """INSERT INTO series (title, quality, tmdb_id, tmdb_id_hash, season_number, episode_number, tmdb_episode_id, streamtape_urls, telegram_file_ids, file_size, original_filename, bandwidth_used, admin_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""" params = ( data["title"], data["quality"], data["tmdb_id"], data["tmdb_id_hash"], data["season_number"], data["episode_number"], data.get("tmdb_episode_id"), json.dumps(data["streamtape_urls"]), json.dumps(data["telegram_file_ids"]), data["file_size"], data["original_filename"], data.get("bandwidth_used", 0), data["admin_id"], ) return await self._execute(sql, params) async def get_total_bandwidth(self) -> int: sql = """SELECT COALESCE(SUM(file_size), 0) as total FROM (SELECT file_size FROM movies UNION ALL SELECT file_size FROM series)""" rows = await self._query(sql) if rows: return rows[0][0] # type: ignore[index] return 0 async def check_duplicate(self, tmdb_id_hash: str) -> bool: sql = "SELECT 1 FROM movies WHERE tmdb_id_hash = ? LIMIT 1" rows = await self._query(sql, (tmdb_id_hash,)) if rows: return True sql = "SELECT 1 FROM series WHERE tmdb_id_hash = ? LIMIT 1" rows = await self._query(sql, (tmdb_id_hash,)) return bool(rows) async def check_duplicate_by_key(self, key: str) -> bool: """Check both movies and series for a composite key (hash+quality). The key is matched against ``tmdb_id_hash`` column, assuming the caller encodes quality into the hash lookup (or uses a separate composite logic). This is a best-effort check; a full composite index would be ideal but the current schema stores them separately. """ # Try movies first — ``tmdb_id_hash`` is UNIQUE there sql = "SELECT 1 FROM movies WHERE (tmdb_id_hash || '+' || quality) = ? LIMIT 1" rows = await self._query(sql, (key,)) if rows: return True sql = "SELECT 1 FROM series WHERE (tmdb_id_hash || '+' || quality) = ? LIMIT 1" rows = await self._query(sql, (key,)) return bool(rows) async def get_stats(self) -> dict: """Return aggregate statistics for the dashboard.""" rows_m = await self._query("SELECT COUNT(*), COALESCE(SUM(file_size), 0) FROM movies") rows_s = await self._query("SELECT COUNT(*), COALESCE(SUM(file_size), 0) FROM series") bw = await self._query("SELECT COALESCE(SUM(file_size), 0) FROM (SELECT file_size FROM movies UNION ALL SELECT file_size FROM series)") m_count = rows_m[0][0] if rows_m else 0 m_size = rows_m[0][1] if rows_m else 0 s_count = rows_s[0][0] if rows_s else 0 s_size = rows_s[0][1] if rows_s else 0 total_bw = bw[0][0] if bw else 0 # Last 10 entries recent: list[dict] = [] for table in ("movies", "series"): sql = f"SELECT title, quality, processed_at FROM {table} ORDER BY processed_at DESC LIMIT 5" rows = await self._query(sql) for r in rows: recent.append({"title": r[0], "quality": r[1], "at": str(r[2]) if len(r) > 2 else "", "type": table}) recent.sort(key=lambda x: x["at"], reverse=True) return { "movies_count": m_count, "movies_size": m_size, "series_count": s_count, "series_size": s_size, "total_bandwidth": total_bw, "recent": recent[:10], } # ------------------------------------------------------------------ # Internal — sync Turso calls via run_in_executor # ------------------------------------------------------------------ async def _connect(self): if self._conn is not None: return # Method 1: Try libsql client first try: import libsql_client # type: ignore[import-untyped] self._conn = libsql_client.create_client_sync( url=settings.turso_url, auth_token=settings.turso_token, ) # Verify connection with a simple query self._conn.execute("SELECT 1") logger.info("✅ Turso DB connected (libsql client)") return except Exception as exc: logger.warning( "libsql client failed: {err}", err=exc, ) # Method 2: Fallback to HTTP API (more compatible) try: import httpx # Convert libsql:// URL to https:// http_url = settings.turso_url.replace("libsql://", "https://") self._conn = _HttpTursoClient(http_url, settings.turso_token) # Test the connection test = await self._query("SELECT 1 AS test") if test: logger.info("✅ Turso DB connected (HTTP API fallback)") return except Exception as exc: logger.warning("HTTP API fallback also failed: {err}", err=exc) logger.warning( "❌ All Turso DB connection methods failed. " "All DB operations will use local JSON fallback. " "Check TURSO_URL and TURSO_TOKEN secrets." ) self._conn = None # fallback only async def _execute(self, sql: str, params: tuple = ()) -> bool: await self._connect() if self._conn is not None: try: self._conn.execute(sql, params) return True except Exception as exc: logger.warning("Database execute failed, falling back to local storage: {}", exc) self._write_fallback(sql, params) return False async def _query(self, sql: str, params: tuple = ()) -> list: await self._connect() if self._conn is not None: try: return self._conn.execute(sql, params).fetchall() except Exception as exc: logger.warning("Database query failed, returning empty result: {}", exc) return [] # ------------------------------------------------------------------ # Fallback persistence & replay # ------------------------------------------------------------------ def _mask_sql(self, sql: str, params: tuple) -> tuple[str, tuple]: """Replace sensitive parameter values with ``'***'``. Scans column names in the SQL for known sensitive keywords and masks the corresponding positional parameters so secrets are never written to disk. """ if not params: return sql, params # Simple heuristic: check if any param looks sensitive by key # context. We split the SQL by placeholders and check the # preceding column/alias names. masked_params: list[Any] = [] # Split on '?' placeholders — each part corresponds to a param parts = sql.split("?") for i, param in enumerate(params): # Look at the segment immediately before this placeholder segment = parts[i] if i < len(parts) else "" if any(p.search(segment) for p in _SENSITIVE_PARAM_PATTERNS): masked_params.append("***") else: masked_params.append(param) return sql, tuple(masked_params) def _write_fallback(self, sql: str, params: tuple): """Persist query locally when Turso is unavailable. Sensitive data is masked before writing. The file is capped at ``max_fallback_entries`` entries to keep the file size manageable. """ sql, params = self._mask_sql(sql, params) entry = {"sql": sql, "params": params} entries: list[dict] = [] if self._fallback_path.exists(): try: raw = self._fallback_path.read_text() if raw.strip(): entries = json.loads(raw) except Exception: entries = [] entries.append(entry) # Cap at max_fallback_entries — keep the most recent ones max_entries = settings.max_fallback_entries if len(entries) > max_entries: entries = entries[-max_entries:] with open(self._fallback_path, "w") as f: json.dump(entries, f, indent=2) async def _replay_fallback(self) -> None: """Replay pending fallback queries now that the DB is back. Reads the fallback file, executes each query, and clears the file on success. If *any* query fails the replay stops so we don't lose ordering; remaining entries stay in the file. """ if not self._fallback_path.exists(): return try: raw = self._fallback_path.read_text() if not raw.strip(): return entries: list[dict] = json.loads(raw) except (json.JSONDecodeError, OSError): logger = __import__("loguru", fromlist=["logger"]).logger logger.warning("Corrupt fallback file — skipping replay") return if not entries: return logger = __import__("loguru", fromlist=["logger"]).logger logger.info(f"Replaying {len(entries)} fallback query(ies)") remaining: list[dict] = [] for entry in entries: try: self._conn.execute(entry["sql"], tuple(entry["params"])) except Exception: logger.exception("Fallback replay failed — stopping") remaining.append(entry) break if remaining: # Write back the un-replayed entries with open(self._fallback_path, "w") as f: json.dump(remaining, f, indent=2) else: # All replayed successfully — clear file self._fallback_path.write_text("") logger.info("All fallback queries replayed successfully") db = DatabaseClient()