feat(cache): снапшот-кеш истории вместо pickle + префиксный limit + джиттер TTL + age-sweep
Docker Image CI / build (pull_request) Waiting to run
Docker Image CI / build (pull_request) Waiting to run
Кеш истории (data/tgcache/) больше не хранит pickle живых объектов pyrogram. Новый модуль message_snapshot.py: при записи из Message извлекается JSON-словарь по явному allowlist (снапшот), при чтении восстанавливается duck-typed CachedMessage (CachedStr с .html, mutable, deepcopy-safe), неотличимый для пайплайна рендера от Message. post_parser.py и rss_generator.py НЕ изменены. tg_cache.py переписан на generic JSON-store (pickle удалён полностью): - _store_entry/_load_entry: атомарная запись через уникальный .tmp.<uuid4> + os.replace, finally чистит свой tmp; версия/TTL/битый JSON -> None; - пути <safe_key>.history.json / .chatinfo.json; - пакет B: префиксный limit (miss только когда cached_limit<limit и есть ещё); - пакет F: джиттер TTL при записи (random.uniform(0.8,1.0), на чтении random не зовётся) + sweep_tgcache(age) + cleanup_legacy_cache_files; - api_server.py: старт-cleanup+sweep в lifespan, per-pass sweep в cache_media_files, удаление legacy data/media_file_ids.json. closes #23 Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -43,6 +43,7 @@ from file_io import (DB_PATH, init_db_sync, get_all_media_file_ids_sync,
|
||||
update_media_file_access_sync, update_media_file_access_bulk_sync,
|
||||
remove_media_file_ids_sync,
|
||||
get_mime_type_sync, set_mime_type_sync)
|
||||
from tg_cache import cleanup_legacy_cache_files, sweep_tgcache
|
||||
|
||||
# Global python-magic instance for MIME type detection
|
||||
magic_mime = magic.Magic(mime=True)
|
||||
@@ -240,6 +241,19 @@ async def lifespan(_: FastAPI):
|
||||
# Initialize SQLite database (creates table if not present)
|
||||
await asyncio.to_thread(init_db_sync, DB_PATH)
|
||||
|
||||
# One-shot startup maintenance of the history/chatinfo cache: drop legacy pickle files
|
||||
# (which are now always a miss), age-sweep stale tgcache entries, and remove the
|
||||
# pre-SQLite data/media_file_ids.json legacy dump (no longer read by any code).
|
||||
await asyncio.to_thread(cleanup_legacy_cache_files)
|
||||
await asyncio.to_thread(sweep_tgcache)
|
||||
legacy_media_ids = os.path.join("data", "media_file_ids.json")
|
||||
if os.path.exists(legacy_media_ids):
|
||||
try:
|
||||
await asyncio.to_thread(os.remove, legacy_media_ids)
|
||||
logger.info(f"legacy_media_file_ids_removed: {legacy_media_ids}")
|
||||
except OSError as e:
|
||||
logger.warning(f"legacy_media_file_ids_remove_error: {e}")
|
||||
|
||||
await client.start()
|
||||
# Supervise the background tasks: if either dies (not via cancellation) it is logged
|
||||
# CRITICAL and restarted, so a crash can no longer silently stop cache sweeping or downloads.
|
||||
@@ -911,6 +925,11 @@ async def cache_media_files() -> None:
|
||||
logger.error(f"Failed to remove old entries from SQLite: {str(e)}")
|
||||
|
||||
await download_new_files(updated_media_files, cache_dir)
|
||||
|
||||
# Age-sweep the history/chatinfo cache once per pass (dead channels, orphaned
|
||||
# uuid tmp files). MUST run off-loop via to_thread (blocking filesystem walk).
|
||||
await asyncio.to_thread(sweep_tgcache)
|
||||
|
||||
await asyncio.sleep(delay) # Check every delay seconds
|
||||
|
||||
except Exception as e:
|
||||
|
||||
@@ -0,0 +1,385 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
|
||||
# flake8: noqa
|
||||
# pylint: disable=broad-exception-raised, raise-missing-from, too-many-arguments, redefined-outer-name
|
||||
# pylint: disable=multiple-statements, logging-fstring-interpolation, trailing-whitespace, line-too-long
|
||||
# pylint: disable=broad-exception-caught, missing-function-docstring, missing-class-docstring
|
||||
# pylance: disable=reportMissingImports, reportMissingModuleSource
|
||||
|
||||
"""JSON snapshot / restore for pyrogram Message objects (issue #23).
|
||||
|
||||
The history cache no longer pickles live pyrogram objects. On write a plain JSON
|
||||
snapshot is extracted from a Message via an explicit allowlist (schema v1); on read
|
||||
a duck-typed CachedMessage is restored that is indistinguishable from a real Message
|
||||
to the render pipeline (post_parser.py / rss_generator.py are NOT changed).
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from datetime import datetime
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from pyrogram.enums import MessageMediaType
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Bump when the snapshot schema changes in a backwards-incompatible way; a mismatch
|
||||
# makes _load_entry treat the cache file as a miss (old files are simply re-fetched).
|
||||
SNAPSHOT_VERSION = 1
|
||||
|
||||
|
||||
class CachedStr(str):
|
||||
"""A ``str`` subclass carrying a precomputed ``.html`` rendering.
|
||||
|
||||
Mirrors pyrogram's ``Str`` (which exposes ``.html`` computed from entities). Here the
|
||||
HTML is computed once at snapshot time and stored, so the render pipeline can read
|
||||
``message.text.html`` on a restored message exactly as on a live one. The instance
|
||||
``__dict__`` (holding ``.html``) is preserved across deepcopy/pickle via ``__reduce__``.
|
||||
"""
|
||||
|
||||
html: str
|
||||
|
||||
@classmethod
|
||||
def build(cls, plain: str, html: str) -> "CachedStr":
|
||||
obj = cls(plain)
|
||||
obj.html = html
|
||||
return obj
|
||||
|
||||
def __reduce__(self):
|
||||
# Preserve the precomputed .html across copy.deepcopy() and pickle. Without this a
|
||||
# str subclass would be reconstructed as a bare str and lose the instance __dict__.
|
||||
return (CachedStr.build, (str(self), self.html))
|
||||
|
||||
|
||||
def _unwrap_text(value: Any) -> Optional[str]:
|
||||
"""Unwrap a styled-text container to a plain string.
|
||||
|
||||
In kurigram 2.2.23 Poll.question and every PollOption.text are always FormattedText
|
||||
objects (``.text`` holds the string). The rule ``v.text if hasattr(v, 'text') else
|
||||
str(v)`` is correct for FormattedText, for a bare Str and for a plain string. Storing
|
||||
a FormattedText as-is would make json.dump raise TypeError, silently killing the cache
|
||||
for any channel that has polls.
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
inner = value.text if hasattr(value, "text") else value
|
||||
return str(inner)
|
||||
|
||||
|
||||
def _snapshot_str(value: Any) -> Optional[dict]:
|
||||
"""Snapshot a text/caption value: {'plain', 'html'} or None if absent.
|
||||
|
||||
``.html`` is computed at write time (live pyrogram Str exposes it); if unavailable the
|
||||
html falls back to the plain text.
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
plain = str(value)
|
||||
html = getattr(value, "html", None)
|
||||
if html is None:
|
||||
html = plain
|
||||
return {"plain": plain, "html": str(html)}
|
||||
|
||||
|
||||
def _enum_name(value: Any) -> Any:
|
||||
"""Return an enum member's ``.name`` (JSON-safe) or the value unchanged."""
|
||||
if value is None:
|
||||
return None
|
||||
return value.name if hasattr(value, "name") else value
|
||||
|
||||
|
||||
def _snapshot_obj(obj: Any, keys: List[str]) -> Optional[dict]:
|
||||
"""Snapshot the allowlisted ``keys`` off ``obj`` (getattr-only), or None if obj is None."""
|
||||
if obj is None:
|
||||
return None
|
||||
return {k: getattr(obj, k, None) for k in keys}
|
||||
|
||||
|
||||
def _snapshot_chat(chat: Any) -> Optional[dict]:
|
||||
if chat is None:
|
||||
return None
|
||||
usernames = getattr(chat, "usernames", None)
|
||||
snap_usernames = None
|
||||
if usernames:
|
||||
snap_usernames = [
|
||||
{"username": getattr(u, "username", None), "active": getattr(u, "active", None)}
|
||||
for u in usernames
|
||||
]
|
||||
return {
|
||||
"id": getattr(chat, "id", None),
|
||||
"username": getattr(chat, "username", None),
|
||||
"title": getattr(chat, "title", None),
|
||||
"usernames": snap_usernames,
|
||||
}
|
||||
|
||||
|
||||
# forward_origin is the ONE object whose key SET mirrors the live object: the
|
||||
# MessageOrigin* classes each carry a DIFFERENT set of attributes and
|
||||
# _format_forward_info branches by hasattr (Case 1-5). Presence of a key in the
|
||||
# snapshot must mirror presence of the attribute on the source object.
|
||||
_FORWARD_ORIGIN_KEYS = ["type", "chat", "sender_user_name", "sender_user", "chat_id", "title"]
|
||||
|
||||
|
||||
def _snapshot_forward_origin(fo: Any) -> Optional[dict]:
|
||||
if fo is None:
|
||||
return None
|
||||
result: dict = {}
|
||||
for key in _FORWARD_ORIGIN_KEYS:
|
||||
if not hasattr(fo, key):
|
||||
continue
|
||||
value = getattr(fo, key)
|
||||
if key == "type":
|
||||
result[key] = _enum_name(value)
|
||||
elif key == "chat":
|
||||
result[key] = _snapshot_obj(value, ["id", "title", "username"])
|
||||
elif key == "sender_user":
|
||||
result[key] = _snapshot_obj(value, ["first_name", "last_name", "username"])
|
||||
else:
|
||||
result[key] = value
|
||||
return result
|
||||
|
||||
|
||||
def _snapshot_reactions(reactions: Any) -> Optional[list]:
|
||||
if reactions is None:
|
||||
return None
|
||||
rlist = getattr(reactions, "reactions", None)
|
||||
if rlist is None:
|
||||
return None
|
||||
out = []
|
||||
for r in rlist:
|
||||
cid = getattr(r, "custom_emoji_id", None)
|
||||
out.append({
|
||||
# Custom-emoji reactions carry no unicode emoji: emoji is null and the
|
||||
# custom_emoji_id is stored as a STRING = str(document_id).
|
||||
"emoji": None if cid is not None else getattr(r, "emoji", None),
|
||||
"custom_emoji_id": str(cid) if cid is not None else None,
|
||||
"count": getattr(r, "count", None),
|
||||
"is_paid": getattr(r, "is_paid", None),
|
||||
})
|
||||
return out
|
||||
|
||||
|
||||
def _snapshot_poll(poll: Any) -> Optional[dict]:
|
||||
if poll is None:
|
||||
return None
|
||||
options = getattr(poll, "options", None)
|
||||
snap_options = None
|
||||
if options:
|
||||
snap_options = [{"text": _unwrap_text(getattr(o, "text", None))} for o in options]
|
||||
return {
|
||||
"question": _unwrap_text(getattr(poll, "question", None)),
|
||||
"options": snap_options,
|
||||
}
|
||||
|
||||
|
||||
def _snapshot_web_page(wp: Any) -> Optional[dict]:
|
||||
if wp is None:
|
||||
return None
|
||||
return {
|
||||
"type": getattr(wp, "type", None),
|
||||
"url": getattr(wp, "url", None),
|
||||
"display_url": getattr(wp, "display_url", None),
|
||||
"site_name": getattr(wp, "site_name", None),
|
||||
"title": getattr(wp, "title", None),
|
||||
"description": getattr(wp, "description", None),
|
||||
"has_large_media": getattr(wp, "has_large_media", None),
|
||||
"photo": _snapshot_obj(getattr(wp, "photo", None), ["file_unique_id"]),
|
||||
}
|
||||
|
||||
|
||||
def snapshot_message(message: Any) -> dict:
|
||||
"""Extract a JSON-serializable snapshot (schema v1) from a pyrogram Message.
|
||||
|
||||
Every field is read via getattr(..., None). See the module docstring / issue #23 for
|
||||
the field contract.
|
||||
"""
|
||||
date = getattr(message, "date", None)
|
||||
media = getattr(message, "media", None)
|
||||
service = getattr(message, "service", None)
|
||||
return {
|
||||
"id": getattr(message, "id", None),
|
||||
# isoformat()/fromisoformat() preserve naive/aware exactly as pyrogram stores it.
|
||||
"date": date.isoformat() if date is not None else None,
|
||||
"text": _snapshot_str(getattr(message, "text", None)),
|
||||
"caption": _snapshot_str(getattr(message, "caption", None)),
|
||||
"media": media.name if media is not None else None,
|
||||
"service": service.name if service is not None else None,
|
||||
"media_group_id": getattr(message, "media_group_id", None),
|
||||
"views": getattr(message, "views", None),
|
||||
"show_caption_above_media": getattr(message, "show_caption_above_media", None),
|
||||
"reply_to_message_id": getattr(message, "reply_to_message_id", None),
|
||||
"empty": getattr(message, "empty", None),
|
||||
"chat": _snapshot_chat(getattr(message, "chat", None)),
|
||||
"sender_chat": _snapshot_obj(getattr(message, "sender_chat", None), ["id", "title", "username"]),
|
||||
"from_user": _snapshot_obj(getattr(message, "from_user", None), ["first_name", "last_name", "username"]),
|
||||
"forward_origin": _snapshot_forward_origin(getattr(message, "forward_origin", None)),
|
||||
"reactions": _snapshot_reactions(getattr(message, "reactions", None)),
|
||||
"poll": _snapshot_poll(getattr(message, "poll", None)),
|
||||
"web_page": _snapshot_web_page(getattr(message, "web_page", None)),
|
||||
"photo": _snapshot_obj(getattr(message, "photo", None), ["file_unique_id"]),
|
||||
"video": _snapshot_obj(getattr(message, "video", None), ["file_unique_id", "file_size"]),
|
||||
"document": _snapshot_obj(getattr(message, "document", None), ["file_unique_id", "mime_type"]),
|
||||
"audio": _snapshot_obj(getattr(message, "audio", None), ["file_unique_id", "mime_type"]),
|
||||
"voice": _snapshot_obj(getattr(message, "voice", None), ["file_unique_id", "mime_type"]),
|
||||
"video_note": _snapshot_obj(getattr(message, "video_note", None), ["file_unique_id"]),
|
||||
"animation": _snapshot_obj(getattr(message, "animation", None), ["file_unique_id"]),
|
||||
"sticker": _snapshot_obj(getattr(message, "sticker", None), ["file_unique_id", "emoji", "is_video"]),
|
||||
}
|
||||
|
||||
|
||||
def _ns(d: Optional[dict], keys: List[str]) -> Optional[SimpleNamespace]:
|
||||
"""Restore a nested object with the FULL schema key set (None-defaults).
|
||||
|
||||
Live pyrogram objects always set every attribute in __init__, and the pipeline reads
|
||||
them directly (e.g. rss_generator ~131 message.chat.username in an except; post_parser
|
||||
~1087 u.username / u.active without getattr). So restored objects must carry all keys.
|
||||
Returns None if ``d`` is None (the attribute was absent on the live object).
|
||||
"""
|
||||
if d is None:
|
||||
return None
|
||||
ns = SimpleNamespace()
|
||||
for k in keys:
|
||||
setattr(ns, k, d.get(k))
|
||||
return ns
|
||||
|
||||
|
||||
def _restore_chat(d: Optional[dict]) -> Optional[SimpleNamespace]:
|
||||
if d is None:
|
||||
return None
|
||||
usernames = d.get("usernames")
|
||||
restored_usernames = None
|
||||
if usernames is not None:
|
||||
restored_usernames = [_ns(u, ["username", "active"]) for u in usernames]
|
||||
return SimpleNamespace(
|
||||
id=d.get("id"),
|
||||
username=d.get("username"),
|
||||
title=d.get("title"),
|
||||
usernames=restored_usernames,
|
||||
)
|
||||
|
||||
|
||||
def _restore_forward_origin(d: Optional[dict]) -> Optional[SimpleNamespace]:
|
||||
"""Restore forward_origin mirroring ONLY the recorded keys (presence-semantics)."""
|
||||
if d is None:
|
||||
return None
|
||||
ns = SimpleNamespace()
|
||||
for key, value in d.items():
|
||||
if key == "chat":
|
||||
setattr(ns, key, _ns(value, ["id", "title", "username"]))
|
||||
elif key == "sender_user":
|
||||
setattr(ns, key, _ns(value, ["first_name", "last_name", "username"]))
|
||||
else:
|
||||
setattr(ns, key, value)
|
||||
return ns
|
||||
|
||||
|
||||
def _restore_reactions(items: Optional[list]) -> Optional[SimpleNamespace]:
|
||||
if items is None:
|
||||
return None
|
||||
reactions = [_ns(r, ["emoji", "custom_emoji_id", "count", "is_paid"]) for r in items]
|
||||
# Mirror pyrogram: message.reactions is an object exposing a .reactions list.
|
||||
return SimpleNamespace(reactions=reactions)
|
||||
|
||||
|
||||
def _restore_poll(d: Optional[dict]) -> Optional[SimpleNamespace]:
|
||||
if d is None:
|
||||
return None
|
||||
options = d.get("options") or []
|
||||
# options must be namespace objects with a .text string; a bare string would make
|
||||
# getattr(option, 'text', '') return '' and render empty options.
|
||||
restored_options = [SimpleNamespace(text=o.get("text")) for o in options]
|
||||
return SimpleNamespace(question=d.get("question"), options=restored_options)
|
||||
|
||||
|
||||
def _restore_web_page(d: Optional[dict]) -> Optional[SimpleNamespace]:
|
||||
if d is None:
|
||||
return None
|
||||
return SimpleNamespace(
|
||||
type=d.get("type"),
|
||||
url=d.get("url"),
|
||||
display_url=d.get("display_url"),
|
||||
site_name=d.get("site_name"),
|
||||
title=d.get("title"),
|
||||
description=d.get("description"),
|
||||
has_large_media=d.get("has_large_media"),
|
||||
photo=_ns(d.get("photo"), ["file_unique_id"]),
|
||||
)
|
||||
|
||||
|
||||
def _restore_media(name: Optional[str]) -> Optional[MessageMediaType]:
|
||||
if name is None:
|
||||
return None
|
||||
try:
|
||||
return MessageMediaType[name]
|
||||
except KeyError:
|
||||
# A media type unknown to this kurigram build: warn and drop to None so the 31
|
||||
# post_parser sites that compare against / index by the enum keep working.
|
||||
logger.warning(f"snapshot_unknown_media_type: {name}")
|
||||
return None
|
||||
|
||||
|
||||
def _restore_str(d: Optional[dict]) -> Optional[CachedStr]:
|
||||
if d is None:
|
||||
return None
|
||||
return CachedStr.build(d.get("plain", ""), d.get("html", d.get("plain", "")))
|
||||
|
||||
|
||||
class CachedMessage:
|
||||
"""Duck-typed stand-in for a pyrogram Message, restored from a snapshot dict.
|
||||
|
||||
Mutable (reply enrichment assigns ``.reply_to_message``). Every top-level attribute the
|
||||
render pipeline reads exists with a None/False default so getattr never raises.
|
||||
"""
|
||||
|
||||
def __init__(self, data: dict):
|
||||
self._snapshot = data
|
||||
self.id = data.get("id")
|
||||
date = data.get("date")
|
||||
self.date = datetime.fromisoformat(date) if date is not None else None
|
||||
self.text = _restore_str(data.get("text"))
|
||||
self.caption = _restore_str(data.get("caption"))
|
||||
self.media = _restore_media(data.get("media"))
|
||||
# service restored as a plain string; all consumers use truthiness and
|
||||
# `'X' in str(service)`.
|
||||
self.service = data.get("service")
|
||||
self.media_group_id = data.get("media_group_id")
|
||||
self.views = data.get("views")
|
||||
self.show_caption_above_media = data.get("show_caption_above_media")
|
||||
self.reply_to_message_id = data.get("reply_to_message_id")
|
||||
self.reply_to_message = None
|
||||
self.empty = bool(data.get("empty"))
|
||||
self.chat = _restore_chat(data.get("chat"))
|
||||
self.sender_chat = _ns(data.get("sender_chat"), ["id", "title", "username"])
|
||||
self.from_user = _ns(data.get("from_user"), ["first_name", "last_name", "username"])
|
||||
self.forward_origin = _restore_forward_origin(data.get("forward_origin"))
|
||||
self.reactions = _restore_reactions(data.get("reactions"))
|
||||
self.poll = _restore_poll(data.get("poll"))
|
||||
self.web_page = _restore_web_page(data.get("web_page"))
|
||||
self.photo = _ns(data.get("photo"), ["file_unique_id"])
|
||||
self.video = _ns(data.get("video"), ["file_unique_id", "file_size"])
|
||||
self.document = _ns(data.get("document"), ["file_unique_id", "mime_type"])
|
||||
self.audio = _ns(data.get("audio"), ["file_unique_id", "mime_type"])
|
||||
self.voice = _ns(data.get("voice"), ["file_unique_id", "mime_type"])
|
||||
self.video_note = _ns(data.get("video_note"), ["file_unique_id"])
|
||||
self.animation = _ns(data.get("animation"), ["file_unique_id"])
|
||||
self.sticker = _ns(data.get("sticker"), ["file_unique_id", "emoji", "is_video"])
|
||||
|
||||
def __str__(self) -> str:
|
||||
return json.dumps(self._snapshot, default=str)
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return self.__str__()
|
||||
|
||||
|
||||
def restore_message(data: dict) -> CachedMessage:
|
||||
return CachedMessage(data)
|
||||
|
||||
|
||||
def snapshot_messages(messages: List[Any]) -> List[dict]:
|
||||
return [snapshot_message(m) for m in messages]
|
||||
|
||||
|
||||
def restore_messages(items: List[dict]) -> List[CachedMessage]:
|
||||
return [restore_message(d) for d in items]
|
||||
@@ -0,0 +1,309 @@
|
||||
"""Tests for message_snapshot.py (issue #23, Задания 4-5,7).
|
||||
|
||||
These verify that a pyrogram Message survives snapshot -> JSON -> restore as a
|
||||
CachedMessage that the render pipeline (post_parser.py) consumes identically.
|
||||
"""
|
||||
import copy
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from pyrogram.enums import MessageMediaType, MessageServiceType
|
||||
|
||||
from message_snapshot import (
|
||||
snapshot_message,
|
||||
restore_message,
|
||||
snapshot_messages,
|
||||
restore_messages,
|
||||
CachedStr,
|
||||
CachedMessage,
|
||||
)
|
||||
from post_parser import PostParser
|
||||
|
||||
|
||||
class FakeStr(str):
|
||||
"""Stand-in for a live pyrogram Str: a str carrying a .html rendering."""
|
||||
def __new__(cls, plain, html):
|
||||
obj = str.__new__(cls, plain)
|
||||
obj._html = html
|
||||
return obj
|
||||
|
||||
@property
|
||||
def html(self):
|
||||
return self._html
|
||||
|
||||
|
||||
def _roundtrip(message):
|
||||
snap = snapshot_message(message)
|
||||
# The snapshot MUST be JSON-serializable (this is the whole point of the rewrite).
|
||||
json.dumps(snap)
|
||||
return restore_message(snap), snap
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Test 1 — top-level round trip.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_roundtrip_basic_fields():
|
||||
msg = SimpleNamespace(
|
||||
id=42,
|
||||
date=datetime(2020, 1, 2, 3, 4, 5),
|
||||
text=FakeStr("hello", "<b>hello</b>"),
|
||||
caption=FakeStr("cap", "<i>cap</i>"),
|
||||
media=MessageMediaType.PHOTO,
|
||||
views=123,
|
||||
media_group_id="mg-1",
|
||||
)
|
||||
restored, _ = _roundtrip(msg)
|
||||
assert restored.id == 42
|
||||
assert restored.date == datetime(2020, 1, 2, 3, 4, 5)
|
||||
assert restored.text.html == "<b>hello</b>"
|
||||
assert str(restored.text) == "hello"
|
||||
assert restored.caption.html == "<i>cap</i>"
|
||||
assert restored.media is MessageMediaType.PHOTO
|
||||
assert restored.views == 123
|
||||
assert restored.media_group_id == "mg-1"
|
||||
|
||||
|
||||
def test_roundtrip_date_naive_and_aware():
|
||||
naive = SimpleNamespace(date=datetime(2021, 5, 5, 10, 0, 0))
|
||||
aware = SimpleNamespace(date=datetime(2021, 5, 5, 10, 0, 0, tzinfo=timezone.utc))
|
||||
|
||||
r_naive, _ = _roundtrip(naive)
|
||||
r_aware, _ = _roundtrip(aware)
|
||||
|
||||
assert r_naive.date.tzinfo is None
|
||||
assert r_aware.date.tzinfo is not None
|
||||
assert r_aware.date == datetime(2021, 5, 5, 10, 0, 0, tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def test_text_falls_back_to_plain_when_no_html():
|
||||
# A bare str (no .html) -> html defaults to the plain text.
|
||||
msg = SimpleNamespace(text="plain only")
|
||||
restored, _ = _roundtrip(msg)
|
||||
assert restored.text.html == "plain only"
|
||||
assert isinstance(restored.text, CachedStr)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Test 2 — polls with FormattedText fakes (must be JSON-serializable).
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_poll_formatted_text_unwrapped():
|
||||
option = SimpleNamespace(text=SimpleNamespace(text="Opt", entities=[]))
|
||||
poll = SimpleNamespace(
|
||||
question=SimpleNamespace(text="Q?", entities=[]),
|
||||
options=[option],
|
||||
)
|
||||
msg = SimpleNamespace(poll=poll)
|
||||
|
||||
snap = snapshot_message(msg)
|
||||
# This MUST NOT raise. If the snapshot stored the FormattedText objects as-is, the
|
||||
# SimpleNamespace values would make json.dumps raise TypeError (regression guard).
|
||||
json.dumps(snap)
|
||||
|
||||
restored = restore_message(snap)
|
||||
assert restored.poll.question == "Q?"
|
||||
assert isinstance(restored.poll.question, str)
|
||||
opt = restored.poll.options[0]
|
||||
assert getattr(opt, "text", "") == "Opt"
|
||||
|
||||
|
||||
def test_poll_snapshot_would_fail_if_formattedtext_kept_raw():
|
||||
# Guard proving the assertion above actually catches raw FormattedText: a snapshot dict
|
||||
# that carried the FormattedText object would not be JSON-serializable.
|
||||
bad_snap = {"poll": {"question": SimpleNamespace(text="Q?"), "options": []}}
|
||||
with pytest.raises(TypeError):
|
||||
json.dumps(bad_snap)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Test 3 — reactions (normal / paid / custom).
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_reactions_normal_paid_custom():
|
||||
normal = SimpleNamespace(emoji="👍", count=5, is_paid=False)
|
||||
paid = SimpleNamespace(count=3, is_paid=True)
|
||||
custom = SimpleNamespace(custom_emoji_id=1234567890, count=2)
|
||||
msg = SimpleNamespace(reactions=SimpleNamespace(reactions=[normal, paid, custom]))
|
||||
|
||||
restored, _ = _roundtrip(msg)
|
||||
r_normal, r_paid, r_custom = restored.reactions.reactions
|
||||
|
||||
assert r_normal.emoji == "👍"
|
||||
assert r_normal.count == 5
|
||||
assert r_normal.is_paid is False
|
||||
assert r_normal.custom_emoji_id is None
|
||||
|
||||
assert r_paid.is_paid is True
|
||||
assert r_paid.count == 3
|
||||
|
||||
# Custom emoji: emoji is null, custom_emoji_id is a STRING; key still present.
|
||||
assert hasattr(r_custom, "emoji") is True
|
||||
assert r_custom.emoji is None
|
||||
assert isinstance(r_custom.custom_emoji_id, str)
|
||||
assert r_custom.custom_emoji_id == "1234567890"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Test 4 — forward_origin Case 1-5 (presence-semantics drives _format_forward_info).
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _forward_html(forward_origin):
|
||||
msg = SimpleNamespace(forward_origin=forward_origin)
|
||||
restored, _ = _roundtrip(msg)
|
||||
return PostParser(None)._format_forward_info(restored)
|
||||
|
||||
|
||||
def test_forward_origin_cases():
|
||||
# Case 1: channel/supergroup (has .chat)
|
||||
html = _forward_html(SimpleNamespace(
|
||||
type="channel",
|
||||
chat=SimpleNamespace(id=-100, title="Chan", username="chanu"),
|
||||
))
|
||||
assert "Chan" in html and "@chanu" in html
|
||||
|
||||
# Case 2: hidden user (sender_user_name)
|
||||
html = _forward_html(SimpleNamespace(type="hidden_user", sender_user_name="Hidden Guy"))
|
||||
assert "Hidden Guy" in html
|
||||
|
||||
# Case 3: regular user (sender_user)
|
||||
html = _forward_html(SimpleNamespace(
|
||||
type="user",
|
||||
sender_user=SimpleNamespace(first_name="Ann", last_name="Bee", username="annbee"),
|
||||
))
|
||||
assert "Ann Bee" in html and "@annbee" in html
|
||||
|
||||
# Case 4: channel without username (chat_id + title, no chat attr)
|
||||
html = _forward_html(SimpleNamespace(type="channel", chat_id=-1002, title="NoUserChan"))
|
||||
assert "NoUserChan" in html
|
||||
|
||||
# Case 5: anything else
|
||||
html = _forward_html(SimpleNamespace(type="something_else"))
|
||||
assert html == '<div class="message-forward">--- Forwarded message ---</div>'
|
||||
|
||||
|
||||
def test_forward_origin_presence_semantics():
|
||||
# Only keys present on the live object are recorded (hasattr semantics).
|
||||
restored, snap = _roundtrip(SimpleNamespace(
|
||||
forward_origin=SimpleNamespace(type="hidden_user", sender_user_name="X")))
|
||||
assert "chat" not in snap["forward_origin"]
|
||||
assert "sender_user" not in snap["forward_origin"]
|
||||
assert not hasattr(restored.forward_origin, "chat")
|
||||
assert hasattr(restored.forward_origin, "sender_user_name")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Test 5 — service restored as string usable in `'X' in str(service)`.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_service_pinned_message():
|
||||
msg = SimpleNamespace(service=MessageServiceType.PINNED_MESSAGE)
|
||||
restored, _ = _roundtrip(msg)
|
||||
assert "PINNED_MESSAGE" in str(restored.service)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Test 6 — mutability + deepcopy (.text.html survives).
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_mutable_and_deepcopy():
|
||||
msg = SimpleNamespace(id=7, text=FakeStr("body", "<u>body</u>"))
|
||||
restored, _ = _roundtrip(msg)
|
||||
|
||||
# Mutable: reply enrichment assigns reply_to_message.
|
||||
sentinel = object()
|
||||
restored.reply_to_message = sentinel
|
||||
assert restored.reply_to_message is sentinel
|
||||
|
||||
clone = copy.deepcopy(restored)
|
||||
assert clone.text.html == "<u>body</u>"
|
||||
assert str(clone.text) == "body"
|
||||
assert clone.id == 7
|
||||
|
||||
|
||||
def test_cachedstr_survives_pickle():
|
||||
import pickle as _pkl
|
||||
s = CachedStr.build("plain", "<b>plain</b>")
|
||||
back = _pkl.loads(_pkl.dumps(s))
|
||||
assert str(back) == "plain"
|
||||
assert back.html == "<b>plain</b>"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Test 7 — unknown media name -> None, no exception.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_unknown_media_type():
|
||||
snap = snapshot_message(SimpleNamespace())
|
||||
snap["media"] = "FUTURE_TYPE"
|
||||
restored = restore_message(snap) # must not raise
|
||||
assert restored.media is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Test 9 — >100MB video is not collected by _save_media_file_ids.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_large_video_not_collected():
|
||||
snap = snapshot_message(SimpleNamespace())
|
||||
snap["media"] = "VIDEO"
|
||||
snap["video"] = {"file_unique_id": "vid1", "file_size": 200 * 1024 * 1024}
|
||||
snap["chat"] = {"id": -1001, "username": "chan", "title": "t", "usernames": None}
|
||||
restored = restore_message(snap)
|
||||
|
||||
pp = PostParser(None)
|
||||
pp._save_media_file_ids(restored)
|
||||
assert pp._pending_media_ids == []
|
||||
|
||||
|
||||
def test_normal_video_is_collected():
|
||||
snap = snapshot_message(SimpleNamespace())
|
||||
snap["id"] = 10
|
||||
snap["media"] = "VIDEO"
|
||||
snap["video"] = {"file_unique_id": "vid2", "file_size": 1024}
|
||||
snap["chat"] = {"id": -1001, "username": "chan", "title": "t", "usernames": None}
|
||||
restored = restore_message(snap)
|
||||
|
||||
pp = PostParser(None)
|
||||
pp._save_media_file_ids(restored)
|
||||
assert len(pp._pending_media_ids) == 1
|
||||
assert pp._pending_media_ids[0][2] == "vid2"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Test 10 — restored chat without username -> None, not AttributeError.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_chat_without_username():
|
||||
snap = snapshot_message(SimpleNamespace())
|
||||
snap["chat"] = {"id": -100, "username": None, "title": "T", "usernames": None}
|
||||
restored = restore_message(snap)
|
||||
assert restored.chat.username is None
|
||||
|
||||
|
||||
def test_chat_usernames_restored_as_objects():
|
||||
chat = SimpleNamespace(
|
||||
id=-100, username=None, title="T",
|
||||
usernames=[SimpleNamespace(username="alt", active=True),
|
||||
SimpleNamespace(username="old", active=False)],
|
||||
)
|
||||
restored, _ = _roundtrip(SimpleNamespace(chat=chat))
|
||||
active = [u.username for u in restored.chat.usernames if u.active]
|
||||
assert active == ["alt"]
|
||||
# get_channel_username uses exactly this path.
|
||||
assert PostParser(None).get_channel_username(restored) == "alt"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Defaults: every pipeline-consumed attribute exists (getattr never raises).
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_empty_message_defaults():
|
||||
restored = restore_message(snapshot_message(SimpleNamespace()))
|
||||
for attr in ["id", "date", "text", "caption", "media", "service", "media_group_id",
|
||||
"views", "show_caption_above_media", "reply_to_message_id",
|
||||
"reply_to_message", "chat", "sender_chat", "from_user", "forward_origin",
|
||||
"reactions", "poll", "web_page", "photo", "video", "document", "audio",
|
||||
"voice", "video_note", "animation", "sticker"]:
|
||||
assert hasattr(restored, attr)
|
||||
assert restored.empty is False
|
||||
|
||||
|
||||
def test_snapshot_messages_list_roundtrip():
|
||||
msgs = [SimpleNamespace(id=i) for i in range(3)]
|
||||
restored = restore_messages(snapshot_messages(msgs))
|
||||
assert [m.id for m in restored] == [0, 1, 2]
|
||||
assert all(isinstance(m, CachedMessage) for m in restored)
|
||||
@@ -0,0 +1,202 @@
|
||||
"""Tests for the tg_cache generic JSON store, prefix-limit, jitter and sweep (issue #23).
|
||||
|
||||
Задания 2 / 6 / 17 / 18, verification item 8, 11, 12.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
import tg_cache
|
||||
from message_snapshot import SNAPSHOT_VERSION
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cache_dir(tmp_path, monkeypatch):
|
||||
d = tmp_path / "tgcache"
|
||||
d.mkdir()
|
||||
monkeypatch.setattr(tg_cache, "CACHE_DIR", str(d))
|
||||
return d
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# _store_entry / _load_entry.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_store_load_roundtrip(cache_dir):
|
||||
path = str(cache_dir / "x.history.json")
|
||||
tg_cache._store_entry(path, {"limit": 5, "messages": [{"id": 1}]})
|
||||
loaded = tg_cache._load_entry(path, max_age_hours=8)
|
||||
assert loaded is not None
|
||||
assert loaded["limit"] == 5
|
||||
assert loaded["messages"] == [{"id": 1}]
|
||||
assert loaded["version"] == SNAPSHOT_VERSION
|
||||
assert 0.8 <= loaded["jitter"] <= 1.0
|
||||
|
||||
|
||||
def test_load_missing_file(cache_dir):
|
||||
assert tg_cache._load_entry(str(cache_dir / "nope.json"), 8) is None
|
||||
|
||||
|
||||
def test_load_expired_ttl(cache_dir):
|
||||
path = str(cache_dir / "x.history.json")
|
||||
tg_cache._store_entry(path, {"data": 1})
|
||||
# max_age_hours=0 -> adjusted max age 0 -> any positive age is expired.
|
||||
assert tg_cache._load_entry(path, max_age_hours=0) is None
|
||||
# Still fresh with a real TTL.
|
||||
assert tg_cache._load_entry(path, max_age_hours=8) is not None
|
||||
|
||||
|
||||
def test_load_version_mismatch(cache_dir):
|
||||
path = str(cache_dir / "x.history.json")
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
json.dump({"version": SNAPSHOT_VERSION + 999, "timestamp": time.time(), "data": 1}, f)
|
||||
assert tg_cache._load_entry(path, 8) is None
|
||||
|
||||
|
||||
def test_load_corrupt_json(cache_dir):
|
||||
path = str(cache_dir / "x.history.json")
|
||||
with open(path, "w", encoding="utf-8") as f:
|
||||
f.write("{not valid json")
|
||||
assert tg_cache._load_entry(path, 8) is None
|
||||
|
||||
|
||||
def test_store_uses_unique_tmp_names(cache_dir, monkeypatch):
|
||||
path = str(cache_dir / "x.history.json")
|
||||
seen = []
|
||||
|
||||
real_replace = os.replace
|
||||
|
||||
def spy_replace(src, dst):
|
||||
seen.append(src)
|
||||
return real_replace(src, dst)
|
||||
|
||||
monkeypatch.setattr(tg_cache.os, "replace", spy_replace)
|
||||
tg_cache._store_entry(path, {"data": 1})
|
||||
tg_cache._store_entry(path, {"data": 2})
|
||||
|
||||
assert len(seen) == 2
|
||||
assert seen[0] != seen[1] # each writer used a distinct tmp path
|
||||
for src in seen:
|
||||
assert src.startswith(path + ".tmp.")
|
||||
# No tmp leftovers.
|
||||
assert [n for n in os.listdir(cache_dir) if ".tmp." in n] == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# cleanup_legacy_cache_files / sweep_tgcache.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_cleanup_legacy_keeps_new_format(cache_dir):
|
||||
(cache_dir / "chan.cache").write_text("x")
|
||||
(cache_dir / "chan_history.cache").write_text("x")
|
||||
(cache_dir / "chan.chatinfo").write_text("x")
|
||||
(cache_dir / "chan.history.json").write_text("x")
|
||||
(cache_dir / "chan.chatinfo.json").write_text("x")
|
||||
|
||||
removed = tg_cache.cleanup_legacy_cache_files()
|
||||
assert removed == 3
|
||||
remaining = set(os.listdir(cache_dir))
|
||||
assert remaining == {"chan.history.json", "chan.chatinfo.json"}
|
||||
|
||||
|
||||
def test_sweep_removes_only_old(cache_dir):
|
||||
old = cache_dir / "old.history.json"
|
||||
new = cache_dir / "new.history.json"
|
||||
old.write_text("x")
|
||||
new.write_text("x")
|
||||
old_time = time.time() - 8 * 86400 # 8 days old, threshold is 7
|
||||
os.utime(old, (old_time, old_time))
|
||||
|
||||
removed = tg_cache.sweep_tgcache(max_age_days=7)
|
||||
assert removed == 1
|
||||
remaining = set(os.listdir(cache_dir))
|
||||
assert remaining == {"new.history.json"}
|
||||
|
||||
|
||||
def test_sweep_removes_orphan_tmp(cache_dir):
|
||||
orphan = cache_dir / "x.history.json.tmp.deadbeef"
|
||||
orphan.write_text("x")
|
||||
old_time = time.time() - 10 * 86400
|
||||
os.utime(orphan, (old_time, old_time))
|
||||
assert tg_cache.sweep_tgcache(max_age_days=7) == 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# History prefix-limit (Задание 6).
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _fake_messages(n):
|
||||
return [SimpleNamespace(id=i) for i in range(n)]
|
||||
|
||||
|
||||
def test_prefix_larger_cache_serves_smaller_request(cache_dir):
|
||||
tg_cache._save_history_to_cache("chan", _fake_messages(100), limit=100)
|
||||
served = tg_cache._get_history_from_cache("chan", limit=50)
|
||||
assert served is not None
|
||||
assert len(served) == 50
|
||||
assert [m.id for m in served] == list(range(50)) # order preserved (newest-first slice)
|
||||
|
||||
|
||||
def test_prefix_smaller_cache_is_miss(cache_dir):
|
||||
tg_cache._save_history_to_cache("chan", _fake_messages(50), limit=50)
|
||||
assert tg_cache._get_history_from_cache("chan", limit=100) is None
|
||||
|
||||
|
||||
def test_prefix_exhausted_channel_serves_larger_request(cache_dir):
|
||||
# Fetched with limit 100 but the channel only has 37 messages -> entire history cached.
|
||||
tg_cache._save_history_to_cache("chan", _fake_messages(37), limit=100)
|
||||
served = tg_cache._get_history_from_cache("chan", limit=200)
|
||||
assert served is not None
|
||||
assert len(served) == 37
|
||||
|
||||
|
||||
def test_save_stores_fetch_limit_not_len(cache_dir):
|
||||
tg_cache._save_history_to_cache("chan", _fake_messages(37), limit=100)
|
||||
path = tg_cache._cache_file_path("chan", "history.json")
|
||||
payload = tg_cache._load_entry(path, max_age_hours=8)
|
||||
assert payload["limit"] == 100
|
||||
assert len(payload["messages"]) == 37
|
||||
|
||||
|
||||
def test_history_ttl_independent_of_prefix(cache_dir):
|
||||
tg_cache._save_history_to_cache("chan", _fake_messages(100), limit=100)
|
||||
# Expired regardless of prefix match.
|
||||
assert tg_cache._get_history_from_cache("chan", limit=50, max_age_hours=0) is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Jitter stability (Задание 17 / item 12).
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_jitter_read_is_stable(cache_dir):
|
||||
path = str(cache_dir / "x.history.json")
|
||||
tg_cache._store_entry(path, {"data": 1})
|
||||
# Repeated reads near a TTL boundary must give a stable result: the jitter is fixed at
|
||||
# write time and _load_entry never calls random().
|
||||
results = [tg_cache._load_entry(path, max_age_hours=8) is not None for _ in range(10)]
|
||||
assert all(results)
|
||||
|
||||
|
||||
def test_load_does_not_call_random(cache_dir, monkeypatch):
|
||||
path = str(cache_dir / "x.history.json")
|
||||
tg_cache._store_entry(path, {"data": 1})
|
||||
|
||||
def boom(*a, **k):
|
||||
raise AssertionError("random.uniform must not be called at read time")
|
||||
|
||||
monkeypatch.setattr(tg_cache.random, "uniform", boom)
|
||||
assert tg_cache._load_entry(path, max_age_hours=8) is not None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# chatinfo store round trip.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def test_chatinfo_roundtrip(cache_dir):
|
||||
tg_cache._save_chat_to_cache("chan", {"id": -100, "title": "T", "username": "u"})
|
||||
data = tg_cache._get_chat_from_cache("chan")
|
||||
assert data == {"id": -100, "title": "T", "username": "u"}
|
||||
|
||||
|
||||
def test_no_pickle_import():
|
||||
import pathlib
|
||||
src = pathlib.Path(tg_cache.__file__).read_text()
|
||||
assert "pickle" not in src
|
||||
+191
-101
@@ -11,18 +11,22 @@
|
||||
|
||||
import os
|
||||
import json
|
||||
import pickle
|
||||
import uuid
|
||||
import logging
|
||||
import random
|
||||
import asyncio
|
||||
import time
|
||||
from datetime import datetime, timedelta
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, Optional, Union, List
|
||||
from pyrogram import Client
|
||||
from pyrogram.types import Chat, Message
|
||||
from pyrogram.types import Message
|
||||
from tg_throttle import tg_rpc_bounded
|
||||
from config import get_settings
|
||||
from message_snapshot import (
|
||||
SNAPSHOT_VERSION,
|
||||
snapshot_messages,
|
||||
restore_messages,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -40,105 +44,148 @@ try:
|
||||
except ValueError:
|
||||
CHAT_CACHE_TTL_HOURS = 12
|
||||
|
||||
def _get_history_cache_file_path(channel_id: Union[str, int]) -> str:
|
||||
"""Returns path to the message history cache file for the channel"""
|
||||
|
||||
def _safe_key(key: Union[str, int]) -> str:
|
||||
"""Sanitize a channel id/username into a filesystem-safe basename component."""
|
||||
return str(key).replace('/', '_').replace('\\', '_')
|
||||
|
||||
|
||||
def _cache_file_path(key: Union[str, int], suffix: str) -> str:
|
||||
"""Return the cache file path <safe_key>.<suffix> (e.g. 'history.json' / 'chatinfo.json')."""
|
||||
if not os.path.exists(CACHE_DIR):
|
||||
os.makedirs(CACHE_DIR, exist_ok=True)
|
||||
logger.info(f"cache_dir_created: path {CACHE_DIR}")
|
||||
# Convert to string for uniformity
|
||||
channel_id_str = str(channel_id)
|
||||
# Replace potentially problematic characters
|
||||
safe_filename = channel_id_str.replace('/', '_').replace('\\', '_')
|
||||
return os.path.join(CACHE_DIR, f"{safe_filename}.cache")
|
||||
return os.path.join(CACHE_DIR, f"{_safe_key(key)}.{suffix}")
|
||||
|
||||
def _save_history_to_cache(channel_id: Union[str, int], messages: List[Message], limit: int) -> None:
|
||||
"""Saves message history to cache"""
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Generic JSON entry store.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _store_entry(path: str, payload: dict) -> None:
|
||||
"""Atomically write {version, timestamp, jitter, **payload} as JSON to ``path``.
|
||||
|
||||
The document is written to a unique '<path>.tmp.<uuid4>' and os.replace()d into place
|
||||
so a concurrent reader never observes a half-written file. The unique per-writer tmp
|
||||
name means two concurrent writers to the same path do not clobber each other's temp
|
||||
file. This writer's own tmp file is always removed in the finally block.
|
||||
"""
|
||||
tmp_path = f"{path}.tmp.{uuid.uuid4().hex}"
|
||||
entry = {
|
||||
'version': SNAPSHOT_VERSION,
|
||||
'timestamp': time.time(),
|
||||
# Per-write TTL jitter (§17): decided ONCE at write time, so repeated reads of the
|
||||
# same file are stable (the reader never calls random()).
|
||||
'jitter': random.uniform(0.8, 1.0),
|
||||
}
|
||||
entry.update(payload)
|
||||
try:
|
||||
cache_file = _get_history_cache_file_path(channel_id)
|
||||
|
||||
# Create cache metadata — store messages directly (no inner pickle.dumps)
|
||||
cache_data = {
|
||||
'timestamp': time.time(),
|
||||
'limit': limit,
|
||||
'messages': messages
|
||||
}
|
||||
|
||||
with open(cache_file, 'wb') as f:
|
||||
pickle.dump(cache_data, f)
|
||||
|
||||
with open(tmp_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(entry, f)
|
||||
os.replace(tmp_path, path)
|
||||
finally:
|
||||
# Remove our own leftover tmp file if os.replace didn't consume it (e.g. it raised
|
||||
# or json.dump failed). Never touches another writer's uniquely-named tmp file.
|
||||
if os.path.exists(tmp_path):
|
||||
try:
|
||||
os.remove(tmp_path)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _load_entry(path: str, max_age_hours: float) -> Optional[dict]:
|
||||
"""Return the stored payload dict, or None on missing / version mismatch / expired / bad JSON.
|
||||
|
||||
TTL uses the jitter written into the entry (no random() at read time) so repeated
|
||||
reads near the boundary give a stable result.
|
||||
"""
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
try:
|
||||
with open(path, 'r', encoding='utf-8') as f:
|
||||
entry = json.load(f)
|
||||
except (json.JSONDecodeError, OSError, ValueError) as e:
|
||||
logger.warning(f"cache_entry_read_error: path {path}, error {str(e)}")
|
||||
return None
|
||||
|
||||
if not isinstance(entry, dict) or entry.get('version') != SNAPSHOT_VERSION:
|
||||
logger.info(f"cache_entry_version_mismatch: path {path}")
|
||||
return None
|
||||
|
||||
timestamp = entry.get('timestamp')
|
||||
if not isinstance(timestamp, (int, float)):
|
||||
return None
|
||||
age = time.time() - timestamp
|
||||
adjusted_max_age = max_age_hours * 3600 * entry.get('jitter', 1.0)
|
||||
if age > adjusted_max_age:
|
||||
logger.info(f"cache_entry_expired: path {path}, age {age:.1f}s > adjusted max {adjusted_max_age:.1f}s")
|
||||
return None
|
||||
return entry
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# History cache.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _save_history_to_cache(channel_id: Union[str, int], messages: List[Message], limit: int) -> None:
|
||||
"""Save message history (as JSON snapshots) to cache. Stores the fetch limit, not len()."""
|
||||
try:
|
||||
cache_file = _cache_file_path(channel_id, 'history.json')
|
||||
payload = {'limit': limit, 'messages': snapshot_messages(messages)}
|
||||
_store_entry(cache_file, payload)
|
||||
logger.info(f"history_cache_saved: channel {channel_id}, limit {limit}, messages {len(messages)}, file {cache_file}")
|
||||
except Exception as e:
|
||||
logger.error(f"history_cache_save_error: channel {channel_id}, limit {limit}, error {str(e)}")
|
||||
|
||||
|
||||
def _get_history_from_cache(channel_id: Union[str, int], limit: int, max_age_hours: int = 8) -> Optional[List[Message]]:
|
||||
"""
|
||||
Retrieves message history from cache if not older than specified age and matches the limit
|
||||
|
||||
Args:
|
||||
channel_id: Channel ID or username
|
||||
limit: Required message limit
|
||||
max_age_hours: Maximum cache age in hours (default 8 hours)
|
||||
|
||||
Returns:
|
||||
List of messages or None if cache not found, expired or limit doesn't match
|
||||
Retrieve message history from cache if fresh and the cached fetch covers ``limit``.
|
||||
|
||||
Messages are stored newest-first. A cached entry fetched with an equal-or-larger limit
|
||||
serves a smaller request by slicing (prefix). A smaller cached fetch is a miss UNLESS
|
||||
the channel is exhausted (fewer messages exist than were asked for), in which case the
|
||||
cache already holds the entire recent history and can serve any larger request.
|
||||
"""
|
||||
try:
|
||||
cache_file = _get_history_cache_file_path(channel_id)
|
||||
|
||||
if not os.path.exists(cache_file):
|
||||
logger.info(f"history_cache_miss: channel {channel_id}, limit {limit}, cache file not found")
|
||||
cache_file = _cache_file_path(channel_id, 'history.json')
|
||||
payload = _load_entry(cache_file, max_age_hours)
|
||||
if payload is None:
|
||||
logger.info(f"history_cache_miss: channel {channel_id}, limit {limit}")
|
||||
return None
|
||||
|
||||
with open(cache_file, 'rb') as f:
|
||||
cache_data = pickle.load(f)
|
||||
|
||||
# Check cache age with randomization
|
||||
cache_age = time.time() - cache_data['timestamp']
|
||||
# Add randomness up to 20% of max_age_seconds
|
||||
random_factor = 1 - random.uniform(0, 0.2)
|
||||
adjusted_max_age = max_age_hours * 3600 * random_factor
|
||||
|
||||
if cache_age > adjusted_max_age:
|
||||
logger.info(f"history_cache_expired: channel {channel_id}, limit {limit}, age {cache_age:.1f}s > adjusted max {adjusted_max_age:.1f}s (random factor: {random_factor:.2f})")
|
||||
|
||||
cached_limit = payload.get('limit', 0)
|
||||
raw_messages = payload['messages']
|
||||
# Serve when fetched with an equal-or-larger limit, OR when the channel is exhausted
|
||||
# (fewer messages exist than asked -> cache holds entire recent history).
|
||||
if cached_limit < limit and len(raw_messages) >= cached_limit:
|
||||
logger.info(f"history_cache_limit_short: channel {channel_id}, cached limit {cached_limit}, requested {limit}")
|
||||
return None
|
||||
|
||||
# Check if limit matches
|
||||
cached_limit = cache_data.get('limit', 0)
|
||||
if cached_limit != limit:
|
||||
logger.info(f"history_cache_limit_mismatch: channel {channel_id}, cached limit {cached_limit}, requested limit {limit}")
|
||||
return None
|
||||
|
||||
# Restore message list; handle old cache files that used double-pickle (bytes = old format)
|
||||
raw = cache_data['messages']
|
||||
if isinstance(raw, bytes):
|
||||
messages = pickle.loads(raw)
|
||||
else:
|
||||
messages = raw
|
||||
logger.info(f"history_cache_hit: channel {channel_id}, limit {limit}, messages {len(messages)}, age {cache_age:.1f}s")
|
||||
|
||||
messages = restore_messages(raw_messages[:limit])
|
||||
logger.info(f"history_cache_hit: channel {channel_id}, served {limit} of cached {cached_limit}, messages {len(messages)}")
|
||||
return messages
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"history_cache_read_error: channel {channel_id}, limit {limit}, error {str(e)}")
|
||||
# In case of cache read error, better return None and request fresh data
|
||||
return None
|
||||
|
||||
|
||||
async def cached_get_chat_history(client: Client, channel_id: Union[str, int], limit: int = 20) -> List[Message]:
|
||||
"""
|
||||
Gets chat message history with caching.
|
||||
|
||||
|
||||
Args:
|
||||
client: Pyrogram client
|
||||
channel_id: Channel ID or username
|
||||
limit: Maximum number of messages to retrieve
|
||||
|
||||
|
||||
Returns:
|
||||
List of messages, same as original client.get_chat_history()
|
||||
List of messages, same as original client.get_chat_history(). On a cache miss the
|
||||
live pyrogram Messages are returned; on a hit, restored CachedMessage objects.
|
||||
"""
|
||||
cached_messages = await asyncio.to_thread(_get_history_from_cache, channel_id, limit)
|
||||
|
||||
|
||||
if cached_messages is not None:
|
||||
return cached_messages
|
||||
|
||||
|
||||
try:
|
||||
logger.info(f"history_cache_request: fetching fresh history for channel {channel_id}, limit {limit}")
|
||||
# Hold the global RPC gate for the live fetch and bound the RPC body with the
|
||||
@@ -148,57 +195,39 @@ async def cached_get_chat_history(client: Client, channel_id: Union[str, int], l
|
||||
async with tg_rpc_bounded(Config["tg_rpc_timeout"]):
|
||||
messages = [m async for m in client.get_chat_history(channel_id, limit=limit)]
|
||||
await asyncio.to_thread(_save_history_to_cache, channel_id, messages, limit)
|
||||
|
||||
|
||||
return messages
|
||||
except Exception as e:
|
||||
logger.error(f"history_cache_request_error: channel {channel_id}, limit {limit}, error {str(e)}")
|
||||
raise
|
||||
|
||||
|
||||
def _get_chat_cache_file_path(channel_id: Union[str, int]) -> str:
|
||||
"""Returns path to the channel-info cache file for the channel."""
|
||||
if not os.path.exists(CACHE_DIR):
|
||||
os.makedirs(CACHE_DIR, exist_ok=True)
|
||||
logger.info(f"cache_dir_created: path {CACHE_DIR}")
|
||||
channel_id_str = str(channel_id)
|
||||
safe_filename = channel_id_str.replace('/', '_').replace('\\', '_')
|
||||
# Distinct suffix so channel-info files never collide with history (.cache) files.
|
||||
return os.path.join(CACHE_DIR, f"{safe_filename}.chatinfo")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Channel-info cache.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def _save_chat_to_cache(channel_id: Union[str, int], data: dict) -> None:
|
||||
"""Saves channel-info (id/title/username) to cache."""
|
||||
try:
|
||||
cache_file = _get_chat_cache_file_path(channel_id)
|
||||
cache_data = {'timestamp': time.time(), 'data': data}
|
||||
with open(cache_file, 'wb') as f:
|
||||
pickle.dump(cache_data, f)
|
||||
cache_file = _cache_file_path(channel_id, 'chatinfo.json')
|
||||
_store_entry(cache_file, {'data': data})
|
||||
logger.info(f"chatinfo_cache_saved: channel {channel_id}, file {cache_file}")
|
||||
except Exception as e:
|
||||
logger.error(f"chatinfo_cache_save_error: channel {channel_id}, error {str(e)}")
|
||||
|
||||
|
||||
def _get_chat_from_cache(channel_id: Union[str, int], max_age_hours: int = CHAT_CACHE_TTL_HOURS) -> Optional[dict]:
|
||||
"""Retrieves channel-info from cache if not older than max_age_hours (with up-to-20% randomization)."""
|
||||
"""Retrieves channel-info from cache if fresh (jittered TTL)."""
|
||||
try:
|
||||
cache_file = _get_chat_cache_file_path(channel_id)
|
||||
if not os.path.exists(cache_file):
|
||||
logger.info(f"chatinfo_cache_miss: channel {channel_id}, cache file not found")
|
||||
cache_file = _cache_file_path(channel_id, 'chatinfo.json')
|
||||
payload = _load_entry(cache_file, max_age_hours)
|
||||
if payload is None:
|
||||
logger.info(f"chatinfo_cache_miss: channel {channel_id}")
|
||||
return None
|
||||
with open(cache_file, 'rb') as f:
|
||||
cache_data = pickle.load(f)
|
||||
cache_age = time.time() - cache_data['timestamp']
|
||||
# Add randomness up to 20% of max age, same approach as the history cache.
|
||||
random_factor = 1 - random.uniform(0, 0.2)
|
||||
adjusted_max_age = max_age_hours * 3600 * random_factor
|
||||
if cache_age > adjusted_max_age:
|
||||
logger.info(f"chatinfo_cache_expired: channel {channel_id}, age {cache_age:.1f}s > adjusted max {adjusted_max_age:.1f}s (random factor: {random_factor:.2f})")
|
||||
return None
|
||||
data = cache_data.get('data')
|
||||
data = payload.get('data')
|
||||
if not isinstance(data, dict):
|
||||
logger.warning(f"chatinfo_cache_invalid: channel {channel_id}, unexpected payload type {type(data).__name__}")
|
||||
return None
|
||||
logger.info(f"chatinfo_cache_hit: channel {channel_id}, age {cache_age:.1f}s")
|
||||
logger.info(f"chatinfo_cache_hit: channel {channel_id}")
|
||||
return data
|
||||
except Exception as e:
|
||||
logger.error(f"chatinfo_cache_read_error: channel {channel_id}, error {str(e)}")
|
||||
@@ -228,3 +257,64 @@ async def cached_get_chat(client: Client, channel_id: Union[str, int]) -> Simple
|
||||
}
|
||||
await asyncio.to_thread(_save_chat_to_cache, channel_id, data)
|
||||
return SimpleNamespace(**data)
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Maintenance: legacy cleanup + age sweep.
|
||||
# --------------------------------------------------------------------------- #
|
||||
def cleanup_legacy_cache_files() -> int:
|
||||
"""Delete legacy binary cache files (*.cache incl. *_history.cache, and *.chatinfo).
|
||||
|
||||
The new store uses *.history.json / *.chatinfo.json, so these old-format files are
|
||||
dead weight and would otherwise never be reclaimed. Returns the number removed.
|
||||
"""
|
||||
removed = 0
|
||||
if not os.path.isdir(CACHE_DIR):
|
||||
return 0
|
||||
try:
|
||||
names = os.listdir(CACHE_DIR)
|
||||
except OSError as e:
|
||||
logger.warning(f"cleanup_legacy_list_error: dir {CACHE_DIR}, error {str(e)}")
|
||||
return 0
|
||||
for name in names:
|
||||
if name.endswith('.cache') or name.endswith('.chatinfo'):
|
||||
try:
|
||||
os.remove(os.path.join(CACHE_DIR, name))
|
||||
removed += 1
|
||||
except OSError as e:
|
||||
logger.warning(f"cleanup_legacy_remove_error: file {name}, error {str(e)}")
|
||||
if removed:
|
||||
logger.info(f"cleanup_legacy_cache_files: removed {removed} legacy files from {CACHE_DIR}")
|
||||
return removed
|
||||
|
||||
|
||||
def sweep_tgcache(max_age_days: int = 7) -> int:
|
||||
"""Delete files in CACHE_DIR whose mtime is older than ``max_age_days``.
|
||||
|
||||
Reclaims cache for dead channels and orphaned uuid tmp files. A race with an in-flight
|
||||
writer is POSSIBLE (stat -> unlink is not atomic against os.replace) but harmless: the
|
||||
worst outcome is one extra cache miss. This is NOT an atomicity guarantee. Returns the
|
||||
number of files removed.
|
||||
"""
|
||||
removed = 0
|
||||
if not os.path.isdir(CACHE_DIR):
|
||||
return 0
|
||||
cutoff = time.time() - max_age_days * 86400
|
||||
try:
|
||||
names = os.listdir(CACHE_DIR)
|
||||
except OSError as e:
|
||||
logger.warning(f"sweep_tgcache_list_error: dir {CACHE_DIR}, error {str(e)}")
|
||||
return 0
|
||||
for name in names:
|
||||
path = os.path.join(CACHE_DIR, name)
|
||||
try:
|
||||
if not os.path.isfile(path):
|
||||
continue
|
||||
if os.path.getmtime(path) < cutoff:
|
||||
os.remove(path)
|
||||
removed += 1
|
||||
except OSError as e:
|
||||
logger.warning(f"sweep_tgcache_remove_error: file {name}, error {str(e)}")
|
||||
if removed:
|
||||
logger.info(f"sweep_tgcache: removed {removed} stale files (> {max_age_days}d) from {CACHE_DIR}")
|
||||
return removed
|
||||
|
||||
Reference in New Issue
Block a user