spb/company-atlas

1"""Operations commands: `catlas stats`, `catlas status`, `catlas backup`."""2from __future__ import annotations34import json5import os6import subprocess7import time8from datetime import UTC, datetime9from pathlib import Path1011import typer12from rich.table import Table1314from companyatlas.config import settings151617def register(app: typer.Typer) -> None:18    app.command()(stats)19    app.command()(status)20    app.command()(backup)212223def stats() -> None:24    """Dataset counters (companies, sensors, observations, changes, events, archive)."""25    from companyatlas.archive import store_stats26    from companyatlas.cli import out, run_async27    from companyatlas.db import connection, fetch_one2829    async def go():  # type: ignore[no-untyped-def]30        async with connection() as conn:31            return await fetch_one(conn, """32                select (select count(*) from companies) as companies,33                       (select count(*) from companies where onboarding_status = 'active') as companies_active,34                       (select count(*) from sensors) as sensors,35                       (select count(*) from sensors where status = 'active') as sensors_active,36                       (select count(*) from observations) as observations,37                       (select count(*) from snapshots) as snapshots,38                       (select count(*) from changes) as changes,39                       (select count(*) from changes where kind in ('meaningful','major','critical')) as meaningful_changes,40                       (select count(*) from events where status = 'active') as events,41                       (select count(*) from jobs where status = 'open') as jobs_open,42                       (select count(*) from observations where fetched_at > now() - interval '24 hours') as observations_24h,43                       (select count(*) from events where detected_at > now() - interval '24 hours' and status = 'active') as events_24h,44                       (select value from settings_kv where key = 'dataset_started_at') as dataset_started_at45            """)4647    row = run_async(go()) or {}48    row["archive"] = store_stats()49    started = row.get("dataset_started_at")50    if started:51        try:52            dt = datetime.fromisoformat(str(started).strip('"'))53            row["dataset_age_days"] = (datetime.now(UTC) - dt).days54        except ValueError:55            pass56    out.print(json.dumps(row, default=str, indent=1))575859def status() -> None:60    """Scheduler heartbeat, queue depth, failures by class (last 24 h), periodic task health."""61    from companyatlas.cli import out, run_async62    from companyatlas.db import connection, fetch_all, fetch_one6364    async def go():  # type: ignore[no-untyped-def]65        async with connection() as conn:66            hb = await fetch_one(conn, "select value, updated_at from settings_kv where key = 'scheduler:heartbeat'")67            due = await fetch_one(conn, """select count(*) filter (where next_run_at <= now()) as due, count(*) as total,68                                                 count(*) filter (where status = 'failing') as failing, count(*) filter (where status = 'stale') as stale,69                                                 count(*) filter (where status = 'blocked') as blocked, count(*) filter (where status = 'retired') as retired70                                          from sensors where status <> 'retired' or status = 'retired'""")71            queue = await fetch_all(conn, "select kind, status, count(*) as n from queue_jobs group by 1, 2 order by 1, 2")72            fails = await fetch_all(conn, "select failure_class, count(*) as n from failures where at > now() - interval '24 hours' group by 1 order by 2 desc")73            onboarding = await fetch_all(conn, "select onboarding_status, count(*) as n from companies group by 1 order by 2 desc")74            return hb, due, queue, fails, onboarding7576    hb, due, queue, fails, onboarding = run_async(go())77    if hb:78        age = (datetime.now(UTC) - hb["updated_at"]).total_seconds() if hb.get("updated_at") else None79        out.print(f"[bold]scheduler[/] heartbeat {int(age)}s ago" if age is not None else "[bold]scheduler[/] heartbeat unknown")80        val = hb["value"] if isinstance(hb["value"], dict) else json.loads(hb["value"])81        out.print(f"  worker={val.get('worker')} inflight={val.get('inflight')} due={val.get('due')} tick_ms={val.get('tick_ms')}")82        for t in val.get("tasks") or []:83            out.print(f"  task {t.get('name'):<22} runs={t.get('runs')} failures={t.get('failures')} err={t.get('last_error') or '-'}")84    else:85        out.print("[yellow]scheduler[/] no heartbeat yet")86    out.print(f"[bold]sensors[/] {due}")87    t = Table(title="onboarding")88    t.add_column("status")89    t.add_column("n", justify="right")90    for r in onboarding:91        t.add_row(str(r["onboarding_status"]), str(r["n"]))92    out.print(t)93    t = Table(title="queue")94    for c in ("kind", "status", "n"):95        t.add_column(c)96    for r in queue:97        t.add_row(r["kind"], r["status"], str(r["n"]))98    out.print(t)99    t = Table(title="failures 24h")100    t.add_column("class")101    t.add_column("n", justify="right")102    for r in fails:103        t.add_row(str(r["failure_class"]), str(r["n"]))104    out.print(t)105106107def backup(keep: int = typer.Option(14, help="number of dumps to keep")) -> None:108    """pg_dump (custom format, compressed) into CA_DATA_DIR/backups; prunes old dumps. Copy off-node with scripts/backup-offnode.sh."""109    from companyatlas.cli import out110111    settings.ensure_dirs()112    ts = time.strftime("%Y%m%d-%H%M%S")113    target = settings.backups_dir / f"companyatlas-{ts}.dump"114    env = dict(os.environ)115    cmd = ["pg_dump", "--format=custom", "--compress=zstd:6", "--no-owner", "--dbname", settings.sync_database_url, "--file", str(target)]116    try:117        subprocess.run(cmd, check=True, env=env, timeout=3600)118    except FileNotFoundError:119        cmd[0] = "/opt/homebrew/opt/postgresql@17/bin/pg_dump"120        subprocess.run(cmd, check=True, env=env, timeout=3600)121    dumps = sorted(Path(settings.backups_dir).glob("companyatlas-*.dump"))122    for old in dumps[:-keep]:123        old.unlink(missing_ok=True)124    out.print(f"[green]backup ok[/] {target} ({target.stat().st_size // 1024} kB), kept {min(len(dumps), keep)}")125126127try:  # register the nightly backup as a periodic task when the scheduler imports this module128    from companyatlas.services.periodic import periodic129130    @periodic("backup-nightly", cron=settings.backup_cron)131    async def _backup_task() -> None:132        import asyncio133134        await asyncio.to_thread(backup, 14)135except Exception:  # noqa: BLE001, S110 — optional: scheduler not present in this process136    pass137