update 26.7.25
This commit is contained in:
Executable
+302
@@ -0,0 +1,302 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user