fix oom during flask dataset check

Author: FuhuXia

scripts/sync_datasets.py

@@ -65,7 +65,7 @@ def _insert_dataset_for_record(interface: HarvesterDBInterface, record: HarvestR
             raise click.ClickException(str(exc))
 
 
-def _records_missing_datasets(session) -> List[HarvestRecord]:
+def _records_missing_datasets(session):
     return (
         session.query(HarvestRecord)
         .outerjoin(Dataset, Dataset.harvest_record_id == HarvestRecord.id)
@@ -74,11 +74,10 @@ def _records_missing_datasets(session) -> List[HarvestRecord]:
             HarvestRecord.action.in_(["create", "update"]),
             Dataset.id.is_(None),
         )
-        .all()
     )
 
 
-def _datasets_with_unexpected_records(session) -> List[Dataset]:
+def _datasets_with_unexpected_records(session):
     return (
         session.query(Dataset)
         .join(HarvestRecord, Dataset.harvest_record_id == HarvestRecord.id)
@@ -88,54 +87,85 @@ def _datasets_with_unexpected_records(session) -> List[Dataset]:
                 HarvestRecord.action.notin_(["create", "update"]),
             )
         )
-        .all()
     )
 
 
-def _report(records_missing, datasets_bad):
+def _report(records_missing_count: int, datasets_bad_count: int):
     click.echo("Dataset Sync Report\n====================")
-    click.echo(f"Records needing datasets: {len(records_missing)}")
+    click.echo(f"Records needing datasets: {records_missing_count}")
     click.echo(
-        f"Datasets tied to non-success/non-create records: {len(datasets_bad)}"
+        f"Datasets tied to non-success/non-create records: {datasets_bad_count}"
     )
 
 
 def _sync_impl(apply_changes: bool):
     interface = HarvesterDBInterface(session=db.session)
 
     try:
-        records_missing = _records_missing_datasets(db.session)
-        datasets_bad = _datasets_with_unexpected_records(db.session)
+        records_missing_query = _records_missing_datasets(db.session)
+        datasets_bad_query = _datasets_with_unexpected_records(db.session)
 
-        _report(records_missing, datasets_bad)
+        records_missing_count = records_missing_query.count()
+        datasets_bad_count = datasets_bad_query.count()
+
+        _report(records_missing_count, datasets_bad_count)
 
         if apply_changes:
             synced = 0
-            total = len(records_missing)
-            for start in range(0, total, BATCH_SIZE):
-                batch = records_missing[start : start + BATCH_SIZE]
+            batches = (records_missing_count + BATCH_SIZE - 1) // BATCH_SIZE
+            current_batch = 0
+            batch_records: List[HarvestRecord] = []
+            for record in records_missing_query.yield_per(BATCH_SIZE):
+                batch_records.append(record)
+                if len(batch_records) == BATCH_SIZE:
+                    current_batch += 1
+                    click.echo(
+                        f"Processing batch {current_batch} "
+                        f"({len(batch_records)} records)..."
+                    )
+                    for record_in_batch in batch_records:
+                        try:
+                            slug = _insert_dataset_for_record(
+                                interface, record_in_batch
+                            )
+                            synced += 1
+                            click.echo(
+                                f"Created dataset for record {record_in_batch.id} "
+                                f"(slug: {slug})"
+                            )
+                        except click.ClickException as exc:
+                            click.echo(
+                                f"Failed to sync record {record_in_batch.id}: {exc}"
+                            )
+                    batch_records = []
+
+            if batch_records:
+                current_batch += 1
                 click.echo(
-                    f"Processing batch {start // BATCH_SIZE + 1}"
-                    f" ({len(batch)} records)..."
+                    f"Processing batch {current_batch} "
+                    f"({len(batch_records)} records)..."
                 )
-                for record in batch:
+                for record_in_batch in batch_records:
                     try:
-                        slug = _insert_dataset_for_record(interface, record)
+                        slug = _insert_dataset_for_record(interface, record_in_batch)
                         synced += 1
                         click.echo(
-                            f"Created dataset for record {record.id} (slug: {slug})"
+                            f"Created dataset for record {record_in_batch.id} "
+                            f"(slug: {slug})"
                         )
                     except click.ClickException as exc:
-                        click.echo(f"Failed to sync record {record.id}: {exc}")
+                        click.echo(
+                            f"Failed to sync record {record_in_batch.id}: {exc}"
+                        )
             click.echo(f"Datasets created: {synced}")
 
             deleted = 0
-            if datasets_bad:
+            if datasets_bad_count:
                 click.echo(
-                    f"Deleting {len(datasets_bad)} dataset(s) tied "
+                    f"Deleting {datasets_bad_count} dataset(s) tied "
                     "to invalid harvest records..."
                 )
-                for dataset in datasets_bad:
+                for dataset in datasets_bad_query.yield_per(BATCH_SIZE):
                     try:
                         interface.db.delete(dataset)
                         interface.db.commit()