From 6e3966aeb53c33656053366c67e1aec9882d9136 Mon Sep 17 00:00:00 2001 From: mr-forust Date: Sat, 27 Jun 2026 15:20:09 +0200 Subject: [PATCH] feat: add multi-account data layer with legacy migration - GlobalStateStore for account registry in global state.json - AccountStateStore for per-account settings (channels, credentials, continuous config) - migrate_legacy_state() copies legacy data/ to data/accounts/default/ with progress logging; renames session.session to default.session - Helpers: account_data_dir(), account_session_path(), list_accounts(), load_account(), save_account(), get_account_store(), get_global_store() - StateStore accepts defaults parameter for flexible default values --- app_state.py | 244 +++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 226 insertions(+), 18 deletions(-) diff --git a/app_state.py b/app_state.py index 6d27411..3d6a910 100644 --- a/app_state.py +++ b/app_state.py @@ -1,11 +1,23 @@ import json +import logging +import shutil import threading import time +from copy import deepcopy from pathlib import Path -from typing import Any, Callable, Dict, Optional +from typing import Any, Callable, Dict, List, Optional +logger = logging.getLogger(__name__) -DEFAULT_STATE: Dict[str, Any] = { +# ── defaults ────────────────────────────────────────────────────────── + +GLOBAL_DEFAULTS: Dict[str, Any] = { + "accounts": [], + "version": 2, +} + +ACCOUNT_DEFAULTS: Dict[str, Any] = { + "label": "", "api_id": None, "api_hash": None, "channels": {}, @@ -20,10 +32,15 @@ DEFAULT_STATE: Dict[str, Any] = { }, } +# ── StateStore (thread-safe JSON) ───────────────────────────────────── + class StateStore: - def __init__(self, path: Path): + """Thread-safe JSON state store with TTL cache.""" + + def __init__(self, path: Path, defaults: Optional[Dict[str, Any]] = None): self.path = path + self.defaults = defaults or {} self.lock = threading.RLock() self._cache: Optional[Dict[str, Any]] = None self._cache_time: float = 0 @@ -35,15 +52,15 @@ class StateStore: if self._cache is not None and (now - self._cache_time) < self._cache_ttl: return dict(self._cache) if not self.path.exists(): - result = self._default_state() + result = deepcopy(self.defaults) self._cache = result self._cache_time = now return result try: with self.path.open("r", encoding="utf-8") as handle: - state = json.load(handle) + state: Dict[str, Any] = json.load(handle) except (json.JSONDecodeError, OSError): - result = self._default_state() + result = deepcopy(self.defaults) self._cache = result self._cache_time = now return result @@ -56,10 +73,9 @@ class StateStore: with self.lock: self.path.parent.mkdir(parents=True, exist_ok=True) tmp_path = self.path.with_suffix(self.path.suffix + ".tmp") + merged = self._merge_defaults(state) with tmp_path.open("w", encoding="utf-8") as handle: - json.dump( - self._merge_defaults(state), handle, ensure_ascii=False, indent=2 - ) + json.dump(merged, handle, ensure_ascii=False, indent=2) handle.write("\n") tmp_path.replace(self.path) self._cache = None @@ -73,9 +89,7 @@ class StateStore: def continuous_config(self) -> Dict[str, Any]: state = self.load() - return dict( - state.get("continuous_scraping") or DEFAULT_STATE["continuous_scraping"] - ) + return deepcopy(state.get("continuous_scraping") or ACCOUNT_DEFAULTS["continuous_scraping"]) def save_continuous_config(self, config: Dict[str, Any]) -> Dict[str, Any]: def mutate(state: Dict[str, Any]) -> None: @@ -92,14 +106,208 @@ class StateStore: return self.update(mutate)["continuous_scraping"] - def _default_state(self) -> Dict[str, Any]: - return json.loads(json.dumps(DEFAULT_STATE)) - def _merge_defaults(self, state: Dict[str, Any]) -> Dict[str, Any]: - merged = self._default_state() + merged = deepcopy(self.defaults) for key, value in state.items(): - if key == "continuous_scraping" and isinstance(value, dict): - merged[key].update(value) + if isinstance(value, dict) and isinstance(merged.get(key), dict): + nested = deepcopy(merged[key]) + nested.update(value) + merged[key] = nested else: merged[key] = value return merged + + +# ── global state helpers ────────────────────────────────────────────── + +_GLOBAL_STORE: Optional[StateStore] = None + + +def get_global_store(data_dir: Path) -> StateStore: + global _GLOBAL_STORE + if _GLOBAL_STORE is None: + _GLOBAL_STORE = StateStore(data_dir / "state.json", defaults=GLOBAL_DEFAULTS) + return _GLOBAL_STORE + + +def load_global(data_dir: Path) -> Dict[str, Any]: + return get_global_store(data_dir).load() + + +def save_global(data_dir: Path, state: Dict[str, Any]) -> None: + get_global_store(data_dir).save(state) + + +def list_accounts(data_dir: Path) -> List[str]: + return list(load_global(data_dir).get("accounts", [])) + + +def account_exists(data_dir: Path, account_id: str) -> bool: + return account_id in list_accounts(data_dir) + + +# ── per-account state helpers ───────────────────────────────────────── + +_ACCOUNT_STORES: Dict[str, StateStore] = {} +_account_stores_lock = threading.Lock() + + +def get_account_store(data_dir: Path, account_id: str) -> StateStore: + global _ACCOUNT_STORES + with _account_stores_lock: + key = f"{data_dir}:{account_id}" + if key not in _ACCOUNT_STORES: + store = StateStore( + data_dir / "accounts" / account_id / "state.json", + defaults=ACCOUNT_DEFAULTS, + ) + _ACCOUNT_STORES[key] = store + return _ACCOUNT_STORES[key] + + +def load_account(data_dir: Path, account_id: str) -> Dict[str, Any]: + return get_account_store(data_dir, account_id).load() + + +def save_account(data_dir: Path, account_id: str, state: Dict[str, Any]) -> None: + get_account_store(data_dir, account_id).save(state) + + +def account_data_dir(data_dir: Path, account_id: str) -> Path: + """Per-account directory for channel DBs, media, exports.""" + return data_dir / "accounts" / account_id + + +def account_session_path(session_dir: Path, account_id: str) -> str: + """Per-account Telethon session file. Always under session/.session.""" + return str(session_dir / f"{account_id}.session") + + +# ── MIGRATION (with data copy) ───────────────────────────────────────── + + +def _copy_channel_data(src_root: Path, dst_root: Path, channel_id: str) -> None: + """Copy a single channel's DB + media from src_root to dst_root.""" + src_ch = src_root / channel_id + dst_ch = dst_root / channel_id + if not src_ch.exists(): + return + dst_ch.mkdir(parents=True, exist_ok=True) + + # SQLite DB + db_name = f"{channel_id}.db" + src_db = src_ch / db_name + dst_db = dst_ch / db_name + if src_db.exists() and not dst_db.exists(): + shutil.copy2(src_db, dst_db) + logger.info(" copied %s", db_name) + + # Media directory + src_media = src_ch / "media" + dst_media = dst_ch / "media" + if src_media.exists() and src_media.is_dir(): + dst_media.mkdir(parents=True, exist_ok=True) + for item in src_media.iterdir(): + if item.is_file(): + dst_file = dst_media / item.name + if not dst_file.exists(): + shutil.copy2(item, dst_file) + + # Exports (csv, json) + for ext in (".csv", ".json"): + for f in src_ch.glob(f"*{ext}"): + dst_f = dst_ch / f.name + if not dst_f.exists(): + shutil.copy2(f, dst_f) + + +def migrate_legacy_state(data_dir: Path, session_dir: Path) -> bool: + """ + Run once at startup. Detects legacy data/state.json that has api_id + and migrates to multi-account format, COPYING channel DBs/media from + data// → data/accounts/default// . + + Returns True if migration was performed. + """ + state_path = data_dir / "state.json" + if not state_path.exists(): + return False + + try: + with state_path.open("r", encoding="utf-8") as f: + raw: Dict[str, Any] = json.load(f) + except (json.JSONDecodeError, OSError): + return False + + # Already migrated or fresh install with no legacy api_id + if raw.get("version") == 2 or not raw.get("api_id"): + return False + + logger.info( + "=== Legacy state detected — migrating to multi-account format ===" + ) + logger.info("This may take a while if you have many channels with media.") + + # ── 1. Create per-account state for "default" ────────────────────── + acc_dir = data_dir / "accounts" / "default" + acc_dir.mkdir(parents=True, exist_ok=True) + + acc_state = { + "label": "Default", + "api_id": raw.get("api_id"), + "api_hash": raw.get("api_hash"), + "channels": raw.get("channels", {}), + "channel_names": raw.get("channel_names", {}), + "scrape_media": raw.get("scrape_media", True), + "forwarding_rules": raw.get("forwarding_rules", []), + "continuous_scraping": raw.get( + "continuous_scraping", + { + "enabled": True, + "interval_minutes": 1, + "channels": [], + "run_all_tracked": True, + }, + ), + } + + acc_state_path = acc_dir / "state.json" + with open(acc_state_path, "w", encoding="utf-8") as f: + json.dump(acc_state, f, ensure_ascii=False, indent=2) + f.write("\n") + logger.info(" ✓ data/accounts/default/state.json created") + + # ── 2. Copy channel data ─────────────────────────────────────────── + legacy_channels = list(raw.get("channels", {}).keys()) + if legacy_channels: + logger.info( + " Copying %d channel(s) from data// → data/accounts/default// ...", + len(legacy_channels), + ) + for i, channel_id in enumerate(legacy_channels, 1): + logger.info(" [%d/%d] channel %s", i, len(legacy_channels), channel_id) + _copy_channel_data(data_dir, acc_dir, channel_id) + logger.info(" ✓ Channel data copied") + else: + logger.info(" (no tracked channels to copy)") + + # ── 3. Rewrite global state.json ─────────────────────────────────── + global_state = { + "accounts": ["default"], + "version": 2, + } + with open(state_path, "w", encoding="utf-8") as f: + json.dump(global_state, f, ensure_ascii=False, indent=2) + f.write("\n") + logger.info(" ✓ data/state.json rewritten (accounts list only)") + + # ── 4. Rename legacy session file ───────────────────────────────── + legacy_session = session_dir / "session.session" + if legacy_session.exists(): + new_session = session_dir / "default.session" + if not new_session.exists(): + legacy_session.rename(new_session) + logger.info(" ✓ session/session.session → session/default.session") + + logger.info("=== Migration complete ===") + return True