File size: 18,290 Bytes
f913c73
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c432b4b
 
 
 
 
 
 
 
 
 
 
 
 
 
75688cd
 
d359454
 
 
 
 
75688cd
d359454
 
 
 
 
75688cd
c432b4b
 
 
 
75688cd
 
 
 
c432b4b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dcdc936
 
 
 
 
 
 
 
 
 
 
 
 
 
c432b4b
 
 
f913c73
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f6d563d
f913c73
f6d563d
 
f913c73
 
 
 
 
 
 
f6d563d
f913c73
 
 
 
 
 
 
 
 
 
f6d563d
 
 
 
 
 
f913c73
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f6d563d
 
 
 
f913c73
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c432b4b
 
f913c73
 
 
 
 
 
 
c432b4b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f913c73
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
"""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()