feat: download feed after creation (#1437)

Author: davidgamez

functions-python/helpers/pub_sub.py

@@ -15,13 +15,19 @@
 #
 import json
 import logging
+import os
 import uuid
 from typing import Dict, List
 
 from google.cloud import pubsub_v1
 from google.cloud.pubsub_v1 import PublisherClient
 from google.cloud.pubsub_v1.publisher.futures import Future
 
+from shared.database_gen.sqlacodegen_models import Feed, Gtfsfeed
+
+PROJECT_ID = os.getenv("PROJECT_ID")
+DATASET_BATCH_TOPIC = os.getenv("DATASET_PROCESSING_TOPIC_NAME")
+
 
 def get_pubsub_client():
     """
@@ -43,7 +49,13 @@ def get_execution_id(request, prefix: str) -> str:
     @param request: HTTP request object
     @param prefix: prefix for the execution ID. Example: "batch-datasets"
     """
-    trace_id = request.headers.get("X-Cloud-Trace-Context")
+    trace_id = (
+        request.headers.get("X-Cloud-Trace-Context")
+        if hasattr(request, "headers")
+        else None
+    )
+    if not trace_id:
+        trace_id = request.trace_id if hasattr(request, "trace_id") else None
     execution_id = f"{prefix}-{trace_id}" if trace_id else f"{prefix}-{uuid.uuid4()}"
     return execution_id
 
@@ -58,3 +70,41 @@ def publish_messages(data: List[Dict], project_id, topic_name) -> None:
         message_data = json.dumps(element).encode("utf-8")
         future = publish(publisher, topic_path, message_data)
         logging.info(f"Published message to Pub/Sub with ID: {future.result()}")
+
+
+def trigger_dataset_download(
+    feed: Feed | Gtfsfeed,
+    execution_id: str,
+    topic_name: str = DATASET_BATCH_TOPIC,
+) -> None:
+    """Publishes the feed to the configured Pub/Sub topic."""
+    publisher = get_pubsub_client()
+    topic_path = publisher.topic_path(PROJECT_ID, topic_name)
+    logging.debug("Publishing to Pub/Sub topic: %s", topic_path)
+
+    message_data = {
+        "execution_id": execution_id,
+        "producer_url": feed.producer_url,
+        "feed_stable_id": feed.stable_id,
+        "feed_id": feed.id,
+        "dataset_id": None,
+        "dataset_hash": None,
+        "authentication_type": feed.authentication_type,
+        "authentication_info_url": feed.authentication_info_url,
+        "api_key_parameter_name": feed.api_key_parameter_name,
+    }
+
+    try:
+        # Convert to JSON string
+        json_message = json.dumps(message_data)
+        future = publisher.publish(topic_path, data=json_message.encode("utf-8"))
+        future.add_done_callback(
+            lambda _: logging.info(
+                "Published feed %s to dataset batch topic", feed.stable_id
+            )
+        )
+        future.result()
+        logging.info("Message published for feed %s", feed.stable_id)
+    except Exception as e:
+        logging.error("Error publishing to dataset batch topic: %s", str(e))
+        raise

functions-python/operations_api/requirements.txt

@@ -24,6 +24,7 @@ uvloop==0.19.0
 
 # Additional packages
 google-cloud-logging==3.10.0
+google-cloud-pubsub
 functions-framework==3.*
 SQLAlchemy==2.0.23
 geoalchemy2==0.14.7

functions-python/operations_api/src/feeds_operations/impl/feeds_operations_impl.py

@@ -13,6 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import JSONResponse
 
@@ -47,13 +48,15 @@
 from feeds_gen.models.update_request_gtfs_rt_feed import (
     UpdateRequestGtfsRtFeed,
 )
+from middleware.request_context_oauth2 import get_request_context
 from shared.database.database import with_db_session, refresh_materialized_view
 from shared.database_gen.sqlacodegen_models import (
     Gtfsfeed,
     t_feedsearch,
     Feed,
     Gtfsrealtimefeed,
 )
+from shared.helpers.pub_sub import get_execution_id, trigger_dataset_download
 from shared.helpers.query_helper import (
     query_feed_by_stable_id,
     get_feeds_query,
@@ -365,6 +368,10 @@ async def create_gtfs_feed(
         db_session.add(new_feed)
         db_session.commit()
         created_feed = db_session.get(Gtfsfeed, new_feed.id)
+        trigger_dataset_download(
+            created_feed,
+            get_execution_id(get_request_context(), "feed-created-process"),
+        )
         logging.info("Created new GTFS feed with ID: %s", new_feed.stable_id)
         payload = OperationGtfsFeedImpl.from_orm(created_feed).model_dump()
         return JSONResponse(status_code=201, content=jsonable_encoder(payload))

functions-python/operations_api/tests/feeds_operations/impl/test_create_feeds_operations_impl_gtfs.py

@@ -22,6 +22,7 @@
 )
 import json
 import uuid
+from unittest.mock import patch
 
 
 @pytest.fixture
@@ -61,7 +62,8 @@ def db_session():
 
 
 @pytest.mark.asyncio
-async def test_create_gtfs_feed_success(db_session):
+@patch("feeds_operations.impl.feeds_operations_impl.trigger_dataset_download")
+async def test_create_gtfs_feed_success(mock_publish_messages, db_session):
     api = OperationsApiImpl()
     unique_url = f"https://new-feed.example.com/{uuid.uuid4()}"
     request = OperationCreateRequestGtfsFeed(
@@ -105,6 +107,24 @@ async def test_create_gtfs_feed_success(db_session):
         assert created.data_type == "gtfs"
         assert created.provider == "New Provider"
         assert created.operational_status == "wip"
+
+        # Assert publish_messages was called exactly once with expected payload
+        assert mock_publish_messages.call_count == 1
+        args, kwargs = mock_publish_messages.call_args
+        assert len(args) == 2  # data list, project_id, topic_name
+        feed, execution_id = args
+
+        # Validate message payload shape and values
+        # assert isinstance(feed, list) and len(data_list) == 1
+        # message = feed[0]
+        assert feed.producer_url == unique_url
+        assert feed.stable_id == payload["stable_id"]
+        assert feed.id == payload["id"]
+        assert feed.authentication_type == "0"
+        assert feed.authentication_info_url is None
+        assert feed.api_key_parameter_name is None
+        # Non-deterministic but must start with expected prefix
+        assert execution_id.startswith("feed-created-process-")
     finally:
         # Cleanup to avoid impacting other tests
         stable_id = payload.get("stable_id") if isinstance(payload, dict) else None

functions-python/tasks_executor/src/tasks/data_import/data_import_utils.py

@@ -39,8 +39,8 @@
     Externalid,
     Officialstatushistory,
 )
-from tasks.data_import.data_import_utils import trigger_dataset_download
-from google.cloud import pubsub_v1
+
+from shared.helpers.pub_sub import trigger_dataset_download
 
 T = TypeVar("T", bound="Feed")
 
@@ -648,9 +648,8 @@ def commit_changes(
         logger.info("Commit after processing items (count=%d)", total_processed)
         db_session.commit()
         execution_id = str(uuid.uuid4())
-        publisher = pubsub_v1.PublisherClient()
         for feed in feeds_to_publish:
-            trigger_dataset_download(feed, execution_id, publisher)
+            trigger_dataset_download(feed, execution_id)
     except IntegrityError:
         db_session.rollback()
         logger.exception("Commit failed with IntegrityError; rolled back")

functions-python/tasks_executor/src/tasks/data_import/import_jbda_feeds.py

@@ -1,8 +1,7 @@
 import os
-import json
 import unittest
 from typing import Any, Dict, List
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 
 from sqlalchemy.orm import Session
 
@@ -205,8 +204,6 @@ def _head_side_effect(url, allow_redirects=True, timeout=15):
 class TestImportJBDA(unittest.TestCase):
     @with_db_session(db_url=default_db_url)
     def test_import_creates_gtfs_rt_and_related_links(self, db_session: Session):
-        fake_pub = _FakePublisher()
-
         # The importer will call HEAD on these URLs for org1/feed1
         base = (
             "https://api.gtfs-data.jp/v2/organizations/org1/feeds/feed1/files/feed.zip"
@@ -221,6 +218,8 @@ def _head_side_effect(url, allow_redirects=True, timeout=15):
             # fail for anything else (e.g., feed3 current)
             return _FakeResponse(status=404)
 
+        # Patch requests.Session and head; replace old pubsub mocks with trigger_dataset_download
+        mock_trigger = MagicMock()
         with patch(
             "tasks.data_import.import_jbda_feeds.requests.Session",
             return_value=_FakeSessionOK(),
@@ -230,12 +229,8 @@ def _head_side_effect(url, allow_redirects=True, timeout=15):
         ), patch(
             "tasks.data_import.import_jbda_feeds.REQUEST_TIMEOUT_S", 0.01
         ), patch(
-            "tasks.data_import.import_jbda_feeds.pubsub_v1.PublisherClient",
-            return_value=fake_pub,
-        ), patch(
-            "tasks.data_import.data_import_utils.PROJECT_ID", "test-project"
-        ), patch(
-            "tasks.data_import.data_import_utils.DATASET_BATCH_TOPIC", "dataset-batch"
+            "tasks.data_import.import_jbda_feeds.trigger_dataset_download",
+            mock_trigger,
         ), patch.dict(
             os.environ, {"COMMIT_BATCH_SIZE": "1"}, clear=False
         ):
@@ -262,6 +257,9 @@ def _head_side_effect(url, allow_redirects=True, timeout=15):
             .first()
         )
         self.assertIsNotNone(sched)
+        # ensure the instance is attached to this test session before accessing relationships
+        sched = db_session.merge(sched)
+
         self.assertEqual(sched.feed_name, "Feed One")
         # producer_url now points to the verified JBDA URL (HEAD-checked)
         self.assertEqual(sched.producer_url, url_current)
@@ -273,7 +271,7 @@ def _head_side_effect(url, allow_redirects=True, timeout=15):
         next1 = next(link for link in links if link.code == "jbda-next_1")
         self.assertEqual(next1.url, url_next1)
 
-        # RT feeds + entity types + back-links
+        # RT feeds + entity types & back-links
         tu = (
             db_session.query(Gtfsrealtimefeed)
             .filter(Gtfsrealtimefeed.stable_id == "jbda-org1-feed1-tu")
@@ -284,6 +282,11 @@ def _head_side_effect(url, allow_redirects=True, timeout=15):
             .filter(Gtfsrealtimefeed.stable_id == "jbda-org1-feed1-vp")
             .first()
         )
+
+        # merge to ensure attached before accessing lazy attrs
+        tu = db_session.merge(tu)
+        vp = db_session.merge(vp)
+
         self.assertIsNotNone(tu)
         self.assertIsNotNone(vp)
         self.assertEqual(len(tu.entitytypes), 1)
@@ -295,16 +298,15 @@ def _head_side_effect(url, allow_redirects=True, timeout=15):
         self.assertEqual(tu.producer_url, "https://rt.example/one/tu.pb")
         self.assertEqual(vp.producer_url, "https://rt.example/one/vp.pb")
 
-        # Pub/Sub was called exactly once (only 1 new GTFS feed)
-        self.assertEqual(len(fake_pub.published), 1)
-        topic_path, data_bytes = fake_pub.published[0]
-        self.assertEqual(topic_path, "projects/test-project/topics/dataset-batch")
-
-        payload = json.loads(data_bytes.decode("utf-8"))
-        self.assertEqual(payload["feed_stable_id"], "jbda-org1-feed1")
-        self.assertEqual(payload["producer_url"], url_current)
-        self.assertIsNone(payload["dataset_id"])
-        self.assertIsNone(payload["dataset_hash"])
+        # trigger_dataset_download (new API) should have been called exactly once for the new GTFS feed
+        mock_trigger.assert_called_once()
+        called_args = mock_trigger.call_args[0]
+        detached_feed = called_args[0]
+        merged_feed = db_session.merge(detached_feed)
+        # self.assertGreaterEqual(len(called_args), 1)
+        self.assertEqual(getattr(merged_feed, "stable_id", None), "jbda-org1-feed1")
+        # second arg should be an execution id (string)
+        self.assertIsInstance(called_args[1], str)
 
     @with_db_session(db_url=default_db_url)
     def test_import_http_failure_graceful(self, db_session: Session):

functions-python/tasks_executor/tests/tasks/data_import/test_jbda_import.py

@@ -723,6 +723,7 @@ resource "google_cloudfunctions2_function" "operations_api" {
       PROJECT_ID = var.project_id
       PYTHONNODEBUGRANGES = 0
       GOOGLE_CLIENT_ID = var.operations_oauth2_client_id
+      DATASET_PROCESSING_TOPIC_NAME = "datasets-batch-topic-${var.environment}"
     }
     available_memory = local.function_operations_api_config.memory
     timeout_seconds = local.function_operations_api_config.timeout