"""HuggingFace Dataset storage for the hosting panel. Mirrors the pattern from bot/storage.py but adapted for the panel's needs. The panel stores its SQLite DB + any uploaded config files in a private HF Dataset. Layout in the dataset repo: latest/panel.sqlite latest/backups/... snapshots//panel.sqlite Persistence model: **write-through with debounce**. Every committed DB change schedules a push; the coordinator collapses bursts of writes into a single HF upload after ``storage_debounce_sec`` of quiet time. A low-frequency safety-net loop also flushes if any write got lost between the listener and the loop. """ from __future__ import annotations import asyncio import logging import os import shutil import tempfile from datetime import datetime, timezone from pathlib import Path from typing import Optional from config import PANEL_CONFIG from hf_auth import get_read_token, get_write_token, hf_read_api, hf_write_api log = logging.getLogger("panel-storage") class PanelStorage: """Panel-level HF Dataset backup. Two-token aware: ``pull`` uses ``HF_TOKEN_READ`` (full read scope), ``push`` (and any other write) uses ``HF_TOKEN_WRITE`` (scoped to a pre-approved allowlist of Spaces). If a single legacy ``HF_TOKEN`` is configured it transparently serves both roles. """ def __init__( self, repo: str | None = None, token: str | None = None, # legacy single-token override (deprecated) data_dir: Path | None = None, ) -> None: self.repo = repo or PANEL_CONFIG.hf_dataset_repo self.data_dir = data_dir or PANEL_CONFIG.data_dir # If a legacy ``token`` arg was passed, override both roles # (preserves backward compatibility for in-test injection). if token: import hf_auth as _hf _hf._read_api = None _hf._write_api = None self._legacy_token = token # remembered only for ``_get_api`` self._api = None @property def enabled(self) -> bool: # Either read OR write token is enough to "have a dataset" \u2014 # readonly callers (pull from a snapshot) don't need write scope. return bool( self.repo and (get_read_token() or get_write_token()) ) def _get_api(self): """Return a fresh HfApi for *push* (write token). Most callers in this module are pushing, so the default is the write API. ``pull`` overrides with ``_get_read_api``. """ if self._legacy_token: from huggingface_hub import HfApi return HfApi(token=self._legacy_token) if self._api is None: self._api = hf_write_api() or hf_read_api() return self._api def _get_read_api(self): """Return a fresh HfApi for *read* operations (read token).""" if self._legacy_token: from huggingface_hub import HfApi return HfApi(token=self._legacy_token) return hf_read_api() or hf_write_api() async def pull(self) -> bool: if not self.enabled: log.info("panel-storage: HF Dataset not configured — skipping pull.") return False loop = asyncio.get_event_loop() try: return await loop.run_in_executor(None, self._sync_pull) except Exception as exc: log.warning("panel-storage pull failed: %s", exc) return False def _sync_pull(self) -> bool: """Atomic snapshot restore. Uses the read-only HF token so this works even if the operator hasn't added this Space to the write token's allowlist (a stale ``latest/`` snapshot is still readable). Downloads into a temp directory, then either renames the SQLite file into ``data_dir/panel.sqlite`` (overwriting only after a successful copy) or returns ``False`` if nothing usable came back. If the local DB is missing AND pull failed, callers MUST refuse to start -- pushing an empty schema would otherwise overwrite live data on the dataset. """ # snapshot_download is a read; route through the read token. import os as _os _read_tok = get_read_token() or self._legacy_token or "" _saved_token = _os.environ.get("HF_TOKEN") if _read_tok: _os.environ["HF_TOKEN"] = _read_tok try: from huggingface_hub import snapshot_download with tempfile.TemporaryDirectory() as tmp: path = snapshot_download( repo_id=self.repo, repo_type="dataset", allow_patterns=["latest/*"], local_dir=tmp, token=_read_tok or None, ) src = Path(path) / "latest" if not src.exists(): log.info("panel-storage: nothing to pull (first run?).") return False db_src = src / "panel.sqlite" if not db_src.exists(): log.info("panel-storage: latest/ contains no panel.sqlite") return False self.data_dir.mkdir(parents=True, exist_ok=True) self.data_dir.chmod(0o700) db_dst = self.data_dir / "panel.sqlite" # Remove any stale sidecar files from a previous/throwaway # db generation before swapping the main file in -- SQLite # gets confused (and silently falls back to readonly) if # it finds a -journal/-wal/-shm that doesn't belong to the # file it's about to open. for suffix in ("-journal", "-wal", "-shm"): stale = self.data_dir / f"panel.sqlite{suffix}" stale.unlink(missing_ok=True) # Stage beside the live file so we only swap atomically on success. staged = self.data_dir / f"panel.sqlite.incoming-{os.getpid()}" try: shutil.copy2(db_src, staged) # shutil.copy2 preserves the source file's permission # bits. huggingface_hub's snapshot_download sources # files from its shared local cache, where blobs are # deliberately read-only (mode 444) to prevent # accidental corruption of the shared cache. Without # this chmod, the restored panel.sqlite inherits that # read-only bit and every subsequent write (including # the very next audit-log insert) fails with # "attempt to write a readonly database". staged.chmod(0o600) os.replace(staged, db_dst) except Exception: try: staged.unlink() except FileNotFoundError: pass raise log.info( "panel-storage: restored panel.sqlite (%d B)", db_dst.stat().st_size, ) # ------------------------------------------------------------------ # Banner images (admin-uploaded marketing assets). # # These live in ``data_dir/banners/`` and are mirrored to the # dataset via the snapshot push. On restore we drop the local # copy and replace it wholesale so renamed / deleted files # propagate from the dataset mirror to the running Space. # ------------------------------------------------------------------ banners_src = src / "banners" banners_dst = self.data_dir / "banners" if banners_src.exists(): if banners_dst.exists(): shutil.rmtree(banners_dst, ignore_errors=True) shutil.copytree(banners_src, banners_dst) try: banners_dst.chmod(0o755) except Exception: pass log.info( "panel-storage: restored banners dir (%d file(s))", sum(1 for _ in banners_dst.iterdir()), ) return True finally: if _saved_token is None: _os.environ.pop("HF_TOKEN", None) else: _os.environ["HF_TOKEN"] = _saved_token @property def safe_to_push(self) -> bool: """Whether writes may be mirrored to the dataset. False on cold start when pull failed and no local DB exists -- pushing an empty schema would clobber existing data on the dataset. Routes / listeners should check this before scheduling a push. """ if not self.enabled: return True # disabled storage -> no risk return self.data_dir.joinpath("panel.sqlite").exists() async def push(self, message: str = "") -> bool: if not self.enabled: return False loop = asyncio.get_event_loop() try: return await loop.run_in_executor(None, self._sync_push, message) except Exception as exc: log.warning("panel-storage push failed: %s", exc) return False def _sync_push(self, message: str) -> bool: api = self._get_api() try: api.create_repo( repo_id=self.repo, repo_type="dataset", private=True, exist_ok=True, ) except Exception as exc: log.warning("panel-storage: create_repo failed: %s", exc) ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") db_file = self.data_dir / "panel.sqlite" banners_dir = self.data_dir / "banners" with tempfile.TemporaryDirectory() as tmp: tmp_path = Path(tmp) if db_file.exists(): shutil.copy2(db_file, tmp_path / "panel.sqlite") # Mirror banner images so admin uploads survive Space rebuilds. if banners_dir.exists(): tmp_banners = tmp_path / "banners" shutil.copytree(banners_dir, tmp_banners) # Upload snapshot api.upload_folder( folder_path=str(tmp_path), repo_id=self.repo, repo_type="dataset", path_in_repo=f"snapshots/{ts}", commit_message=message or f"snapshot {ts}", token=self.token, ) # Mirror to latest api.upload_folder( folder_path=str(tmp_path), repo_id=self.repo, repo_type="dataset", path_in_repo="latest", commit_message=f"latest -> {ts}", token=self.token, ) log.info("panel-storage: pushed snapshot %s", ts) # Retention try: self._enforce_retention(api) except Exception as exc: log.warning("panel-storage: retention failed: %s", exc) return True def _enforce_retention(self, api) -> None: keep = PANEL_CONFIG.snapshot_retention_count if keep <= 0: return try: files = api.list_repo_files( repo_id=self.repo, repo_type="dataset", token=self.token ) except Exception: return prefix = "snapshots/" snapshot_names = set() for f in files: if f.startswith(prefix) and "/" not in f[len(prefix):]: snapshot_names.add(f[len(prefix):]) ordered = sorted(snapshot_names, reverse=True) for name in ordered[keep:]: try: api.delete_folder( repo_id=self.repo, repo_type="dataset", path_in_repo=f"snapshots/{name}", token=self.token, commit_message=f"retention: drop {name}", ) log.info("panel-storage: pruned old snapshot %s", name) except Exception as exc: log.warning("panel-storage: failed to prune %s: %s", name, exc) async def panel_storage_loop(storage: PanelStorage, interval_sec: int) -> None: if not storage.enabled: log.info("panel-storage: background loop disabled.") return log.info("panel-storage: loop running every %ds", interval_sec) while True: try: await asyncio.sleep(interval_sec) await storage.push(message="periodic snapshot") except asyncio.CancelledError: return except Exception as exc: log.warning("panel-storage loop error: %s", exc) await asyncio.sleep(30) # --------------------------------------------------------------------------- # Write-through coordinator # --------------------------------------------------------------------------- class PanelStorageCoordinator: """Debounced write-through layer on top of :class:`PanelStorage`. Routes / models schedule a push after every DB commit (via SQLAlchemy's ``after_commit`` event). The coordinator waits for ``debounce_sec`` of quiet time, then performs a single HTTP upload. Subsequent commits during a push simply mark the buffer dirty again so the next loop iteration will pick them up. Public surface: - :meth:`schedule_push` -- non-blocking; call from event listeners. - :meth:`flush_now` -- blocking; await in shutdown / tests. - :meth:`run_safety_net` -- long-running task; periodic flush, never replaces write-through. """ def __init__( self, storage: PanelStorage, debounce_sec: float | None = None, ) -> None: self.storage = storage self.debounce_sec = ( debounce_sec if debounce_sec is not None else PANEL_CONFIG.storage_debounce_sec ) self._lock = asyncio.Lock() self._dirty: bool = False self._last_reason: str = "change" self._write_count: int = 0 self._push_count: int = 0 self._timer_task: Optional[asyncio.Task] = None self._stopping: bool = False # -- call from after_commit (sync context) ------------------------------ def schedule_push_sync(self, reason: str = "commit") -> None: """Schedule a debounced push from sync code (e.g. a SQLAlchemy ``after_commit`` listener). Fire-and-forget; safe to call from anywhere an event loop is reachable. """ if self._stopping or not self.storage.enabled: return try: loop = asyncio.get_running_loop() except RuntimeError: # No running loop (e.g. unit test outside pytest-asyncio). # Best-effort: spawn a thread to push synchronously. try: self.storage.push(message=reason) except Exception as exc: log.warning("panel-storage sync push failed: %s", exc) return loop.create_task(self.schedule_push(reason)) async def schedule_push(self, reason: str = "commit") -> None: """Mark state dirty and ensure the debounced worker is scheduled.""" if self._stopping or not self.storage.enabled: return if not self.storage.safe_to_push: # Cold start without a successful pull -- refuse to push. log.warning( "panel-storage: refusing to push (%s) -- local DB absent " "and storage pull did not restore one. Data could be lost.", reason, ) return async with self._lock: self._dirty = True self._last_reason = reason self._write_count += 1 if self._timer_task is None or self._timer_task.done(): self._timer_task = asyncio.create_task( self._debounced_loop(), name="storage-debounce", ) # -- internals ---------------------------------------------------------- async def _debounced_loop(self) -> None: """Wait for quiet, push, repeat while dirty. Exits when clean.""" try: while True: await asyncio.sleep(self.debounce_sec) async with self._lock: if not self._dirty: return self._dirty = False reason = self._last_reason pushed = await self._do_push(reason) if not pushed: # Restore dirty so the next iteration retries. async with self._lock: if not self._dirty: self._dirty = True self._last_reason = f"retry: {reason}" except asyncio.CancelledError: raise async def _do_push(self, reason: str) -> bool: try: ok = await self.storage.push(message=reason) except Exception as exc: log.warning("panel-storage debounced push failed: %s", exc) return False if ok: async with self._lock: self._push_count += 1 return ok # -- long-running safety-net ------------------------------------------- async def run_safety_net(self, interval_sec: int) -> None: """Periodic flush as a safety net in case the debounced worker ever falls behind (e.g. an exception path that didn't mark dirty, or a write race). Lower frequency than the listener -- this is a belt, not the suspenders.""" if not self.storage.enabled: return log.info( "panel-storage: safety-net running every %ds (debounce=%ss)", interval_sec, self.debounce_sec, ) while not self._stopping: try: await asyncio.sleep(interval_sec) if self._stopping: return async with self._lock: if not self._dirty: continue self._dirty = False reason = f"safety-net: {self._last_reason}" await self._do_push(reason) except asyncio.CancelledError: return except Exception as exc: log.warning("panel-storage safety-net error: %s", exc) await asyncio.sleep(30) # -- shutdown / tests -------------------------------------------------- async def flush_now(self, reason: str = "flush") -> bool: """Force a sync push right now (await in shutdown / tests). Any already-scheduled timer is cancelled first.""" if not self.storage.enabled: return False async with self._lock: if self._timer_task and not self._timer_task.done(): self._timer_task.cancel() try: await self._timer_task except (asyncio.CancelledError, Exception): pass dirty = self._dirty self._dirty = False if not dirty: return False return await self._do_push(reason) async def shutdown(self) -> None: """Cancel timer, push any remaining dirty state, mark stopped.""" self._stopping = True if self._timer_task and not self._timer_task.done(): self._timer_task.cancel() try: await self._timer_task except (asyncio.CancelledError, Exception): pass if self.storage.enabled: await self.flush_now(reason="shutdown") # -- diagnostics -------------------------------------------------------- @property def stats(self) -> dict: return { "writes": self._write_count, "pushes": self._push_count, "dirty": self._dirty, "enabled": self.storage.enabled, "debounce_sec": self.debounce_sec, } # Module-level coordinator singleton. Initialised lazily on first access so # importing this module never opens HF API connections during tests. PANEL_STORAGE_COORD: PanelStorageCoordinator = PanelStorageCoordinator( PanelStorage() ) # --------------------------------------------------------------------------- # Audit-log pruning # --------------------------------------------------------------------------- async def audit_log_prune_loop(retention_days: int, interval_sec: int = 6 * 3600) -> None: """Background loop that trims ``audit_logs`` older than ``retention_days``. Runs once every ``interval_sec`` (default 6h). A separate task so an audit-log table bloat doesn't push out the snapshot cadence. Uses a fresh session because the per-request session is closed by the time the task fires. """ from datetime import timedelta from sqlalchemy import delete from database import async_session from models import AuditLog if retention_days <= 0: log.info("audit-log pruning: disabled (retention_days=%d)", retention_days) return log.info("audit-log pruning: retention=%d days, every %ds", retention_days, interval_sec) while True: try: await asyncio.sleep(interval_sec) cutoff = datetime.now(timezone.utc) - timedelta(days=retention_days) async with async_session() as session: result = await session.execute( delete(AuditLog).where(AuditLog.created_at < cutoff) ) await session.commit() log.info( "audit-log pruning: deleted %d row(s) older than %s", result.rowcount, cutoff.isoformat(), ) except asyncio.CancelledError: return except Exception as exc: log.warning("audit-log pruning error: %s", exc) await asyncio.sleep(60)