0pt/services/analytics/ingest_aud0.py

"""
Backfill/refresh: AUD0 events -> 0pt ClickHouse analytics events.

One audit event becomes one analytics event with metric `event_count` and
dimensions `event_type`, `app_id`, `capability_id`, `risk_level`. The write is
idempotent per event_id: already-present ids are skipped before the batch is
inserted. No synthetic event is ever created; an empty AUD0 page writes
nothing.

Usage:
  AUD0_TOKEN=... AUD0_TENANT_ID=379... CLICKHOUSE_URL=http://... \
    python ingest_aud0.py --limit 50000
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from typing import Any, Iterable

import httpx

AUD0_API_URL = os.getenv("AUD0_API_URL", "https://api.aud0.siax.io").rstrip("/")
AUD0_TOKEN = os.getenv("AUD0_TOKEN", "")
AUD0_TENANT_ID = os.getenv("AUD0_TENANT_ID", "")
CLICKHOUSE_URL = os.getenv("CLICKHOUSE_URL", "").rstrip("/")
CLICKHOUSE_USER = os.getenv("CLICKHOUSE_USER", "default")
CLICKHOUSE_PASSWORD = os.getenv("CLICKHOUSE_PASSWORD", "")
CLICKHOUSE_DATABASE = os.getenv("CLICKHOUSE_DATABASE", "opt_analytics")


def dims(event: dict[str, Any]) -> tuple[list[str], list[str]]:
    names: list[str] = []
    values: list[str] = []
    for key in ("event_type", "app_id", "capability_id", "risk_level"):
        value = event.get(key)
        if isinstance(value, str) and value:
            names.append(key)
            values.append(value)
    return names, values


def to_row(tenant: str, event: dict[str, Any]) -> dict[str, Any] | None:
    event_id = event.get("event_id")
    created_at = event.get("created_at")
    event_type = event.get("event_type")
    if not isinstance(event_id, str) or not event_id:
        return None
    if not isinstance(created_at, str) or not created_at:
        return None
    if not isinstance(event_type, str) or not event_type:
        return None
    names, values = dims(event)
    return {
        "tenant_id": tenant,
        "project_id": "aud0",
        "dataset_id": "aud0.events",
        "metric": "event_count",
        "event_id": event_id,
        "event_time": created_at,
        "value": 1,
        "dim_names": names,
        "dim_values": values,
    }


def ch_insert(rows: Iterable[dict[str, Any]]) -> None:
    payload = "\n".join(json.dumps(row, ensure_ascii=False) for row in rows)
    if not payload:
        return
    auth = (CLICKHOUSE_USER, CLICKHOUSE_PASSWORD) if CLICKHOUSE_PASSWORD else None
    with httpx.Client(base_url=CLICKHOUSE_URL, auth=auth, timeout=60.0) as client:
        response = client.post(
            "/",
            params={"database": CLICKHOUSE_DATABASE, "query": "INSERT INTO opt_analytics_event FORMAT JSONEachRow"},
            content=payload,
        )
        response.raise_for_status()


def existing_ids(tenant: str, ids: list[str]) -> set[str]:
    if not ids:
        return set()
    auth = (CLICKHOUSE_USER, CLICKHOUSE_PASSWORD) if CLICKHOUSE_PASSWORD else None
    found: set[str] = set()
    sql = (
        f"SELECT event_id FROM {CLICKHOUSE_DATABASE}.opt_analytics_event "
        "WHERE tenant_id = {tenant:String} AND event_id IN ({ids:Array(String)}) FORMAT JSON"
    )
    with httpx.Client(base_url=CLICKHOUSE_URL, auth=auth, timeout=30.0) as client:
        # ClickHouse query parameters travel in the URL; 1000 ids overflow it.
        for start in range(0, len(ids), 50):
            chunk = ids[start : start + 50]
            response = client.post(
                "/",
                params={
                    "database": CLICKHOUSE_DATABASE,
                    "param_tenant": tenant,
                    "param_ids": "[" + ",".join("'" + value.replace("'", "\\'") + "'" for value in chunk) + "]",
                },
                content=sql,
            )
            response.raise_for_status()
            found.update(row["event_id"] for row in response.json().get("data", []))
    return found


def fetch_page(offset: int, limit: int) -> dict[str, Any]:
    if not AUD0_TOKEN or not AUD0_TENANT_ID:
        raise SystemExit("AUD0_TOKEN and AUD0_TENANT_ID are required")
    with httpx.Client(base_url=AUD0_API_URL, timeout=30.0) as client:
        response = client.get(
            "/v1/aud0/events",
            params={"tenant_id": AUD0_TENANT_ID, "limit": limit, "offset": offset},
            headers={"authorization": f"Bearer {AUD0_TOKEN}", "x-tenant-id": AUD0_TENANT_ID},
        )
        response.raise_for_status()
        return response.json()


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit", type=int, default=50_000, help="maximum events to read")
    parser.add_argument("--page", type=int, default=1000)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    if not CLICKHOUSE_URL:
        raise SystemExit("CLICKHOUSE_URL is required")

    written = skipped = seen = 0
    for offset in range(0, args.limit, args.page):
        page = fetch_page(offset, min(args.page, args.limit - offset))
        events = page.get("events") or []
        if not events:
            break
        rows = [row for row in (to_row(AUD0_TENANT_ID, event) for event in events) if row]
        present = existing_ids(AUD0_TENANT_ID, [row["event_id"] for row in rows])
        fresh = [row for row in rows if row["event_id"] not in present]
        skipped += len(rows) - len(fresh)
        seen += len(events)
        if not args.dry_run:
            ch_insert(fresh)
        written += len(fresh)
        print(json.dumps({"offset": offset, "read": len(events), "written": len(fresh), "skipped": skipped}))
        if len(events) < args.page:
            break
    print(json.dumps({"seen": seen, "written": written, "skipped": skipped, "dry_run": args.dry_run}))
    return 0


if __name__ == "__main__":
    sys.exit(main())