#!/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()