add support for filtering feed ids in the gbfs batch function

Author: davidgamez

functions-python/gbfs_validator/src/main.py

@@ -101,9 +101,11 @@ def gbfs_validator_pubsub(cloud_event: CloudEvent):
 
 @with_db_session
 @functions_framework.http
-def gbfs_validator_batch(_, db_session: Session):
+def gbfs_validator_batch(request, db_session: Session):
     """
-    HTTP Cloud Function to trigger the GBFS Validator function for multiple datasets.
+    HTTP Cloud  Function to trigger the GBFS Validator function for multiple datasets.
+    When the request is a POST request with a JSON body containing `feed_ids`,
+    it processes only those feeds. Otherwise, it processes all feeds in the database.
     @param _: The request object.
     @return: The response of the function.
     """
@@ -113,9 +115,23 @@ def gbfs_validator_batch(_, db_session: Session):
         logging.error("PUBSUB_TOPIC_NAME environment variable not set.")
         return "PUBSUB_TOPIC_NAME environment variable not set.", 500
 
-    # Get all GBFS feeds from the database
     try:
-        gbfs_feeds = fetch_all_gbfs_feeds(db_session)
+        feed_ids = None
+        if request and request.method == "POST" and request.is_json:
+            request_json = request.get_json()
+            feed_ids = request_json.get("feed_ids") if request_json else None
+        else:
+            logging.info("Request body not provided or not a valid JSON.")
+    except Exception as e:
+        logging.error("Error parsing request body: %s", e)
+        return "Invalid request body.", 400
+
+    try:
+        if feed_ids:
+            gbfs_feeds = fetch_gbfs_feeds_by_ids(db_session, feed_ids)
+        else:
+            # Get all GBFS feeds from the database
+            gbfs_feeds = fetch_all_gbfs_feeds(db_session)
     except Exception:
         return "Error getting all GBFS feeds.", 500
 
@@ -150,3 +166,13 @@ def gbfs_validator_batch(_, db_session: Session):
         f"GBFS Validator batch function triggered successfully for {len(feeds_data)} feeds.",
         200,
     )
+
+
+def fetch_gbfs_feeds_by_ids(db_session, feed_ids):
+    """Fetch GBFS feeds by their IDs from the database."""
+    gbfs_feeds = (
+        db_session.query(Gbfsfeed)
+        .filter(Gbfsfeed.id.in_(feed_ids), Gbfsfeed.status != "deprecated")
+        .all()
+    )
+    return gbfs_feeds

functions-python/gbfs_validator/tests/test_gbfs_validator.py

@@ -179,3 +179,40 @@ def test_gbfs_validator_batch_publish_exception(
         # Call the function
         result = gbfs_validator_batch(None)
         self.assertEqual(result[1], 500)
+
+    @patch.dict(
+        os.environ,
+        {
+            "PUBSUB_TOPIC_NAME": "mock-topic",
+        },
+    )
+    @patch("main.pubsub_v1.PublisherClient")
+    @patch("main.fetch_gbfs_feeds_by_ids")
+    def test_gbfs_validator_batch_by_feed_ids(
+        self, mock_fetch_gbfs_feeds_by_ids, mock_publisher_client
+    ):
+        # Prepare mocks
+        mock_session = MagicMock()
+        # mock_database.return_value.start_db_session.return_value = mock_session
+
+        mock_publisher = MagicMock()
+        mock_publisher_client.return_value = mock_publisher
+
+        mock_feed = MagicMock()
+        mock_feed.stable_id = "mock-stable-id"
+        mock_feed.id = str(uuid.uuid4())
+        mock_feed.auto_discovery_url = "http://mock-url.com"
+        mock_feed.gbfsversions = [MagicMock(version="1.0")]
+        mock_feed_2 = copy.deepcopy(mock_feed)
+        mock_feed_2.gbfsversions = []
+        mock_fetch_gbfs_feeds_by_ids.return_value = [mock_feed, mock_feed_2]
+        request = MagicMock()
+        request.method = "POST"
+        request.is_json = True
+        request.get_json.return_value = {"feed_ids": [mock_feed.id, mock_feed_2.id]}
+        # Call the function
+        result = gbfs_validator_batch(request, db_session=mock_session)
+        self.assertEqual(result[1], 200)
+
+        mock_fetch_gbfs_feeds_by_ids.assert_called_once()
+        self.assertEqual(mock_publisher.publish.call_count, 2)