303 lines
9.8 KiB
Python
Executable File
303 lines
9.8 KiB
Python
Executable File
"""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()
|