Page the org backfill instead of streaming whole collections
A bare .stream() over calls/ returned 503 "Query timed out. Please try either limiting the entities scanned", which the Firestore client then re-raised as an AttributeError from its own retry path -- so the real cause was only visible in the chained traceback. The collection has simply outgrown a single scan. Both passes now walk each collection in 500-document pages ordered by document id, which needs no composite index. The counting pass also stops building a list of every document just to count the ones missing org_id. Documents written mid-run may be missed or seen twice; neither matters, since post-tenancy code stamps org_id at write time and the update is idempotent, so a second run cleans up anything the first skipped. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 5
parent
1a563c995c
commit
c3fe2a3466
@@ -76,6 +76,38 @@ TENANT_COLLECTIONS = ["nodes", "systems", "calls", "incidents", "alert_rules", "
|
||||
_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__,
|
||||
@@ -126,11 +158,15 @@ def main() -> None:
|
||||
counts: dict[str, tuple[int, int]] = {}
|
||||
total_missing = 0
|
||||
for collection in TENANT_COLLECTIONS:
|
||||
docs = list(db.collection(collection).stream())
|
||||
missing = [d for d in docs if not (d.to_dict() or {}).get("org_id")]
|
||||
counts[collection] = (len(docs), len(missing))
|
||||
total_missing += len(missing)
|
||||
print(f"{collection}: {len(docs)} docs total, {len(missing)} missing org_id")
|
||||
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}")
|
||||
|
||||
@@ -179,7 +215,7 @@ def main() -> None:
|
||||
batch = db.batch()
|
||||
batch_count = 0
|
||||
written = 0
|
||||
for doc in db.collection(collection).stream():
|
||||
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})
|
||||
|
||||
Reference in New Issue
Block a user