pyrogram-bridge/file_io.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# flake8: noqa
# pylint: disable=missing-function-docstring

import os
import sqlite3
from contextlib import contextmanager
from typing import List

# Path to SQLite database file
DB_PATH = os.path.join(os.path.abspath("./data"), "media_file_ids.db")


def _open_db(db_path: str) -> sqlite3.Connection:
    """Open a SQLite connection with WAL journal mode enabled."""
    conn = sqlite3.connect(db_path)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute("PRAGMA busy_timeout = 5000")  # Wait up to 5 seconds on lock
    return conn


@contextmanager
def _db_connection(db_path: str):
    """Context manager that opens a WAL-mode SQLite connection and ensures it is closed."""
    conn = _open_db(db_path)
    try:
        yield conn
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def init_db_sync(db_path: str) -> None:
    """Create the media_file_ids table if it does not exist."""
    with _db_connection(db_path) as conn:
        conn.execute(
            """
            CREATE TABLE IF NOT EXISTS media_file_ids (
                channel        TEXT    NOT NULL,
                post_id        INTEGER NOT NULL,
                file_unique_id TEXT    NOT NULL,
                added          REAL    NOT NULL,
                PRIMARY KEY (channel, post_id, file_unique_id)
            )
            """
        )
        # Add mime_type column if it does not exist yet (idempotent migration)
        try:
            conn.execute("ALTER TABLE media_file_ids ADD COLUMN mime_type TEXT")
        except sqlite3.OperationalError as e:
            if "duplicate column name" not in str(e):
                raise  # Re-raise any OperationalError that is not "duplicate column name"


def upsert_media_file_id_sync(db_path: str, channel: str, post_id: int, file_unique_id: str, added: float) -> None:
    """Insert or replace a single media file ID record."""
    with _db_connection(db_path) as conn:
        conn.execute(
            """INSERT INTO media_file_ids (channel, post_id, file_unique_id, added)
               VALUES (?, ?, ?, ?)
               ON CONFLICT(channel, post_id, file_unique_id)
               DO UPDATE SET added = excluded.added""",
            (channel, post_id, file_unique_id, added),
        )


def upsert_media_file_ids_bulk_sync(db_path: str, entries: List[tuple]) -> None:
    """Insert or replace multiple media file ID records in a single transaction.

    entries: iterable of (channel, post_id, file_unique_id, added) tuples.
    Uses executemany for batched upserts (one connection, one commit).
    """
    if not entries:
        return
    with _db_connection(db_path) as conn:
        conn.executemany(
            """INSERT INTO media_file_ids (channel, post_id, file_unique_id, added)
               VALUES (?, ?, ?, ?)
               ON CONFLICT(channel, post_id, file_unique_id)
               DO UPDATE SET added = excluded.added""",
            entries,
        )


def update_media_file_access_sync(db_path: str, channel: str, post_id: int, file_unique_id: str, added: float) -> None:
    """Update the access timestamp for an existing media file ID record."""
    with _db_connection(db_path) as conn:
        conn.execute(
            "UPDATE media_file_ids SET added = ? WHERE channel = ? AND post_id = ? AND file_unique_id = ?",
            (added, channel, post_id, file_unique_id),
        )


def get_all_media_file_ids_sync(db_path: str) -> List[dict]:
    """Return all rows from media_file_ids as a list of dicts."""
    with _db_connection(db_path) as conn:
        conn.row_factory = sqlite3.Row
        cursor = conn.execute("SELECT channel, post_id, file_unique_id, added FROM media_file_ids")
        rows = cursor.fetchall()
    return [dict(row) for row in rows]


def remove_media_file_ids_sync(db_path: str, entries: List[tuple]) -> None:
    """Remove media file ID records identified by (channel, post_id, file_unique_id) tuples."""
    with _db_connection(db_path) as conn:
        conn.executemany(
            "DELETE FROM media_file_ids WHERE channel = ? AND post_id = ? AND file_unique_id = ?",
            entries,
        )


def get_mime_type_sync(db_path: str, channel: str, post_id: int, file_unique_id: str) -> str | None:
    """Return the cached MIME type for a given media key, or None if not stored yet."""
    with _db_connection(db_path) as conn:
        cursor = conn.execute(
            "SELECT mime_type FROM media_file_ids WHERE channel = ? AND post_id = ? AND file_unique_id = ?",
            (channel, post_id, file_unique_id),
        )
        row = cursor.fetchone()
    if row is None:
        return None
    return row[0]  # May be None if the column value was never set


def set_mime_type_sync(db_path: str, channel: str, post_id: int, file_unique_id: str, mime_type: str) -> None:
    """Persist a detected MIME type for an existing media file ID record."""
    with _db_connection(db_path) as conn:
        conn.execute(
            "UPDATE media_file_ids SET mime_type = ? WHERE channel = ? AND post_id = ? AND file_unique_id = ?",
            (mime_type, channel, post_id, file_unique_id),
        )