Spaces:
Paused
Paused
| """Putting a backup made by /export back into the bot (/import). | |
| The backup is a JSON file: {"exported_at", "channels": [...], "posts": [...], "ignored": [...]}. Importing only ADDS: | |
| * a channel the bot already has is left as it is (its title, username and access come from Telegram, which is newer | |
| than the file); a channel it does not have is added; | |
| * a post at a message id the bot already has is left as it is; every other published post is added, and so is every | |
| draft that is not already here; | |
| * a post the owner has made the bot forget (/posts -> Only forget it in the bot) is not brought back; | |
| * nothing is ever deleted or overwritten. | |
| The file is read as data only: every value is checked before it goes near the database. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import logging | |
| from dataclasses import dataclass, field | |
| from datetime import datetime, timezone | |
| from typing import Optional | |
| from sqlalchemy import select | |
| from sqlalchemy.exc import IntegrityError | |
| from .db import Channel, Database, IgnoredPost, Post, as_utc, utcnow | |
| log = logging.getLogger(__name__) | |
| MAX_BYTES = 40 * 1024 * 1024 # a backup of tens of thousands of posts is a few megabytes | |
| CHUNK = 500 # rows saved in one go | |
| POST_KINDS = ("draft", "sent") | |
| SOURCES = ("bot", "adopted") | |
| class BackupError(ValueError): | |
| """The file can't be used; the message says why (and is safe to show).""" | |
| # ------------------------------------------------------------------------------------------------ reading | |
| def _int(value, what: str) -> int: | |
| if isinstance(value, bool) or not isinstance(value, int): | |
| raise ValueError(f"{what} is not a whole number") | |
| return value | |
| def _text(value, what: str, limit: Optional[int] = None) -> str: | |
| if value is None: | |
| return "" | |
| if not isinstance(value, str): | |
| raise ValueError(f"{what} is not text") | |
| return value[:limit] if limit else value | |
| def _when(value, what: str, default: Optional[datetime] = None) -> Optional[datetime]: | |
| if value in (None, ""): | |
| return default | |
| if not isinstance(value, str): | |
| raise ValueError(f"{what} is not a date") | |
| try: | |
| d = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) | |
| except ValueError: | |
| raise ValueError(f"{what} is not a date") | |
| return as_utc(d) | |
| def _rows_of(value, what: str) -> list: | |
| if value is None: | |
| return [] | |
| if not isinstance(value, list): | |
| raise BackupError(f"“{what}” should be a list.") | |
| return value | |
| def channel_row(row) -> dict: | |
| """One channel of the file as the values for a Channel row (ValueError if it is not usable).""" | |
| if not isinstance(row, dict): | |
| raise ValueError("not an object") | |
| cid = _int(row.get("id"), "id") | |
| if cid == 0: | |
| raise ValueError("id is 0") | |
| username = _text(row.get("username"), "username", 64).strip().lstrip("@") or None | |
| created = _when(row.get("created_at"), "created_at", utcnow()) | |
| return dict( | |
| id=cid, | |
| access_hash=_int(row.get("access_hash", 0) or 0, "access_hash"), | |
| title=_text(row.get("title"), "title", 256), | |
| username=username, | |
| active=bool(row.get("active", True)), | |
| added_by=_int(row.get("added_by", 0) or 0, "added_by"), | |
| created_at=created, | |
| ) | |
| def _buttons(value) -> list: | |
| if value is None: | |
| return [] | |
| if not isinstance(value, list): | |
| raise ValueError("buttons are not a list") | |
| for line in value: | |
| if not isinstance(line, list) or not all(isinstance(b, dict) for b in line): | |
| raise ValueError("buttons are not rows of buttons") | |
| return value | |
| def _entities(value) -> list: | |
| if value is None: | |
| return [] | |
| if not isinstance(value, list) or not all(isinstance(e, dict) for e in value): | |
| raise ValueError("formatting is not a list") | |
| return value | |
| def post_row(row) -> dict: | |
| """One post of the file as the values for a Post row (ValueError if it is not usable).""" | |
| if not isinstance(row, dict): | |
| raise ValueError("not an object") | |
| cid = _int(row.get("channel_id"), "channel_id") | |
| mid = row.get("message_id") | |
| mid = None if mid is None else _int(mid, "message_id") | |
| status = row.get("status") or ("sent" if mid else "draft") | |
| if status not in POST_KINDS: | |
| raise ValueError(f"status “{str(status)[:20]}” is unknown") | |
| if status == "sent" and (mid is None or mid < 1): | |
| raise ValueError("a published post without a message id") | |
| if status == "draft": | |
| mid = None | |
| source = row.get("source") if row.get("source") in SOURCES else "bot" | |
| kind = row.get("media_kind") | |
| if kind is not None and (not isinstance(kind, str) or len(kind) > 16): | |
| raise ValueError("media_kind is not usable") | |
| ref = row.get("media_file_id") | |
| if ref is not None and not isinstance(ref, str): | |
| raise ValueError("media_file_id is not text") | |
| created = _when(row.get("created_at"), "created_at", utcnow()) | |
| return dict( | |
| channel_id=cid, | |
| message_id=mid, | |
| status=status, | |
| source=source, | |
| text=_text(row.get("text"), "text"), | |
| entities=_entities(row.get("entities")), | |
| media_kind=kind or None, | |
| media_file_id=ref or None, | |
| buttons=_buttons(row.get("buttons")), | |
| link_preview=bool(row.get("link_preview")), | |
| created_by=_int(row.get("created_by", 0) or 0, "created_by"), | |
| created_at=created, | |
| updated_at=_when(row.get("updated_at"), "updated_at", created), | |
| sent_at=_when(row.get("sent_at"), "sent_at"), | |
| ) | |
| class Backup: | |
| exported_at: Optional[str] = None | |
| channels: list = field(default_factory=list) # dicts for Channel(**d) | |
| posts: list = field(default_factory=list) # dicts for Post(**d) | |
| ignored: list = field(default_factory=list) # (channel id, message id) | |
| bad_rows: int = 0 # entries of the file that were not usable and are left out | |
| def parse_backup(raw: bytes) -> Backup: | |
| """The file's bytes -> a Backup. BackupError (with a message for the person) when it is not an /export file.""" | |
| if len(raw) > MAX_BYTES: | |
| raise BackupError(f"it is larger than {MAX_BYTES // (1024 * 1024)} MB.") | |
| try: | |
| data = json.loads(raw.decode("utf-8-sig")) | |
| except (UnicodeDecodeError, ValueError): | |
| raise BackupError("it is not a JSON file. Send the file that /export made, unchanged.") | |
| if not isinstance(data, dict) or ("channels" not in data and "posts" not in data): | |
| raise BackupError("it does not look like a backup from /export (no channels and posts in it).") | |
| bk = Backup(exported_at=data.get("exported_at") if isinstance(data.get("exported_at"), str) else None) | |
| for row in _rows_of(data.get("channels"), "channels"): | |
| try: | |
| bk.channels.append(channel_row(row)) | |
| except (ValueError, TypeError): | |
| bk.bad_rows += 1 | |
| for row in _rows_of(data.get("posts"), "posts"): | |
| try: | |
| bk.posts.append(post_row(row)) | |
| except (ValueError, TypeError): | |
| bk.bad_rows += 1 | |
| for row in _rows_of(data.get("ignored"), "ignored"): | |
| try: | |
| if not isinstance(row, dict): | |
| raise ValueError("not an object") | |
| bk.ignored.append((_int(row.get("channel_id"), "channel_id"), _int(row.get("message_id"), "message_id"))) | |
| except (ValueError, TypeError): | |
| bk.bad_rows += 1 | |
| if not (bk.channels or bk.posts): | |
| raise BackupError("there is nothing in it that can be imported." + (" (All entries were unreadable.)" if bk.bad_rows else "")) | |
| return bk | |
| # ------------------------------------------------------------------------------------------------ the plan | |
| class ImportPlan: | |
| backup: Backup | |
| new_channels: list = field(default_factory=list) # channel dicts that are not in the bot yet | |
| channels_present: int = 0 # channels of the file the bot has already (left as they are) | |
| new_posts: list = field(default_factory=list) # published posts and drafts to add | |
| posts_present: int = 0 # published posts at a message id the bot already has | |
| drafts_present: int = 0 # drafts that are here already | |
| posts_forgotten: int = 0 # posts the owner made the bot forget: not brought back | |
| posts_orphaned: int = 0 # posts of a channel that is neither in the bot nor in the file | |
| new_ignored: list = field(default_factory=list) # "forgotten" marks to add | |
| def sent_to_add(self) -> int: | |
| return sum(1 for p in self.new_posts if p["status"] == "sent") | |
| def drafts_to_add(self) -> int: | |
| return sum(1 for p in self.new_posts if p["status"] == "draft") | |
| def nothing_new(self) -> bool: | |
| return not (self.new_channels or self.new_posts or self.new_ignored) | |
| async def plan_import(db: Database, backup: Backup) -> ImportPlan: | |
| """What importing would do - nothing is written.""" | |
| plan = ImportPlan(backup=backup) | |
| async with db.Session() as s: | |
| have_channels = {int(c) for c in (await s.execute(select(Channel.id))).scalars()} | |
| known_marks = {(int(c), int(m)) for c, m in (await s.execute(select(IgnoredPost.channel_id, IgnoredPost.message_id))).all()} | |
| wanted = {p["channel_id"] for p in backup.posts} | |
| have_posts: set = set() | |
| have_drafts: set = set() | |
| if wanted: | |
| q = select(Post.channel_id, Post.message_id).where(Post.message_id.is_not(None)) | |
| have_posts = {(int(c), int(m)) for c, m in (await s.execute(q)).all() if int(c) in wanted} | |
| q = select(Post.channel_id, Post.text, Post.created_at).where(Post.status == "draft") | |
| have_drafts = {(int(c), t, as_utc(at)) for c, t, at in (await s.execute(q)).all() if int(c) in wanted} | |
| in_file: dict = {} | |
| for ch in backup.channels: | |
| if ch["id"] in have_channels: | |
| plan.channels_present += 1 | |
| elif ch["id"] not in in_file: | |
| in_file[ch["id"]] = ch | |
| plan.new_channels.append(ch) | |
| marks = set(known_marks) | |
| for mark in backup.ignored: | |
| if mark[0] in have_channels or mark[0] in in_file: | |
| if mark not in marks: | |
| marks.add(mark) | |
| plan.new_ignored.append(mark) | |
| seen_posts: set = set() | |
| seen_drafts: set = set() | |
| for p in backup.posts: | |
| cid = p["channel_id"] | |
| if cid not in have_channels and cid not in in_file: | |
| plan.posts_orphaned += 1 | |
| elif p["status"] == "sent": | |
| key = (cid, p["message_id"]) | |
| if key in marks: | |
| plan.posts_forgotten += 1 | |
| elif key in have_posts or key in seen_posts: | |
| plan.posts_present += 1 | |
| else: | |
| seen_posts.add(key) | |
| plan.new_posts.append(p) | |
| else: | |
| key = (cid, p["text"], p["created_at"]) | |
| if key in have_drafts or key in seen_drafts: | |
| plan.drafts_present += 1 | |
| else: | |
| seen_drafts.add(key) | |
| plan.new_posts.append(p) | |
| return plan | |
| # ----------------------------------------------------------------------------------------------- doing it | |
| class ImportResult: | |
| channels_added: int = 0 | |
| posts_added: int = 0 | |
| drafts_added: int = 0 | |
| marks_added: int = 0 | |
| skipped_meanwhile: int = 0 # posts that turned up in the bot while the import was running | |
| async def _save_posts(db: Database, rows: list) -> list: | |
| """Save these posts; returns the ones that were saved (one that is there already is skipped, never an error).""" | |
| async with db.Session() as s: | |
| s.add_all([Post(**r) for r in rows]) | |
| try: | |
| await s.commit() | |
| return list(rows) | |
| except IntegrityError: | |
| await s.rollback() | |
| saved = [] | |
| for r in rows: # one is there already: one by one | |
| async with db.Session() as s: | |
| s.add(Post(**r)) | |
| try: | |
| await s.commit() | |
| saved.append(r) | |
| except IntegrityError: | |
| await s.rollback() | |
| return saved | |
| async def apply_import(db: Database, plan: ImportPlan, progress=None) -> ImportResult: | |
| """Write the plan: channels first (posts point at them), then the posts in pieces, then the "forgotten" marks.""" | |
| res = ImportResult() | |
| if plan.new_channels: | |
| async with db.Session() as s: | |
| for d in plan.new_channels: | |
| if await s.get(Channel, d["id"]) is None: | |
| s.add(Channel(**d)) | |
| res.channels_added += 1 | |
| await s.commit() | |
| rows = plan.new_posts | |
| for i in range(0, len(rows), CHUNK): | |
| part = rows[i : i + CHUNK] | |
| saved = await _save_posts(db, part) | |
| res.posts_added += sum(1 for r in saved if r["status"] == "sent") | |
| res.drafts_added += sum(1 for r in saved if r["status"] == "draft") | |
| res.skipped_meanwhile += len(part) - len(saved) | |
| if progress is not None: | |
| await progress(min(i + CHUNK, len(rows)), len(rows)) | |
| for cid, mid in plan.new_ignored: | |
| await db.ignore_post(cid, mid) | |
| res.marks_added += 1 | |
| return res | |
| def stamp(iso: Optional[str]) -> str: | |
| """'2026-10-07T10:00:00+00:00' -> '07 Oct 2026, 10:00 UTC' (the text as it is when it can't be read).""" | |
| if not iso: | |
| return "" | |
| try: | |
| d = as_utc(datetime.fromisoformat(iso.replace("Z", "+00:00"))) | |
| except ValueError: | |
| return iso[:40] | |
| return d.astimezone(timezone.utc).strftime("%d %b %Y, %H:%M UTC") | |