vvzvlad/pyrogram-bridge
1
0
Fork 0
Code Issues 6 Pull Requests Actions 8 Packages Projects Releases Wiki Activity

feat(db): replace media file ID JSON store with SQLite

Browse Source 
Switch persistence of media file identifiers from a JSON file to a
SQLite database. Introduce DB_PATH, init_db_sync, upsert, update, and
removal functions. Remove json-repair dependency and related locking
logic. Update api_server and post_parser to use the new database
helpers.

BREAKING CHANGE: media_file_ids.json is no longer used; existing
data must be migrated to the new SQLite database.
This commit is contained in:
vvzvlad
2026-04-05 20:20:37 +03:00
parent 17ec871d7c
commit c9b0d2ffc9
4 changed files with 146 additions and 203 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api_server.py
+45 -96
View File
@@ -24,7 +24,6 @@ import uvloop
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.background import BackgroundTask
import json_repair
import sys
import magic
@@ -37,7 +36,8 @@ from config import get_settings, setup_logging
from rss_generator import generate_channel_rss, generate_channel_html
from post_parser import PostParser
from url_signer import verify_media_digest, generate_media_digest
from file_io import read_json_file_sync, write_json_file_sync, get_media_file_ids_lock
from file_io import (DB_PATH, init_db_sync, get_all_media_file_ids_sync,
update_media_file_access_sync, remove_media_file_ids_sync)
# Global python-magic instance for MIME type detection
magic_mime = magic.Magic(mime=True)
@@ -78,7 +78,10 @@ async def lifespan(_: FastAPI):
base_cache_dir = os.path.abspath("./data/cache")
os.makedirs(base_cache_dir, exist_ok=True) # Create cache directory
# Initialize SQLite database (creates table if not present)
await asyncio.to_thread(init_db_sync, DB_PATH)
await client.start()
background_task = asyncio.create_task(cache_media_files()) # Start background task
worker_task = asyncio.create_task(background_download_worker()) # Start download worker
@@ -290,21 +293,14 @@ async def download_media_file(channel: Union[str, int], post_id: int, file_uniqu
logger.error(f"cleanup_error: Failed to remove zero-size cached file {cache_path}: {e}")
# Do not raise error here, proceed to download below
else:
# File exists and is not zero size, update timestamp and return
# File exists and is not zero size, update access timestamp and return
logger.info(f"Found cached media file: {cache_path}")
# Update access timestamp in media_file_ids.json
file_ids_path = os.path.join(os.path.abspath("./data"), 'media_file_ids.json')
try:
if os.path.exists(file_ids_path):
media_files = await asyncio.to_thread(read_json_file_sync, file_ids_path)
for file_data in media_files:
if (file_data.get('channel') == str(channel) and # Ensure channel is string for comparison
file_data.get('post_id') == post_id and
file_data.get('file_unique_id') == file_unique_id):
file_data['added'] = datetime.now().timestamp()
break
async with get_media_file_ids_lock():
await asyncio.to_thread(write_json_file_sync, file_ids_path, media_files)
await asyncio.to_thread(
update_media_file_access_sync,
DB_PATH, str(channel), post_id, file_unique_id,
datetime.now().timestamp()
)
except Exception as e:
logger.error(f"Failed to update timestamp for {channel}/{post_id}/{file_unique_id}: {str(e)}")
return cache_path, False
@@ -313,31 +309,17 @@ async def download_media_file(channel: Union[str, int], post_id: int, file_uniqu
if not file_id:
error_message = f"Media with file_unique_id {file_unique_id} not found in message {post_id} for channel {channel}"
logger.error(error_message)
# Attempt to remove the invalid entry from media_file_ids.json using threadpool I/O
file_ids_path = os.path.join(os.path.abspath("./data"), 'media_file_ids.json')
try:
if os.path.exists(file_ids_path):
media_files = await asyncio.to_thread(read_json_file_sync, file_ids_path)
initial_count = len(media_files)
media_files_updated = [
f_data for f_data in media_files
if not (
f_data.get('channel') == str(channel) and
f_data.get('post_id') == post_id and
f_data.get('file_unique_id') == file_unique_id
)
]
if len(media_files_updated) < initial_count:
async with get_media_file_ids_lock():
await asyncio.to_thread(write_json_file_sync, file_ids_path, media_files_updated)
logger.info(f"Removed invalid entry for {channel}/{post_id}/{file_unique_id} from media_file_ids.json")
else:
logger.warning(f"Entry for {channel}/{post_id}/{file_unique_id} not found in media_file_ids.json for removal")
# Remove the invalid entry from the SQLite database
try:
await asyncio.to_thread(
remove_media_file_ids_sync,
DB_PATH, [(str(channel), post_id, file_unique_id)]
)
logger.info(f"Removed invalid entry for {channel}/{post_id}/{file_unique_id} from SQLite")
except Exception as e:
logger.error(f"Failed to remove entry for {channel}/{post_id}/{file_unique_id} from media_file_ids.json: {str(e)}")
logger.error(f"Failed to remove entry for {channel}/{post_id}/{file_unique_id} from SQLite: {str(e)}")
raise HTTPException(status_code=404, detail="File not found in message")
try:
@@ -431,7 +413,7 @@ def remove_old_cached_files_sync(media_files: list, cache_dir: str) -> tuple[lis
files_removed += 1
logger.info(f"Removed temporary file: {file_path}")
# Also remove entry from media_file_ids.json if exists
# Also remove the temporary file entry from the in-memory list
updated_media_files = [
f for f in updated_media_files
if not (f.get('channel') == channel and
@@ -491,32 +473,6 @@ async def download_new_files(media_files: list, cache_dir: str) -> None:
logger.info(f"Queued {files_queued} files for background download")
def fix_corrupted_json(file_path: str) -> list:
"""
Attempt to fix corrupted JSON file using json-repair library
Returns list of valid entries
"""
try:
with open(file_path, 'r', encoding='utf-8') as f:
content = f.read()
# Try to repair and parse the JSON
fixed_data = json_repair.loads(content)
# Validate entries
valid_entries = []
if isinstance(fixed_data, list):
for entry in fixed_data:
if isinstance(entry, dict) and all(key in entry for key in ['channel', 'post_id', 'file_unique_id']):
valid_entries.append(entry)
logger.info(f"Fixed JSON file: {file_path}, found {len(valid_entries)} valid entries")
return valid_entries
except Exception as e:
logger.error(f"Failed to fix JSON file {file_path}: {str(e)}")
return []
async def background_download_worker():
"""Worker that processes downloads from queue"""
while True:
@@ -541,37 +497,33 @@ async def cache_media_files() -> None:
delay = 60
while True:
try:
file_path = os.path.join(os.path.abspath("./data"), 'media_file_ids.json')
if not os.path.exists(file_path):
await asyncio.sleep(delay)
continue
try:
media_files = await asyncio.to_thread(read_json_file_sync, file_path)
except json.JSONDecodeError:
logger.error(f"JSON decode error in {file_path}, attempting to fix")
media_files = await asyncio.to_thread(fix_corrupted_json, file_path)
if media_files:
await asyncio.to_thread(write_json_file_sync, file_path, media_files)
else:
logger.error("Failed to fix JSON file, skipping cache update")
await asyncio.sleep(delay)
continue
# Load all media file ID records from SQLite
media_files = await asyncio.to_thread(get_all_media_file_ids_sync, DB_PATH)
cache_dir = os.path.abspath("./data/cache")
updated_media_files, files_removed = await asyncio.to_thread(remove_old_cached_files_sync, media_files, cache_dir)
if files_removed > 0:
try:
async with get_media_file_ids_lock():
await asyncio.to_thread(write_json_file_sync, file_path, updated_media_files)
# Determine which entries were removed and delete them from SQLite
updated_set = {
(r['channel'], r['post_id'], r['file_unique_id'])
for r in updated_media_files
}
removed_entries = [
(r['channel'], r['post_id'], r['file_unique_id'])
for r in media_files
if (r['channel'], r['post_id'], r['file_unique_id']) not in updated_set
]
if removed_entries:
await asyncio.to_thread(remove_media_file_ids_sync, DB_PATH, removed_entries)
logger.info(f"Removed {files_removed} old files from cache")
except Exception as e:
logger.error(f"Failed to update media_file_ids.json: {str(e)}")
logger.error(f"Failed to remove old entries from SQLite: {str(e)}")
await download_new_files(updated_media_files, cache_dir)
await asyncio.sleep(delay) # Check every delay seconds
except Exception as e:
logger.error(f"Cache media files error: {str(e)}")
await asyncio.sleep(delay)
@@ -616,17 +568,14 @@ def calculate_cache_stats() -> dict[str, Any]:
cache_files_count = 0
cache_total_size_mb = 0
media_file_ids_path = os.path.join(os.path.abspath("./data"), "media_file_ids.json")
cache_times = []
if os.path.exists(media_file_ids_path):
try:
with open(media_file_ids_path, "r", encoding="utf-8") as json_file:
media_files = json.load(json_file)
for entry in media_files:
if "added" in entry:
cache_times.append(entry["added"])
except Exception as e:
logger.error(f"Error reading media_file_ids.json: {str(e)}")
try:
media_files = get_all_media_file_ids_sync(DB_PATH)
for entry in media_files:
if "added" in entry:
cache_times.append(entry["added"])
except Exception as e:
logger.error(f"Error reading media file IDs from SQLite: {str(e)}")
if cache_times:
cache_time_diff_seconds = max(cache_times) - min(cache_times)
cache_time_diff_days = round(cache_time_diff_seconds / 86400, 2) # rounded to two decimals
file_io.py
+71 -18
View File
@@ -5,31 +5,84 @@
# pylint: disable=missing-function-docstring
import os
import json
import asyncio
from typing import Any, List
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")
# One shared lock for operations that write to media_file_ids.json
_media_file_ids_lock: asyncio.Lock | None = None
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")
return conn
def get_media_file_ids_lock() -> asyncio.Lock:
global _media_file_ids_lock
if _media_file_ids_lock is None:
_media_file_ids_lock = asyncio.Lock()
return _media_file_ids_lock
@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 read_json_file_sync(file_path: str) -> List[dict[str, Any]]:
with open(file_path, 'r', encoding='utf-8') as f:
return json.load(f)
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)
)
"""
)
def write_json_file_sync(file_path: str, data: List[dict[str, Any]]) -> None:
temp_path = f"{file_path}.tmp"
with open(temp_path, 'w', encoding='utf-8') as f:
json.dump(data, f, ensure_ascii=False, indent=2)
os.replace(temp_path, file_path)
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 OR REPLACE INTO media_file_ids (channel, post_id, file_unique_id, added) VALUES (?, ?, ?, ?)",
(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,
)
post_parser.py
+30 -88
View File
@@ -21,7 +21,7 @@ from pyrogram.enums import MessageMediaType
from bleach.css_sanitizer import CSSSanitizer
from bleach import clean as HTMLSanitizer
from config import get_settings
from file_io import read_json_file_sync, write_json_file_sync, get_media_file_ids_lock
from file_io import upsert_media_file_id_sync, DB_PATH
from url_signer import generate_media_digest
Config = get_settings()
@@ -932,60 +932,16 @@ class PostParser:
logger.error(f"file_id_extraction_error: media_type {message.media}, error {str(e)}")
return None
async def _persist_media_file_id_async(self, file_path: str, file_data: dict) -> None:
import time as _time
task_start = _time.monotonic()
channel = file_data.get('channel', '?')
post_id = file_data.get('post_id', '?')
file_uid = file_data.get('file_unique_id', '?')
async def _persist_media_file_id_async(self, channel: str, post_id: int, file_unique_id: str, added: float) -> None:
"""Persist a single media file ID record to SQLite (fire-and-forget)."""
try:
existing_data = []
if os.path.exists(file_path):
read_start = _time.monotonic()
existing_data = await asyncio.to_thread(read_json_file_sync, file_path)
read_elapsed = _time.monotonic() - read_start
if read_elapsed > 0.1:
logger.warning(f"diag_persist_slow_read: {channel}/{post_id}/{file_uid} read {len(existing_data)} entries in {read_elapsed:.3f}s")
found = False
for item in existing_data:
if (item.get('channel') == file_data['channel'] and
item.get('post_id') == file_data['post_id'] and
item.get('file_unique_id') == file_data['file_unique_id']):
item['added'] = datetime.now().timestamp()
found = True
break
if not found:
existing_data.append(file_data)
lock_wait_start = _time.monotonic()
async with get_media_file_ids_lock():
lock_acquired = _time.monotonic()
lock_wait = lock_acquired - lock_wait_start
if lock_wait > 0.05:
logger.warning(f"diag_persist_lock_wait: {channel}/{post_id}/{file_uid} waited {lock_wait:.3f}s for lock")
write_start = _time.monotonic()
await asyncio.to_thread(write_json_file_sync, file_path, existing_data)
write_elapsed = _time.monotonic() - write_start
if write_elapsed > 0.1:
logger.warning(f"diag_persist_slow_write: {channel}/{post_id}/{file_uid} wrote {len(existing_data)} entries in {write_elapsed:.3f}s")
await asyncio.to_thread(upsert_media_file_id_sync, DB_PATH, channel, post_id, file_unique_id, added)
logger.debug(f"persist_media_file_id: upserted {channel}/{post_id}/{file_unique_id}")
except Exception as e:
logger.error(f"file_id_save_error: error writing to {file_path}, error {str(e)}")
finally:
total = _time.monotonic() - task_start
if total > 0.2:
logger.warning(f"diag_persist_slow_total: {channel}/{post_id}/{file_uid} total {total:.3f}s")
logger.error(f"file_id_save_error: error upserting {channel}/{post_id}/{file_unique_id}, error {str(e)}")
def _save_media_file_ids(self, message: Message) -> None:
try:
file_data = {
'channel': '',
'post_id': 0,
'file_unique_id': '',
'added': .0
}
channel_username = self.get_channel_username(message)
if not channel_username:
logger.error(f"channel_username_error: no username found for chat in message {message.id}")
@@ -993,55 +949,41 @@ class PostParser:
if message.media:
# Skip large videos - they shouldn't be cached permanently
if message.video and message.video.file_size and message.video.file_size > 100 * 1024 * 1024: return
if message.photo: file_data['file_unique_id'] = message.photo.file_unique_id
elif message.video: file_data['file_unique_id'] = message.video.file_unique_id
elif message.document: file_data['file_unique_id'] = message.document.file_unique_id
elif message.audio: file_data['file_unique_id'] = message.audio.file_unique_id
elif message.voice: file_data['file_unique_id'] = message.voice.file_unique_id
elif message.video_note: file_data['file_unique_id'] = message.video_note.file_unique_id
elif message.animation: file_data['file_unique_id'] = message.animation.file_unique_id
elif message.sticker: file_data['file_unique_id'] = message.sticker.file_unique_id
if message.video and message.video.file_size and message.video.file_size > 100 * 1024 * 1024:
return
file_unique_id = ''
if message.photo: file_unique_id = message.photo.file_unique_id
elif message.video: file_unique_id = message.video.file_unique_id
elif message.document: file_unique_id = message.document.file_unique_id
elif message.audio: file_unique_id = message.audio.file_unique_id
elif message.voice: file_unique_id = message.voice.file_unique_id
elif message.video_note: file_unique_id = message.video_note.file_unique_id
elif message.animation: file_unique_id = message.animation.file_unique_id
elif message.sticker: file_unique_id = message.sticker.file_unique_id
elif message.web_page and message.web_page.photo:
file_data['file_unique_id'] = message.web_page.photo.file_unique_id
file_unique_id = message.web_page.photo.file_unique_id
if file_data['file_unique_id']:
file_data['channel'] = channel_username
file_data['post_id'] = message.id
file_data['added'] = datetime.now().timestamp()
file_path = os.path.join(os.path.abspath("./data"), 'media_file_ids.json')
if file_unique_id:
added_ts = datetime.now().timestamp()
try:
loop = asyncio.get_running_loop()
if loop.is_running():
task = loop.create_task(self._persist_media_file_id_async(file_path, file_data))
# Диагностика: считаем сколько fire-and-forget задач висит в очереди
task = loop.create_task(
self._persist_media_file_id_async(channel_username, message.id, file_unique_id, added_ts)
)
# Count pending fire-and-forget tasks to surface backlog issues
all_tasks = asyncio.all_tasks(loop)
persist_tasks = [t for t in all_tasks if '_persist_media_file_id_async' in str(t.get_coro())]
task_count = len(persist_tasks)
if task_count > 5:
logger.warning(f"diag_persist_task_queue: {task_count} pending _persist_media_file_id_async tasks (msg {message.id})")
logger.debug(f"diag_persist_task_created: queued for {channel_username}/{message.id}/{file_data['file_unique_id']}, total pending: {task_count}")
logger.warning(f"persist_task_queue: {task_count} pending _persist_media_file_id_async tasks (msg {message.id})")
logger.debug(f"persist_task_created: queued for {channel_username}/{message.id}/{file_unique_id}, total pending: {task_count}")
else:
# Fallback sync write (should not occur during normal FastAPI runtime)
existing_data = []
if os.path.exists(file_path):
existing_data = read_json_file_sync(file_path) # type: ignore
found = False
for item in existing_data:
if (item.get('channel') == file_data['channel'] and
item.get('post_id') == file_data['post_id'] and
item.get('file_unique_id') == file_data['file_unique_id']):
item['added'] = datetime.now().timestamp()
found = True
break
if not found:
existing_data.append(file_data)
write_json_file_sync(file_path, existing_data) # type: ignore
# Fallback sync path (should not occur during normal FastAPI runtime)
upsert_media_file_id_sync(DB_PATH, channel_username, message.id, file_unique_id, added_ts)
except Exception as e:
logger.error(f"file_id_save_error: error writing to {file_path}, error {str(e)}")
logger.error(f"file_id_save_error: error saving {channel_username}/{message.id}/{file_unique_id}, error {str(e)}")
except Exception as e:
logger.error(f"file_id_collection_error: message_id {message.id}, error {str(e)}")
requirements.txt
-1
View File
@@ -11,4 +11,3 @@ python-dateutil==2.9.0.post0
python-magic==0.4.27
bleach[css]==6.1.0
types-bleach
json-repair