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

Implement media file caching and background download mechanism

Browse Source
This commit is contained in:
vvzvlad
2025-02-03 03:04:50 +03:00
parent b0e03ce6e8
commit c0580b3abf
3 changed files with 113 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1 -1
View File
@@ -2,4 +2,4 @@ __pycache__/
.env
.DS_Store
media_file_ids.json
data/
api_server.py
+90 -12
View File
@@ -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)
post_parser.py
+22 -12
View File
@@ -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'<div class="message-media">')
if message.media in [MessageMediaType.PHOTO, MessageMediaType.DOCUMENT]:
html_content.append(f'<img src="{url}" style="max-width:600px; max-height:600px; object-fit:contain;">')
@@ -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):