From c9b0d2ffc9c426a2e83502f46adc1730435f64a4 Mon Sep 17 00:00:00 2001 From: vvzvlad Date: Sun, 5 Apr 2026 20:20:37 +0300 Subject: [PATCH] feat(db): replace media file ID JSON store with SQLite 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. --- api_server.py | 141 +++++++++++++++-------------------------------- file_io.py | 89 ++++++++++++++++++++++++------ post_parser.py | 118 ++++++++++----------------------------- requirements.txt | 1 - 4 files changed, 146 insertions(+), 203 deletions(-) diff --git a/api_server.py b/api_server.py index 8ab97e3..c53a2af 100644 --- a/api_server.py +++ b/api_server.py @@ -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 diff --git a/file_io.py b/file_io.py index b8c2f8a..269bcb4 100644 --- a/file_io.py +++ b/file_io.py @@ -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, + ) diff --git a/post_parser.py b/post_parser.py index 643219b..edb5281 100644 --- a/post_parser.py +++ b/post_parser.py @@ -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)}") diff --git a/requirements.txt b/requirements.txt index d67bb19..ac8ece5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,4 +11,3 @@ python-dateutil==2.9.0.post0 python-magic==0.4.27 bleach[css]==6.1.0 types-bleach -json-repair