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
|
_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:
|
def main() -> None:
|
||||||
parser = argparse.ArgumentParser(
|
parser = argparse.ArgumentParser(
|
||||||
description=__doc__,
|
description=__doc__,
|
||||||
@@ -126,11 +158,15 @@ def main() -> None:
|
|||||||
counts: dict[str, tuple[int, int]] = {}
|
counts: dict[str, tuple[int, int]] = {}
|
||||||
total_missing = 0
|
total_missing = 0
|
||||||
for collection in TENANT_COLLECTIONS:
|
for collection in TENANT_COLLECTIONS:
|
||||||
docs = list(db.collection(collection).stream())
|
total = 0
|
||||||
missing = [d for d in docs if not (d.to_dict() or {}).get("org_id")]
|
missing = 0
|
||||||
counts[collection] = (len(docs), len(missing))
|
for doc in _iter_collection(db, collection):
|
||||||
total_missing += len(missing)
|
total += 1
|
||||||
print(f"{collection}: {len(docs)} docs total, {len(missing)} missing org_id")
|
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}")
|
print(f"\nTotal documents to backfill: {total_missing}")
|
||||||
|
|
||||||
@@ -179,7 +215,7 @@ def main() -> None:
|
|||||||
batch = db.batch()
|
batch = db.batch()
|
||||||
batch_count = 0
|
batch_count = 0
|
||||||
written = 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"):
|
if (doc.to_dict() or {}).get("org_id"):
|
||||||
continue
|
continue
|
||||||
batch.update(doc.reference, {"org_id": FOUNDING_ORG_ID})
|
batch.update(doc.reference, {"org_id": FOUNDING_ORG_ID})
|
||||||
|
|||||||
Reference in New Issue
Block a user