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
This commit is contained in:
+226
-18
@@ -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/<id>.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/<channel_id>/ → data/accounts/default/<channel_id>/ .
|
||||
|
||||
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/<id>/ → data/accounts/default/<id>/ ...",
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user