"""SQLite persistence for Telegram download records.""" import json import os import sqlite3 from dataclasses import asdict, dataclass from datetime import datetime, timezone from typing import Optional, Tuple from module.app import DownloadStatus from module.pyrogram_extension import get_extension @dataclass class DownloadRecord: """One persisted Telegram download outcome.""" chat_id: str message_id: int download_status: str chat_title: Optional[str] = None message_date: Optional[str] = None message_thread_id: Optional[int] = None topic_id: Optional[int] = None reply_to_message_id: Optional[int] = None media_group_id: Optional[str] = None sender_id: Optional[str] = None sender_name: Optional[str] = None media_type: Optional[str] = None caption: Optional[str] = None message_text: Optional[str] = None file_id: Optional[str] = None file_unique_id: Optional[str] = None original_file_name: Optional[str] = None file_extension: Optional[str] = None mime_type: Optional[str] = None file_size: Optional[int] = None width: Optional[int] = None height: Optional[int] = None duration: Optional[int] = None stored_file_name: Optional[str] = None local_path: Optional[str] = None relative_path: Optional[str] = None original_save_path: Optional[str] = None nfo_path: Optional[str] = None error_message: Optional[str] = None retry_count: int = 0 downloaded_at: Optional[str] = None created_at: Optional[str] = None updated_at: Optional[str] = None extra_json: Optional[str] = None MEDIA_ATTRS = ( "audio", "document", "photo", "sticker", "animation", "video", "voice", "video_note", "new_chat_photo", ) SCHEMA = """ CREATE TABLE IF NOT EXISTS download_records ( id INTEGER PRIMARY KEY AUTOINCREMENT, chat_id TEXT NOT NULL, chat_title TEXT, message_id INTEGER NOT NULL, message_date TEXT, message_thread_id INTEGER, topic_id INTEGER, reply_to_message_id INTEGER, media_group_id TEXT, sender_id TEXT, sender_name TEXT, media_type TEXT, caption TEXT, message_text TEXT, file_id TEXT, file_unique_id TEXT, original_file_name TEXT, file_extension TEXT, mime_type TEXT, file_size INTEGER, width INTEGER, height INTEGER, duration INTEGER, stored_file_name TEXT, local_path TEXT, relative_path TEXT, original_save_path TEXT, nfo_path TEXT, download_status TEXT NOT NULL, error_message TEXT, retry_count INTEGER NOT NULL DEFAULT 0, downloaded_at TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, extra_json TEXT, UNIQUE(chat_id, message_id) ) """ def utc_now() -> str: """Return a compact UTC timestamp for database rows.""" return ( datetime.now(timezone.utc) .replace(microsecond=0) .isoformat() .replace("+00:00", "Z") ) def get_database_path(save_path: str) -> str: """Return the SQLite path under the configured download root.""" return os.path.join(save_path, "download_records.sqlite3") def init_database(db_path: str): """Create the records database and table if needed.""" os.makedirs(os.path.dirname(db_path), exist_ok=True) with sqlite3.connect(db_path) as conn: conn.execute(SCHEMA) conn.commit() def status_to_text(status: DownloadStatus) -> str: """Map internal enum values to persisted status text.""" if status is DownloadStatus.SuccessDownload: return "success" if status is DownloadStatus.SkipDownload: return "skip" if status is DownloadStatus.FailedDownload: return "failed" return "failed" def find_message_media(message) -> Tuple[Optional[str], Optional[object]]: """Return the first media type/object present on the message.""" for media_type in MEDIA_ATTRS: media_obj = getattr(message, media_type, None) if media_obj is not None: return media_type, media_obj if getattr(message, "text", None): return "text", None return None, None def isoformat(value) -> Optional[str]: """Safely convert datetime-like values to ISO strings.""" if not value: return None if hasattr(value, "isoformat"): return value.isoformat() return str(value) def get_file_extension(media_obj) -> Optional[str]: """Return an extension without the leading dot.""" if media_obj is None: return None original_file_name = getattr(media_obj, "file_name", None) if original_file_name: _, suffix = os.path.splitext(os.path.basename(original_file_name)) if suffix: return suffix.lstrip(".") file_id = getattr(media_obj, "file_id", "") mime_type = getattr(media_obj, "mime_type", "") try: extension = get_extension(file_id, mime_type, False) except Exception: return None return extension or None def get_sender_identity(message) -> Tuple[Optional[str], Optional[str]]: """Return sender id/name from user or sender chat metadata.""" from_user = getattr(message, "from_user", None) if from_user: sender_id = getattr(from_user, "id", None) username = getattr(from_user, "username", None) name_parts = [ getattr(from_user, "first_name", None), getattr(from_user, "last_name", None), ] sender_name = username or " ".join(part for part in name_parts if part) or None return str(sender_id) if sender_id is not None else None, sender_name sender_chat = getattr(message, "sender_chat", None) if sender_chat: sender_id = getattr(sender_chat, "id", None) sender_name = ( getattr(sender_chat, "title", None) or getattr(sender_chat, "username", None) or None ) return str(sender_id) if sender_id is not None else None, sender_name return None, None def build_download_record( app, node, message, download_status: DownloadStatus, local_path: Optional[str] = None, nfo_path: Optional[str] = None, error_message: Optional[str] = None, ) -> DownloadRecord: """Build a database record from a message and final download outcome.""" media_type, media_obj = find_message_media(message) chat = getattr(message, "chat", None) sender_id, sender_name = get_sender_identity(message) local_path = local_path or None now = utc_now() relative_path = None stored_file_name = None if local_path: stored_file_name = os.path.basename(local_path) try: relative_path = os.path.relpath(local_path, app.save_path) except ValueError: relative_path = local_path message_thread_id = getattr(message, "message_thread_id", None) caption = getattr(message, "caption", None) if not caption and getattr(message, "media_group_id", None): caption = app.get_caption_name(node.chat_id, message.media_group_id) extra = {} if media_obj is not None and hasattr(media_obj, "date"): extra["media_date"] = isoformat(getattr(media_obj, "date")) return DownloadRecord( chat_id=str(node.chat_id), chat_title=getattr(chat, "title", None), message_id=int(message.id), message_date=isoformat(getattr(message, "date", None)), message_thread_id=message_thread_id, topic_id=message_thread_id, reply_to_message_id=getattr(message, "reply_to_message_id", None), media_group_id=getattr(message, "media_group_id", None), sender_id=sender_id, sender_name=sender_name, media_type=media_type, caption=caption, message_text=getattr(message, "text", None), file_id=getattr(media_obj, "file_id", None) if media_obj is not None else None, file_unique_id=getattr(media_obj, "file_unique_id", None) if media_obj is not None else None, original_file_name=getattr(media_obj, "file_name", None) if media_obj is not None else None, file_extension=get_file_extension(media_obj), mime_type=getattr(media_obj, "mime_type", None) if media_obj is not None else None, file_size=getattr(media_obj, "file_size", None) if media_obj is not None else None, width=getattr(media_obj, "width", None) if media_obj is not None else None, height=getattr(media_obj, "height", None) if media_obj is not None else None, duration=getattr(media_obj, "duration", None) if media_obj is not None else None, stored_file_name=stored_file_name, local_path=local_path, relative_path=relative_path, original_save_path=app.save_path, nfo_path=nfo_path, download_status=status_to_text(download_status), error_message=error_message, retry_count=0, downloaded_at=now if download_status is DownloadStatus.SuccessDownload else None, created_at=now, updated_at=now, extra_json=json.dumps(extra, ensure_ascii=False) if extra else None, ) def upsert_download_record(db_path: str, record: DownloadRecord): """Insert or update one download record.""" init_database(db_path) data = asdict(record) columns = list(data.keys()) placeholders = ", ".join(["?"] * len(columns)) column_sql = ", ".join(columns) update_columns = [ column for column in columns if column not in {"chat_id", "message_id", "created_at"} ] update_sql = ", ".join([f"{column}=excluded.{column}" for column in update_columns]) sql = ( f"INSERT INTO download_records ({column_sql}) VALUES ({placeholders}) " f"ON CONFLICT(chat_id, message_id) DO UPDATE SET {update_sql}" ) with sqlite3.connect(db_path) as conn: conn.execute(sql, [data[column] for column in columns]) conn.commit()