Polish and add code comments

Author: leplatrem

commands/__init__.py

@@ -1,3 +1,4 @@
+import concurrent.futures
 import os
 
 import backoff
@@ -70,3 +71,11 @@ def records_equal(a, b):
     ra = {k: v for k, v in a.items() if k not in ignored_fields}
     rb = {k: v for k, v in b.items() if k not in ignored_fields}
     return ra == rb
+
+
+def call_parallel(func, args_list, max_workers=4):
+    results = []
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        futures = [executor.submit(func, *args) for args in args_list]
+        results = [future.result() for future in futures]
+    return results

commands/build_bundles.py

@@ -1,4 +1,10 @@
-import concurrent.futures
+"""
+This command will create Zip files in order to bundle all collections data,
+and all attachments of collections that have the `attachment.bundle` flag in
+their metadata.
+It then uploads these zip files to Google Cloud Storage.
+"""
+
 import io
 import json
 import os
@@ -8,27 +14,26 @@
 import requests
 from google.cloud import storage
 
-from . import KintoClient, retry_timeout
+from . import KintoClient, call_parallel, retry_timeout
 
 
 SERVER = os.getenv("SERVER")
-REQUESTS_PARALLEL_COUNT = int(os.getenv("REQUESTS_PARALLEL_COUNT", "4"))
+REQUESTS_PARALLEL_COUNT = int(os.getenv("REQUESTS_PARALLEL_COUNT", "8"))
 BUNDLE_MAX_SIZE_BYTES = int(os.getenv("BUNDLE_MAX_SIZE_BYTES", "20_000_000"))
-BUILD_ALL = os.getenv("BUILD_ALL", "0") in "1yY"
-STORAGE_BUCKET_NAME = os.getenv("STORAGE_BUCKET_NAME", "rs-attachments")
+STORAGE_BUCKET_NAME = os.getenv("STORAGE_BUCKET_NAME", "remote-settings-nonprod-stage-attachments")
 DESTINATION_FOLDER = os.getenv("DESTINATION_FOLDER", "bundles")
+# Flags for local development
+BUILD_ALL = os.getenv("BUILD_ALL", "0") in "1yY"
 SKIP_UPLOAD = os.getenv("SKIP_UPLOAD", "0") in "1yY"
 
 
-def call_parallel(func, args_list):
-    results = []
-    with concurrent.futures.ThreadPoolExecutor(max_workers=REQUESTS_PARALLEL_COUNT) as executor:
-        futures = [executor.submit(func, *args) for args in args_list]
-        results = [future.result() for future in futures]
-    return results
-
-
 def fetch_all_changesets(client):
+    """
+    Return the `/changeset` responses for all collections listed
+    in the `monitor/changes` endpoint.
+    The result contains the metadata and all the records of all collections
+    for both preview and main buckets.
+    """
     random_cache_bust = random.randint(999999000000, 999999999999)
     monitor_changeset = client.get_changeset("monitor", "changes", random_cache_bust)
     print("%s collections" % len(monitor_changeset["changes"]))
@@ -37,7 +42,7 @@ def fetch_all_changesets(client):
         (c["bucket"], c["collection"], c["last_modified"]) for c in monitor_changeset["changes"]
     ]
     all_changesets = call_parallel(
-        lambda bid, cid, ts: client.get_changeset(bid, cid, ts), args_list
+        lambda bid, cid, ts: client.get_changeset(bid, cid, ts), args_list, REQUESTS_PARALLEL_COUNT
     )
     return [
         {"bucket": bid, **changeset} for (bid, _, _), changeset in zip(args_list, all_changesets)
@@ -51,7 +56,11 @@ def fetch_attachment(url):
     return resp.content
 
 
-def write_zip(output_path, content):
+def write_zip(output_path: str, content: list[tuple[str, bytes]]):
+    """
+    Write a Zip at the specified `output_path` location with the specified `content`.
+    The content is specified as a list of file names and their binary content.
+    """
     parent_folder = os.path.dirname(output_path)
     os.makedirs(parent_folder, exist_ok=True)
 
@@ -64,11 +73,15 @@ def write_zip(output_path, content):
     print("Wrote %r" % output_path)
 
 
-def sync_cloud_storage(folder):
+def sync_cloud_storage(folder, storage_bucket):
+    """
+    Synchronizes a local folder (eg. `bundles/`) with a remote one in the specified
+    `storage_bucket` name.
+    """
     # Ensure you have set the GOOGLE_APPLICATION_CREDENTIALS environment variable
     # to the path of your Google Cloud service account key file before running this script.
     client = storage.Client()
-    bucket = client.bucket(STORAGE_BUCKET_NAME)
+    bucket = client.bucket(storage_bucket)
     local_files = set()
     for root, _, files in os.walk(folder):
         for file in files:
@@ -77,17 +90,25 @@ def sync_cloud_storage(folder):
 
             blob = bucket.blob(remote_file_path)
             blob.upload_from_filename(local_file_path)
-            print(f"Uploaded {local_file_path} to gs://{STORAGE_BUCKET_NAME}/{remote_file_path}")
+            print(f"Uploaded {local_file_path} to gs://{storage_bucket}/{remote_file_path}")
             local_files.add(remote_file_path)
 
     blobs = bucket.list_blobs(prefix=folder)
     for blob in blobs:
         if blob.name not in local_files:
             blob.delete()
-            print(f"Deleted gs://{STORAGE_BUCKET_NAME}/{blob.name}")
+            print(f"Deleted gs://{storage_bucket}/{blob.name}")
 
 
 def build_bundles(event, context):
+    """
+    Main command entry point that:
+    - fetches all collections changesets
+    - builds a `bundles/changesets.zip`
+    - fetches attachments of all collections with bundle flag
+    - builds `bundles/{bid}--{cid}.zip` for each of them
+    - synchronizes the `bundles/` folder with a remote Cloud storage bucket
+    """
     rs_server = event.get("server") or SERVER
 
     client = KintoClient(server_url=rs_server)
@@ -127,12 +148,12 @@ def build_bundles(event, context):
 
         # Fetch all attachments and build "{bid}--{cid}.zip"
         args_list = [(f'{base_url}{r["attachment"]["location"]}',) for r in records]
-        all_attachments = call_parallel(fetch_attachment, args_list)
+        all_attachments = call_parallel(fetch_attachment, args_list, REQUESTS_PARALLEL_COUNT)
         write_zip(
             f"{DESTINATION_FOLDER}/{bid}--{cid}.zip",
             [(f'{record["id"]}.meta.json', json.dumps(record)) for record in records]
             + [(record["id"], attachment) for record, attachment in zip(records, all_attachments)],
         )
 
     if not SKIP_UPLOAD:
-        sync_cloud_storage(DESTINATION_FOLDER)
+        sync_cloud_storage(DESTINATION_FOLDER, STORAGE_BUCKET_NAME)

tests/test_build_bundles.py

@@ -129,21 +129,36 @@ def test_build_bundles(mock_fetch_all_changesets, mock_write_zip, mock_sync_clou
     mock_fetch_all_changesets.return_value = [
         {
             "bucket": "bucket1",
-            "changes": [{"id": "record1", "attachment": {"location": "file.jpg", "size": 10}}],
+            "changes": [
+                {"id": "record1", "attachment": {"location": "file.jpg", "size": 10}},
+                {"id": "record2"},
+            ],
             "metadata": {"id": "collection1", "attachment": {"bundle": True}},
-        }
+        },
+        {  # collection without bundle flag
+            "bucket": "bucket2",
+            "changes": [{"id": "record2"}],
+            "metadata": {"id": "collection2"},
+        },
+        {  # collection without attachments
+            "bucket": "bucket3",
+            "changes": [{"id": "record3"}],
+            "metadata": {"id": "collection3", "attachment": {"bundle": True}},
+        },
     ]
 
     build_bundles(event, context={})
 
-    assert mock_write_zip.call_count == 2  # One for changesets and one for the attachments
+    assert mock_write_zip.call_count == 2  # One for changesets and only one for the attachments
     calls = mock_write_zip.call_args_list
 
     # Assert the first call (changesets.zip)
     changesets_zip_path, changesets_zip_files = calls[0][0]
     assert changesets_zip_path == "bundles/changesets.zip"
-    assert len(changesets_zip_files) == 1
+    assert len(changesets_zip_files) == 3
     assert changesets_zip_files[0][0] == "bucket1--collection1.json"
+    assert changesets_zip_files[1][0] == "bucket2--collection2.json"
+    assert changesets_zip_files[2][0] == "bucket3--collection3.json"
 
     # Assert the second call (attachments zip)
     attachments_zip_path, attachments_zip_files = calls[1][0]
@@ -153,7 +168,9 @@ def test_build_bundles(mock_fetch_all_changesets, mock_write_zip, mock_sync_clou
     assert attachments_zip_files[1][0] == "record1"
     assert attachments_zip_files[1][1] == b"jpeg_content"
 
-    mock_sync_cloud_storage.assert_called_once_with("bundles")
+    mock_sync_cloud_storage.assert_called_once_with(
+        "bundles", "remote-settings-nonprod-stage-attachments"
+    )
 
 
 @patch("commands.build_bundles.os.walk")
@@ -173,7 +190,7 @@ def test_sync_cloud_storage_upload_and_delete(mock_os_walk, mock_storage_client,
     mock_blob3.name = f"{local_folder}/file3.txt"
     bucket.list_blobs.return_value = [mock_blob1, mock_blob2, mock_blob3]
 
-    sync_cloud_storage(local_folder)
+    sync_cloud_storage(local_folder, "remote-bucket")
 
     # Check uploads
     mock_blob1.upload_from_filename.assert_called_once_with(f"{local_folder}/file1.txt")