Limit to MAX_BLOBS_PER_RUN (#58)

Author: di

main.py

@@ -31,6 +31,7 @@
 DATASETS = os.environ.get("BIGQUERY_DATASET", "").strip().split()
 SIMPLE_TABLE = os.environ.get("BIGQUERY_SIMPLE_TABLE")
 DOWNLOAD_TABLE = os.environ.get("BIGQUERY_DOWNLOAD_TABLE")
+MAX_BLOBS_PER_RUN = 5000  # Cannot exceed 10,000
 
 prefix = {Simple.__name__: "simple_requests", Download.__name__: "file_downloads"}
 
@@ -138,12 +139,16 @@ def load_processed_files_into_bigquery(event, context):
 
     # Get the processed files we're loading
     download_prefix = f"{folder}/downloads-"
-    download_source_blobs = bucket.list_blobs(prefix=download_prefix)
+    download_source_blobs = bucket.list_blobs(
+        prefix=download_prefix, max_results=MAX_BLOBS_PER_RUN
+    )
     download_source_uris = [
         f"gs://{blob.bucket.name}/{blob.name}" for blob in download_source_blobs
     ]
     simple_prefix = f"{folder}/simple-"
-    simple_source_blobs = bucket.list_blobs(prefix=simple_prefix)
+    simple_source_blobs = bucket.list_blobs(
+        prefix=simple_prefix, max_results=MAX_BLOBS_PER_RUN
+    )
     simple_source_uris = [
         f"gs://{blob.bucket.name}/{blob.name}" for blob in simple_source_blobs
     ]

test_functions.py

@@ -151,7 +151,7 @@ def test_load_processed_files_into_bigquery(
 
     blob_stub = pretend.stub(name="blobname", bucket=bucket)
 
-    def _generate_blob_list(prefix):
+    def _generate_blob_list(prefix, max_results):
         blob_list = [blob_stub]
         blob_lists[prefix] = blob_list
         return blob_list
@@ -203,8 +203,8 @@ def _generate_blob_list(prefix):
         pretend.call(RESULT_BUCKET),
     ]
     assert bucket_stub.list_blobs.calls == [
-        pretend.call(prefix=f"processed/{partition}/downloads-"),
-        pretend.call(prefix=f"processed/{partition}/simple-"),
+        pretend.call(prefix=f"processed/{partition}/downloads-", max_results=5000),
+        pretend.call(prefix=f"processed/{partition}/simple-", max_results=5000),
     ]
     assert (
         load_job_stub.result.calls