server-26/drb-c2-core/scripts/backfill_org_id.py

#!/usr/bin/env python3
"""
Backfill org_id onto every pre-tenancy document and create the founding org.

*** WRITE-ONLY REFERENCE — NOT RUN AS PART OF THIS CHANGE. ***
SAAS_PLAN.md B2 explicitly says "write it; do not run it" — this script
touches production Firestore (organizations, org_members, nodes, systems,
calls, incidents, alert_rules, alert_events) and Firebase Auth custom
claims. Read this whole docstring before ever running it anywhere.

WHY IT'S NEEDED: as of this pass, every doc in the six tenant collections
below predates the org_id field entirely (org_id is Optional[...] = None on
every model in app/models.py specifically to allow this). c2-core's read
routes and infra/firestore/firestore.rules now filter/require org_id, so
until this runs, pre-existing docs are invisible through the org-scoped
paths — they still exist, they're just unreachable by a caller whose token
carries an org_id claim. New docs created going forward (enrollment.py,
mqtt_handler.py, upload.py, incident_correlator.py) already stamp org_id
themselves; this script only needs to run ONCE, retroactively, and is safe
to re-run after that (idempotent — see below).

WHAT IT DOES:
  1. Creates organizations/{FOUNDING_ORG_ID} if it doesn't already exist.
     FOUNDING_ORG_ID ("founding") is the same id app/internal/tenancy.py
     defines and the same id enrollment.py's legacy fleet-wide
     ENROLLMENT_TOKEN fallback and mqtt_handler.py's legacy MQTT-checkin
     path both already resolve brand-new nodes into — so a node that
     enrolled the old way and a doc backfilled by this script end up in the
     same org.
  2. Sets --owner-email's org_id/org_role Firebase custom claims and writes
     their org_members doc — the same shape POST /auth/signup writes for a
     self-serve org, so this person becomes the founding org's owner in the
     UI exactly as if they'd signed up normally. Their platform `role`
     claim is left alone if already set, else defaults to "admin" (the
     backfill owner is presumed to be today's single-tenant deployment's
     admin).
  3. Walks nodes / systems / calls / incidents / alert_rules / alert_events
     and stamps org_id = FOUNDING_ORG_ID onto every document that doesn't
     already have one. A document that already has org_id (from the
     post-tenancy code paths that shipped alongside this script) is left
     untouched — this is what makes a second run a no-op rather than a
     re-stamp, so running it twice by accident is harmless.

USAGE (run from the drb-c2-core directory, with GCP_CREDENTIALS_PATH set or
Application Default Credentials available — same auth as set_admin.py):

  python scripts/backfill_org_id.py --owner-email you@example.com --dry-run
  python scripts/backfill_org_id.py --owner-email you@example.com

ALWAYS run with --dry-run first and read every line of its output — it
prints exactly what would be created/changed, with no writes, before you
run it for real. --dry-run performs full collection scans (read-only) to
produce accurate counts; on a large calls/incidents collection this is not
free, but it is the only way to know the real backfill count in advance.

NOT HANDLED: org_api_keys (collection doesn't exist server-side yet — see
DEFERRED.md) and node_keys (deliberately never gets an org_id column; it's
looked up by node_id / api_key value, not read as an org-scoped list).
"""
import argparse
import os
import sys
from datetime import datetime, timezone

import firebase_admin
from firebase_admin import auth, credentials, firestore

# Mirrors app/internal/tenancy.py — duplicated rather than imported so this
# script has no dependency on the app package (or its settings/env) being
# importable from wherever it's actually run.
FOUNDING_ORG_ID = "founding"

TENANT_COLLECTIONS = ["nodes", "systems", "calls", "incidents", "alert_rules", "alert_events"]

# Firestore batched writes cap at 500 operations; stay comfortably under it.
_BATCH_SIZE = 400


_PAGE_SIZE = 500


def _iter_collection(db, collection: str, page_size: int = _PAGE_SIZE):
    """
    Walk a collection in key-ordered pages instead of one open stream.

    A bare .stream() over calls/ died with "503 Query timed out. Please try
    either limiting the entities scanned" -- the collection has outgrown what
    Firestore will serve as a single scan, and the failure surfaced as an
    unrelated-looking AttributeError from the client's own retry path. Paging by
    document id needs no composite index and keeps each request small; it also
    stops the counting pass holding every document in memory at once.

    Documents written while this runs may be missed or seen twice. Neither
    matters: post-tenancy code stamps org_id itself, and the update below is
    idempotent.
    """
    base = db.collection(collection).order_by("__name__").limit(page_size)
    cursor = None
    while True:
        query = base.start_after(cursor) if cursor is not None else base
        page = list(query.stream())
        if not page:
            return
        for doc in page:
            yield doc
        if len(page) < page_size:
            return
        cursor = page[-1]


def main() -> None:
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--owner-email", required=True, help="Firebase user who becomes the founding org's owner")
    parser.add_argument("--org-name", default="Founding Org", help="Display name for the founding org")
    parser.add_argument("--dry-run", action="store_true", help="Print what would change; write nothing")
    args = parser.parse_args()

    # Match app/internal/firestore.py exactly. Two ways this script diverged
    # from the app and would have failed or, worse, half-succeeded:
    #   * credentials — GCP_CREDENTIALS_PATH is unset in production (the server
    #     is a GCE instance and the app uses Application Default Credentials via
    #     the metadata server). Defaulting to a "gcp-key.json" that does not
    #     exist made the script unrunnable there.
    #   * database — the app talks to FIRESTORE_DATABASE (c2-server in prod),
    #     while a bare firestore.client() talks to "(default)". That one is the
    #     dangerous half: the script would have scanned an empty database, found
    #     nothing to backfill, created the founding org in the wrong place and
    #     printed a clean success.
    creds_path = os.getenv("GCP_CREDENTIALS_PATH")
    cred = credentials.Certificate(creds_path) if creds_path else credentials.ApplicationDefault()
    firebase_admin.initialize_app(cred)

    database_id = os.getenv("FIRESTORE_DATABASE", "(default)")
    print(f"Using Firestore database: {database_id}")
    db = firestore.client(database_id=database_id)

    try:
        owner = auth.get_user_by_email(args.owner_email)
    except auth.UserNotFoundError:
        print(f"No Firebase user found for {args.owner_email!r}")
        sys.exit(1)

    now = datetime.now(timezone.utc).isoformat()
    existing_claims = owner.custom_claims or {}

    org_ref = db.collection("organizations").document(FOUNDING_ORG_ID)
    org_exists = org_ref.get().exists
    print(f"organizations/{FOUNDING_ORG_ID}: {'exists — left alone' if org_exists else 'WILL CREATE'}")
    print(f"org_members/{owner.uid}: WILL SET org_role='owner' (email={args.owner_email})")
    print(
        f"Firebase custom claims for {args.owner_email}: WILL SET org_id={FOUNDING_ORG_ID!r} org_role='owner', "
        f"role={existing_claims.get('role', 'admin (default)')!r}"
    )

    counts: dict[str, tuple[int, int]] = {}
    total_missing = 0
    for collection in TENANT_COLLECTIONS:
        total = 0
        missing = 0
        for doc in _iter_collection(db, collection):
            total += 1
            if not (doc.to_dict() or {}).get("org_id"):
                missing += 1
        counts[collection] = (total, missing)
        total_missing += missing
        print(f"{collection}: {total} docs total, {missing} missing org_id")

    print(f"\nTotal documents to backfill: {total_missing}")

    if args.dry_run:
        print("\n--dry-run: no writes performed.")
        return

    if not org_exists:
        org_ref.set({
            "org_id": FOUNDING_ORG_ID,
            "name": args.org_name,
            "created_at": now,
            "created_by_uid": owner.uid,
            # Inert placeholders — no billing model exists yet, see
            # app/internal/tenancy.py and models.py's OrganizationRecord.
            "plan_id": None,
            "subscription_status": None,
            "stripe_customer_id": None,
            "stripe_subscription_id": None,
            "current_period_end": None,
            "seat_limit": None,
            "node_limit": None,
            "retention_days": None,
        })
        print(f"Created organizations/{FOUNDING_ORG_ID}.")

    db.collection("org_members").document(owner.uid).set({
        "uid": owner.uid,
        "org_id": FOUNDING_ORG_ID,
        "org_role": "owner",
        "email": args.owner_email,
        "added_at": now,
    }, merge=True)
    print(f"Set org_members/{owner.uid}.")

    new_claims = {**existing_claims, "org_id": FOUNDING_ORG_ID, "org_role": "owner"}
    new_claims.setdefault("role", "admin")
    auth.set_custom_user_claims(owner.uid, new_claims)
    print(f"Set custom claims for {args.owner_email}.")

    for collection in TENANT_COLLECTIONS:
        _, missing_count = counts[collection]
        if not missing_count:
            print(f"{collection}: nothing to backfill.")
            continue
        batch = db.batch()
        batch_count = 0
        written = 0
        for doc in _iter_collection(db, collection):
            if (doc.to_dict() or {}).get("org_id"):
                continue
            batch.update(doc.reference, {"org_id": FOUNDING_ORG_ID})
            batch_count += 1
            written += 1
            if batch_count >= _BATCH_SIZE:
                batch.commit()
                batch = db.batch()
                batch_count = 0
        if batch_count:
            batch.commit()
        print(f"{collection}: backfilled {written} document(s).")

    print("\nDone. The owner must sign out and back in (or wait up to 1 hour) for the new claims to take effect.")


if __name__ == "__main__":
    main()