From c0580b3abf4a3560ff7aa56c985fc7312eac4a71 Mon Sep 17 00:00:00 2001 From: vvzvlad Date: Mon, 3 Feb 2025 03:04:50 +0300 Subject: [PATCH] Implement media file caching and background download mechanism --- .gitignore | 2 +- api_server.py | 102 +++++++++++++++++++++++++++++++++++++++++++------ post_parser.py | 34 +++++++++++------ 3 files changed, 113 insertions(+), 25 deletions(-) diff --git a/.gitignore b/.gitignore index 4705c9a..08eff2d 100644 --- a/.gitignore +++ b/.gitignore @@ -2,4 +2,4 @@ __pycache__/ .env .DS_Store -media_file_ids.json +data/ diff --git a/api_server.py b/api_server.py index c8eaa8c..e7926c5 100644 --- a/api_server.py +++ b/api_server.py @@ -1,6 +1,9 @@ import logging import os import mimetypes +import asyncio +import json +from datetime import datetime from contextlib import asynccontextmanager from fastapi import FastAPI, HTTPException, Response, BackgroundTasks @@ -12,13 +15,32 @@ from pyrogram import errors from post_parser import PostParser logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + +if not logger.handlers: + handler = logging.StreamHandler() + handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) + logger.addHandler(handler) + client = TelegramClient() Config = get_settings() @asynccontextmanager async def lifespan(_: FastAPI): await client.start() + + # Start background task + background_task = asyncio.create_task(cache_media_files()) + yield + + # Cleanup + background_task.cancel() + try: + await background_task + except asyncio.CancelledError: + pass + await client.stop() app = FastAPI( title="Pyrogram Bridge", lifespan=lifespan) @@ -51,16 +73,14 @@ async def find_file_id_in_message(message, file_unique_id: str): return None async def prepare_file_response(file_path: str, background_tasks: BackgroundTasks): - """Prepare file response with proper headers and cleanup task""" + """Prepare file response with proper headers""" if not os.path.exists(file_path): - raise HTTPException(status_code=404, detail="Temporary file not found") + raise HTTPException(status_code=404, detail="File not found") media_type, _ = mimetypes.guess_type(file_path) if not media_type: media_type = "application/octet-stream" - background_tasks.add_task(lambda: os.remove(file_path) if os.path.exists(file_path) else None) - headers = {"Content-Disposition": f"inline; filename={os.path.basename(file_path)}"} return FileResponse(path=file_path, media_type=media_type, headers=headers) @@ -112,17 +132,34 @@ async def health_check(): logger.error(error_message) raise HTTPException(status_code=500, detail=error_message) from e +async def download_media_file(channel: str, post_id: int, file_unique_id: str) -> str: + """ + Download media file from Telegram and save to cache + Returns path to downloaded file + """ + cache_dir = os.path.abspath("./data/cache") + cache_path = os.path.join(cache_dir, f"{channel}-{post_id}-{file_unique_id}") + os.makedirs(cache_dir, exist_ok=True) + + if os.path.exists(cache_path): + logger.info(f"Found cached media file: {cache_path}") + return cache_path + + message = await client.client.get_messages(channel, post_id) + file_id = await find_file_id_in_message(message, file_unique_id) + + if not file_id: + logger.error(f"Media with file_unique_id {file_unique_id} not found in message {post_id}") + raise HTTPException(status_code=404, detail="File not found in message") + + file_path = await client.client.download_media(file_id, file_name=cache_path) + logger.info(f"Downloaded media file {file_unique_id} to {cache_path}") + return file_path + @app.get("/media/{channel}/{post_id}/{file_unique_id}") async def get_media(channel: str, post_id: int, file_unique_id: str, background_tasks: BackgroundTasks): try: - message = await client.client.get_messages(channel, post_id) - file_id = await find_file_id_in_message(message, file_unique_id) - - if not file_id: - logger.error(f"Media with file_unique_id {file_unique_id} not found in message {post_id}") - raise HTTPException(status_code=404, detail="File not found in message") - file_path = await client.client.download_media(file_id) - logger.info(f"Downloaded media file {file_unique_id} to {file_path}") + file_path = await download_media_file(channel, post_id, file_unique_id) return await prepare_file_response(file_path, background_tasks) except HTTPException: raise @@ -147,3 +184,44 @@ async def get_rss_feed(channel: str, limit: int = 20): error_message = f"Failed to generate RSS feed for channel {channel}: {str(e)}" logger.error(error_message) raise HTTPException(status_code=500, detail=error_message) from e + +async def cache_media_files(): + """Background task to download media files to cache""" + 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(60) + continue + + with open(file_path, 'r', encoding='utf-8') as f: + media_files = json.load(f) + + for file_data in media_files: + try: + file_id = file_data['file_id'] + if not file_id: + continue + + # Parse file_id to get parameters + channel, post_id_str, file_unique_id = file_id.split('-', 2) + post_id = int(post_id_str) # Convert post_id to integer + + # Check if file already exists in cache + cache_dir = os.path.abspath("./data/cache") + cache_path = os.path.join(cache_dir, f"{channel}-{post_id}-{file_unique_id}") + + if not os.path.exists(cache_path): + logger.info(f"Background download started for {file_id}") + await download_media_file(channel, post_id, file_unique_id) + await asyncio.sleep(1) # Delay between downloads + + except Exception as e: + logger.error(f"Background download failed for {file_id}: {str(e)}") + continue + + await asyncio.sleep(60) # Check every minute + + except Exception as e: + logger.error(f"Cache media files error: {str(e)}") + await asyncio.sleep(60) diff --git a/post_parser.py b/post_parser.py index a8accdb..8d4cc2d 100644 --- a/post_parser.py +++ b/post_parser.py @@ -12,6 +12,13 @@ import json Config = get_settings() logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + +if not logger.handlers: + handler = logging.StreamHandler() + handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) + logger.addHandler(handler) + #tests #http://127.0.0.1:8000/post/html/DragorWW_space/114 — video @@ -65,7 +72,7 @@ class PostParser: raise def _format_json(self, message: Message, naked: bool = False) -> Dict[Any, Any]: - html_content = self._format_html(message) + html_content = self._format_html(message, naked=naked) result = { 'channel': message.chat.username, 'message_id': message.id, @@ -148,9 +155,12 @@ class PostParser: base_url = Config['pyrogram_bridge_url'] if message.media: file_unique_id = self._get_file_unique_id(message) - print(file_unique_id) + if file_unique_id is None: + logger.error(f"File unique id not found for message {message.id}: {message}") + return if file_unique_id: url = f"{base_url}/media/{message.chat.username}/{message.id}/{file_unique_id}" + logger.debug(f"Collected media file: {message.chat.username}/{message.id}/{file_unique_id}") html_content.append(f'
') if message.media in [MessageMediaType.PHOTO, MessageMediaType.DOCUMENT]: html_content.append(f'') @@ -330,19 +340,19 @@ class PostParser: } if message.media: - if message.photo: file_data['file_id'] = message.photo.file_unique_id - elif message.video: file_data['file_id'] = message.video.file_unique_id - elif message.document: file_data['file_id'] = message.document.file_unique_id - elif message.audio: file_data['file_id'] = message.audio.file_unique_id - elif message.voice: file_data['file_id'] = message.voice.file_unique_id - elif message.video_note: file_data['file_id'] = message.video_note.file_unique_id - elif message.animation: file_data['file_id'] = message.animation.file_unique_id - elif message.sticker: file_data['file_id'] = message.sticker.file_unique_id + if message.photo: file_data['file_id'] = f"{message.chat.username}-{message.id}-{message.photo.file_unique_id}" + elif message.video: file_data['file_id'] = f"{message.chat.username}-{message.id}-{message.video.file_unique_id}" + elif message.document: file_data['file_id'] = f"{message.chat.username}-{message.id}-{message.document.file_unique_id}" + elif message.audio: file_data['file_id'] = f"{message.chat.username}-{message.id}-{message.audio.file_unique_id}" + elif message.voice: file_data['file_id'] = f"{message.chat.username}-{message.id}-{message.voice.file_unique_id}" + elif message.video_note: file_data['file_id'] = f"{message.chat.username}-{message.id}-{message.video_note.file_unique_id}" + elif message.animation: file_data['file_id'] = f"{message.chat.username}-{message.id}-{message.animation.file_unique_id}" + elif message.sticker: file_data['file_id'] = f"{message.chat.username}-{message.id}-{message.sticker.file_unique_id}" elif message.web_page and message.web_page.photo: - file_data['file_id'] = message.web_page.photo.file_unique_id + file_data['file_id'] = f"{message.chat.username}-{message.id}-{message.web_page.photo.file_unique_id}" if file_data['file_id']: - file_path = 'media_file_ids.json' + file_path = os.path.join(os.path.abspath("./data"), 'media_file_ids.json') try: existing_data = [] if os.path.exists(file_path):