fix: batch process dataset function raises out of memory with large feeds (#1320)

Author: cka-y

functions-python/batch_datasets/README.md

@@ -1,5 +1,12 @@
 # Batch Datasets
 This directory contains the GCP serverless function that enqueue all active feeds to download datasets.
+The function accepts an option request body to limit the feeds to process, otherwise it processes all active feeds:
+```json
+{
+    "feed_stable_ids": ["feed_id_1", "feed_id_2"]
+}
+```
+
 The function publish one Pub/Sub message per active feed with the following format:
 ```json
     {

functions-python/batch_datasets/src/main.py

@@ -19,6 +19,7 @@
 import os
 import uuid
 from datetime import datetime
+from typing import Optional
 
 import functions_framework
 from google.cloud import pubsub_v1
@@ -64,7 +65,9 @@ def publish(publisher: PublisherClient, topic_path: str, data_bytes: bytes) -> F
     return publisher.publish(topic_path, data=data_bytes)
 
 
-def get_non_deprecated_feeds(session: Session):
+def get_non_deprecated_feeds(
+    session: Session, feed_stable_ids: Optional[list[str]] = None
+):
     """
     Returns a list of non deprecated feeds
     :return: list of feeds
@@ -79,14 +82,17 @@ def get_non_deprecated_feeds(session: Session):
             Gtfsfeed.authentication_info_url,
             Gtfsfeed.api_key_parameter_name,
             Gtfsfeed.status,
-            Gtfsdataset.id.label("dataset_id"),
+            Gtfsdataset.stable_id.label("dataset_stable_id"),
             Gtfsdataset.hash.label("dataset_hash"),
         )
         .select_from(Gtfsfeed)
         .outerjoin(Gtfsdataset, (Gtfsdataset.feed_id == Gtfsfeed.id))
         .filter(Gtfsfeed.status != "deprecated")
         .filter(or_(Gtfsdataset.id.is_(None), Gtfsdataset.latest.is_(True)))
     )
+    if feed_stable_ids:
+        # If feed_stable_ids are provided, filter the query by stable IDs
+        query = query.filter(Gtfsfeed.stable_id.in_(feed_stable_ids))
     # Limit the query to 10 feeds (or FEEDS_LIMIT param) for testing purposes and lower environments
     if os.getenv("ENVIRONMENT", "").lower() in ("dev", "test", "qa"):
         limit = os.getenv("FEEDS_LIMIT")
@@ -108,8 +114,17 @@ def batch_datasets(request, db_session: Session):
     :param db_session: database session object
     :return: HTTP response object
     """
+    feed_stable_ids = None
     try:
-        feeds = get_non_deprecated_feeds(db_session)
+        request_json = request.get_json()
+        feed_stable_ids = request_json.get("feed_stable_ids") if request_json else None
+    except Exception:
+        logging.info(
+            "No feed_stable_ids provided in the request, processing all feeds."
+        )
+
+    try:
+        feeds = get_non_deprecated_feeds(db_session, feed_stable_ids=feed_stable_ids)
     except Exception as error:
         logging.error(f"Error retrieving feeds: {error}")
         raise Exception(f"Error retrieving feeds: {error}")
@@ -130,7 +145,7 @@ def batch_datasets(request, db_session: Session):
             "producer_url": feed.producer_url,
             "feed_stable_id": feed.stable_id,
             "feed_id": feed.feed_id,
-            "dataset_id": feed.dataset_id,
+            "dataset_stable_id": feed.dataset_stable_id,
             "dataset_hash": feed.dataset_hash,
             "authentication_type": feed.authentication_type,
             "authentication_info_url": feed.authentication_info_url,

functions-python/batch_datasets/tests/test_batch_datasets_main.py

@@ -54,7 +54,47 @@ def test_batch_datasets(mock_client, mock_publish, db_session):
             "shared.dataset_service.main.BatchExecutionService.save",
             return_value=None,
         ):
-            batch_datasets(Mock())
+            mock_request = MagicMock()
+            mock_request.get_json = MagicMock(return_value={})
+            batch_datasets(mock_request)
+            assert mock_publish.call_count == 5
+            # loop over mock_publish.call_args_list and check that the stable_id of the feed is in the list of
+            # active feeds
+            for i in range(3):
+                message = json.loads(
+                    mock_publish.call_args_list[i][0][2].decode("utf-8")
+                )
+                assert message["feed_stable_id"] in [feed.stable_id for feed in feeds]
+
+
+@mock.patch.dict(
+    os.environ,
+    {
+        "FEEDS_DATABASE_URL": default_db_url,
+        "FEEDS_PUBSUB_TOPIC_NAME": "test_topic",
+        "ENVIRONMENT": "test",
+        "FEEDS_LIMIT": "5",
+    },
+)
+@patch("main.publish")
+@patch("main.get_pubsub_client")
+@with_db_session(db_url=default_db_url)
+def test_batch_datasets_w_feed_ids(mock_client, mock_publish, db_session):
+    mock_client.return_value = MagicMock()
+    feeds = get_non_deprecated_feeds(db_session)
+    with patch(
+        "shared.dataset_service.main.BatchExecutionService.__init__",
+        return_value=None,
+    ):
+        with patch(
+            "shared.dataset_service.main.BatchExecutionService.save",
+            return_value=None,
+        ):
+            mock_request = MagicMock()
+            mock_request.get_json = MagicMock(
+                return_value={"feed_stable_ids": [feed.stable_id for feed in feeds]}
+            )
+            batch_datasets(mock_request)
             assert mock_publish.call_count == 5
             # loop over mock_publish.call_args_list and check that the stable_id of the feed is in the list of
             # active feeds

functions-python/batch_process_dataset/function_config.json

@@ -3,7 +3,7 @@
   "description": "Process datasets from the feed passed in the Pub/Sub event",
   "entry_point": "process_dataset",
   "timeout": 540,
-  "memory": "2Gi",
+  "memory": "8Gi",
   "trigger_http": true,
   "include_folders": ["helpers", "dataset_service"],
   "include_api_folders": ["database_gen", "database", "common"],
@@ -20,5 +20,5 @@
   "max_instance_request_concurrency": 1,
   "max_instance_count": 5,
   "min_instance_count": 0,
-  "available_cpu": 1
+  "available_cpu": 2
 }