Files
pyrogram-bridge/file_io.py
T
vvzvlad f0e38b9776 perf(core): improve async handling and caching performance
Convert delayed_delete_file to async and use asyncio.sleep.
Add pre‑semaphore cache hit to serve cached files without acquiring the download semaphore.
Set SQLite busy_timeout to reduce lock errors.
Introduce a diagnostics counter for pending _persist_media_file_id_async tasks.
Refactor flag extraction to accept pre‑generated HTML and compute html body once.
Batch reply enrichment in RSS generator to minimize API calls.
Store cached messages directly in tg_cache and add fallback for legacy double‑pickle format.

BREAKING CHANGE: delayed_delete_file is now async and must be awaited.
2026-05-18 00:08:46 +03:00

121 lines
4.4 KiB
Python

#!/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 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),
)