feat: updating geolocation trigger schedule for gbfs feeds + trigger feed processing on addition (#1631)

Author: cka-y

.github/workflows/catalog-update.yml

@@ -8,6 +8,16 @@ on:
         required: false
         default: true
         type: boolean
+      ENVIRONMENT:
+        description: Target environment (leave as 'all' to run all environments)
+        required: false
+        default: all
+        type: choice
+        options:
+          - all
+          - qa
+          - prod
+          - dev
   repository_dispatch: # Update on mobility-database-catalog repo dispatch
     types: [ catalog-sources-updated, gbfs-systems-updated ]
 
@@ -18,6 +28,7 @@ env:
 jobs:
   resolve-api-meta-qa:
     name: QA Resolve API commit/version
+    if: github.event_name == 'repository_dispatch' || inputs.ENVIRONMENT == 'all' || inputs.ENVIRONMENT == 'qa'
     runs-on: ubuntu-latest
     outputs:
       CHECKOUT_REF: ${{ steps.resolve.outputs.COMMIT_SHA != '' && steps.resolve.outputs.COMMIT_SHA || 'main' }}
@@ -57,6 +68,7 @@ jobs:
 
   resolve-api-meta-prod:
     name: PROD Resolve API commit/version
+    if: github.event_name == 'repository_dispatch' || inputs.ENVIRONMENT == 'all' || inputs.ENVIRONMENT == 'prod'
     runs-on: ubuntu-latest
     outputs:
       CHECKOUT_REF: ${{ steps.resolve.outputs.COMMIT_SHA != '' && steps.resolve.outputs.COMMIT_SHA || 'main' }}
@@ -96,6 +108,7 @@ jobs:
 
   resolve-api-meta-dev:
     name: DEV Resolve API commit/version
+    if: github.event_name == 'repository_dispatch' || inputs.ENVIRONMENT == 'all' || inputs.ENVIRONMENT == 'dev'
     runs-on: ubuntu-latest
     outputs:
       CHECKOUT_REF: ${{ steps.resolve.outputs.COMMIT_SHA != '' && steps.resolve.outputs.COMMIT_SHA || 'main' }}

api/src/scripts/populate_db_gbfs.py

@@ -1,8 +1,14 @@
+import json
+import os
+import uuid
+from concurrent import futures
 from datetime import datetime
+from typing import List, Dict
 
 import pandas as pd
 import pycountry
 import pytz
+from google.cloud import pubsub_v1
 
 from scripts.gbfs_utils.comparison import generate_system_csv_from_db, compare_db_to_csv
 from scripts.gbfs_utils.fetching import fetch_data, get_data_content
@@ -11,10 +17,13 @@
 from shared.database.database import generate_unique_id, configure_polymorphic_mappers
 from shared.database_gen.sqlacodegen_models import Gbfsfeed, Location, Externalid
 
+GBFS_PUBSUB_TOPIC_NAME = "validate-gbfs-feed"
+
 
 class GBFSDatabasePopulateHelper(DatabasePopulateHelper):
     def __init__(self, filepaths, test_mode=False):
         super().__init__(filepaths, test_mode)
+        self.added_feeds: List[Dict] = []
 
     def filter_data(self):
         """Filter out rows with Authentication Info and duplicate System IDs"""
@@ -73,6 +82,7 @@ def populate_db(self, session, fetch_url=True):
                 else:
                     fetched_data = dict()
                 # If the feed already exists, update it. Otherwise, create a new feed.
+                is_new_feed = gbfs_feed is None
                 if gbfs_feed:
                     self.logger.info(f"Updating feed {stable_id} - {row['Name']}")
                 else:
@@ -116,6 +126,14 @@ def populate_db(self, session, fetch_url=True):
                     gbfs_feed.locations = [location]
 
                 session.flush()
+                if is_new_feed:
+                    self.added_feeds.append(
+                        {
+                            "feed_id": gbfs_feed.id,
+                            "stable_id": gbfs_feed.stable_id,
+                            "url": gbfs_feed.auto_discovery_url,
+                        }
+                    )
                 self.logger.info(80 * "-")
 
                 # self.db.session.commit()
@@ -125,6 +143,42 @@ def populate_db(self, session, fetch_url=True):
             self.logger.error(f"Error populating the database: {e}")
             raise e
 
+    def trigger_downstream_tasks(self):
+        """Trigger GBFS version extraction, validation, and geolocation for newly added feeds."""
+        self.logger.info(
+            f"Triggering downstream tasks for {len(self.added_feeds)} newly added feed(s): "
+            f"{', '.join(f['stable_id'] for f in self.added_feeds)}"
+        )
+        if os.getenv("ENV", "local") == "local":
+            self.logger.info("Skipping downstream tasks in local environment.")
+            return
+        if not self.added_feeds:
+            self.logger.info("No feeds to trigger downstream tasks for.")
+            return
+
+        env = os.getenv("ENV", "dev")
+        project_id = f"mobility-feeds-{env}"
+        publisher = pubsub_v1.PublisherClient()
+        topic_path = publisher.topic_path(project_id, GBFS_PUBSUB_TOPIC_NAME)
+        execution_id = str(uuid.uuid4())
+
+        publish_futures = []
+        for feed_data in self.added_feeds:
+            message = {
+                "execution_id": execution_id,
+                "stable_id": feed_data["stable_id"],
+                "feed_id": feed_data["feed_id"],
+                "url": feed_data["url"],
+                "extract_geolocation": True,
+            }
+            data = json.dumps(message).encode("utf-8")
+            self.logger.info(f"Publishing feed {feed_data['stable_id']} to {topic_path}")
+            future = publisher.publish(topic_path, data)
+            publish_futures.append(future)
+
+        futures.wait(publish_futures, return_when=futures.ALL_COMPLETED)
+        self.logger.info(f"Published {len(self.added_feeds)} feed(s) to {topic_path}.")
+
 
 if __name__ == "__main__":
-    GBFSDatabasePopulateHelper(set_up_configs()).initialize(trigger_downstream_tasks=False)
+    GBFSDatabasePopulateHelper(set_up_configs()).initialize(trigger_downstream_tasks=True)

functions-python/gbfs_validator/src/gbfs_data_processor.py

@@ -53,7 +53,9 @@ def __init__(self, stable_id: str, feed_id: str):
         self.validation_reports: Dict[str, Dict[str, Any]] = {}
         self.logger = get_logger(GBFSDataProcessor.__name__, stable_id)
 
-    def process_gbfs_data(self, autodiscovery_url: str) -> None:
+    def process_gbfs_data(
+        self, autodiscovery_url: str, extract_geolocation: bool = True
+    ) -> None:
         """Process the GBFS data from the autodiscovery URL."""
         # Record the request to the autodiscovery URL
         self.record_autodiscovery_request(autodiscovery_url)
@@ -70,7 +72,8 @@ def process_gbfs_data(self, autodiscovery_url: str) -> None:
         # Update database entities
         self.update_database_entities()
 
-        self.trigger_location_extraction()
+        if extract_geolocation:
+            self.trigger_location_extraction()
 
     @with_db_session()
     def record_autodiscovery_request(

functions-python/gbfs_validator/src/main.py

@@ -32,9 +32,12 @@
 
 def fetch_all_gbfs_feeds(db_session: Session) -> List[Gbfsfeed]:
     try:
-        gbfs_feeds = (
-            db_session.query(Gbfsfeed).filter(Gbfsfeed.status != "deprecated").all()
-        )
+        query = db_session.query(Gbfsfeed).filter(Gbfsfeed.status != "deprecated")
+        limit = os.getenv("FEEDS_LIMIT")
+        if limit is not None:
+            logging.info("Limiting batch to %s feeds (FEEDS_LIMIT is set).", limit)
+            query = query.limit(int(limit))
+        gbfs_feeds = query.all()
         db_session.expunge_all()
         return gbfs_feeds
     except Exception as e:
@@ -60,6 +63,7 @@ def gbfs_validator_pubsub(cloud_event: CloudEvent):
     except KeyError as e:
         logging.error("Missing required field: %s", e)
         return f"Invalid Pub/Sub message data. Missing {e}."
+    extract_geolocation = message_json.get("extract_geolocation", True)
 
     # get logger with stable_id
     logger = get_logger(__name__, stable_id)
@@ -86,7 +90,7 @@ def gbfs_validator_pubsub(cloud_event: CloudEvent):
     # Process GBFS data
     try:
         processor = GBFSDataProcessor(stable_id, feed_id)
-        processor.process_gbfs_data(url)
+        processor.process_gbfs_data(url, extract_geolocation=extract_geolocation)
     except Exception as e:
         error_message = f"Error processing GBFS data: {e}"
         logger.error(error_message)
@@ -117,11 +121,15 @@ def gbfs_validator_batch(request, db_session: Session):
 
     try:
         feed_stable_ids = None
+        extract_geolocation = True
         if request and request.method == "POST" and request.is_json:
             request_json = request.get_json()
             feed_stable_ids = (
                 request_json.get("feed_stable_ids") if request_json else None
             )
+            extract_geolocation = (
+                request_json.get("extract_geolocation", True) if request_json else True
+            )
         else:
             logging.info("Request body not provided or not a valid JSON.")
     except Exception as e:
@@ -146,6 +154,7 @@ def gbfs_validator_batch(request, db_session: Session):
             "stable_id": gbfs_feed.stable_id,
             "feed_id": gbfs_feed.id,
             "url": gbfs_feed.auto_discovery_url,
+            "extract_geolocation": extract_geolocation,
         }
         feeds_data.append(feed_data)
         logging.info("Feed %s added to the batch.", gbfs_feed.stable_id)

functions-python/gbfs_validator/tests/test_gbfs_data_processor.py

@@ -233,7 +233,9 @@ def test_fetch_gbfs_files(
                     },
                 ],
             }
-            self.processor.process_gbfs_data(autodiscovery_url)
+            self.processor.process_gbfs_data(
+                autodiscovery_url, extract_geolocation=True
+            )
             gbfs_feed = (
                 session.query(Gbfsfeed)
                 .filter_by(stable_id=self.stable_id)
@@ -250,3 +252,64 @@ def test_fetch_gbfs_files(
                 )
             self.assertIn("2.2", versions)
             self.assertIn("2.1", versions)
+
+    @with_db_session(db_url=default_db_url)
+    @patch("gbfs_data_processor.create_http_task")
+    @patch("gbfs_data_processor.tasks_v2")
+    @patch(
+        "gbfs_data_processor.GBFSEndpoint.get_request_metadata",
+        side_effect=mock_get_request_metadata,
+    )
+    @patch("google.cloud.storage.Client")
+    @patch("gbfs_data_processor.fetch_gbfs_data", side_effect=mock_fetch_gbfs_data)
+    @patch("requests.post")
+    @patch("requests.get")
+    @patch.dict(
+        os.environ,
+        {
+            "FEEDS_DATABASE_URL": default_db_url,
+            "GOOGLE_APPLICATION_CREDENTIALS": "test",
+        },
+    )
+    def test_process_gbfs_data_skip_geolocation(
+        self,
+        _,
+        mock_post,
+        __,
+        mock_cloud_storage_client,
+        ___,
+        mock_tasks,
+        mock_create_http_task,
+        db_session,
+    ):
+        """Test that trigger_location_extraction is not called when extract_geolocation=False."""
+        autodiscovery_url = "http://example.com/gbfs.json"
+        gbfs_feed = Gbfsfeed(
+            id=self.feed_id,
+            operator=self.faker.company(),
+            operator_url=self.faker.url(),
+            stable_id=self.stable_id,
+            auto_discovery_url=autodiscovery_url,
+            status="active",
+            operational_status="published",
+        )
+        session = db_session
+        session.add(gbfs_feed)
+        session.commit()
+        (
+            mock_cloud_storage_client.return_value.bucket.return_value.blob.return_value
+        ).public_url = self.faker.url()
+        with patch("logging.info"), patch("logging.error"), patch("logging.warning"):
+            mock_post.return_value.json.return_value = {
+                "summary": {
+                    "validatorVersion": "1.0.13",
+                    "version": {"detected": "2.2", "validated": "2.2"},
+                    "hasErrors": False,
+                    "errorsCount": 0,
+                },
+                "filesSummary": [],
+            }
+            self.processor.process_gbfs_data(
+                autodiscovery_url, extract_geolocation=False
+            )
+        mock_create_http_task.assert_not_called()