feat: extract all stop locations from stops.txt (#921)

Author: cka-y

api/pyproject.toml

@@ -19,7 +19,7 @@ extend-exclude = '''
     | build
     | dist
     | src/feeds_gen/*
-    | src/database_gen/*
+    | src/shared/database_gen/*
   )/
 )
 '''

api/src/feeds/impl/feeds_api_impl.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import List, Union, TypeVar
+from typing import List, Union, TypeVar, Optional
 
 from sqlalchemy import or_
 from sqlalchemy import select
@@ -44,7 +44,6 @@
 from utils.date_utils import valid_iso_date
 from utils.location_translation import (
     create_location_translation_object,
-    LocationTranslation,
     get_feeds_location_translations,
 )
 from utils.logger import Logger
@@ -130,42 +129,39 @@ def get_feeds(
     @with_db_session
     def get_gtfs_feed(self, id: str, db_session: Session) -> GtfsFeed:
         """Get the specified gtfs feed from the Mobility Database."""
-        feed, translations = self._get_gtfs_feed(id, db_session)
+        feed = self._get_gtfs_feed(id, db_session)
         if feed:
-            return GtfsFeedImpl.from_orm(feed, translations)
+            return GtfsFeedImpl.from_orm(feed)
         else:
             raise_http_error(404, gtfs_feed_not_found.format(id))
 
     @staticmethod
-    def _get_gtfs_feed(stable_id: str, db_session: Session) -> tuple[Gtfsfeed | None, dict[str, LocationTranslation]]:
+    def _get_gtfs_feed(stable_id: str, db_session: Session) -> Optional[Gtfsfeed]:
         results = (
             FeedFilter(
                 stable_id=stable_id,
                 status=None,
                 provider__ilike=None,
                 producer_url__ilike=None,
             )
-            .filter(db_session.query(Gtfsfeed, t_location_with_translations_en))
+            .filter(db_session.query(Gtfsfeed))
             .filter(
                 or_(
                     Gtfsfeed.operational_status == None,  # noqa: E711
                     Gtfsfeed.operational_status != "wip",
                     not is_user_email_restricted(),  # Allow all feeds to be returned if the user is not restricted
                 )
             )
-            .outerjoin(Location, Feed.locations)
-            .outerjoin(t_location_with_translations_en, Location.id == t_location_with_translations_en.c.location_id)
             .options(
                 joinedload(Gtfsfeed.gtfsdatasets)
                 .joinedload(Gtfsdataset.validation_reports)
                 .joinedload(Validationreport.notices),
                 *get_joinedload_options(),
             )
         ).all()
-        if len(results) > 0 and results[0].Gtfsfeed:
-            translations = {result[1]: create_location_translation_object(result) for result in results}
-            return results[0].Gtfsfeed, translations
-        return None, {}
+        if len(results) == 0:
+            return None
+        return results[0]
 
     @with_db_session
     def get_gtfs_feed_datasets(
@@ -393,8 +389,8 @@ def _get_response(feed_query: Query, impl_cls: type[T], db_session: "Session") -
     @with_db_session
     def get_gtfs_feed_gtfs_rt_feeds(self, id: str, db_session: Session) -> List[GtfsRTFeed]:
         """Get a list of GTFS Realtime related to a GTFS feed."""
-        feed, translations = self._get_gtfs_feed(id, db_session)
+        feed = self._get_gtfs_feed(id, db_session)
         if feed:
-            return [GtfsRTFeedImpl.from_orm(gtfs_rt_feed, translations) for gtfs_rt_feed in feed.gtfs_rt_feeds]
+            return [GtfsRTFeedImpl.from_orm(gtfs_rt_feed) for gtfs_rt_feed in feed.gtfs_rt_feeds]
         else:
             raise_http_error(404, gtfs_feed_not_found.format(id))

functions-python/.flake8

@@ -1,5 +1,5 @@
 [flake8]
 max-line-length = 120
-exclude = .git,__pycache__,__init__.py,.mypy_cache,.pytest_cache,venv,build,.*,database_gen,feeds_operations_gen
+exclude = .git,__pycache__,__init__.py,.mypy_cache,.pytest_cache,venv,build,.*,database_gen,feeds_operations_gen,shared
 # Ignored because conflict with black
 extend-ignore = E203

functions-python/extract_location/src/main.py

@@ -7,10 +7,6 @@
 
 import functions_framework
 from cloudevents.http import CloudEvent
-from google.cloud import pubsub_v1
-from sqlalchemy import or_
-from sqlalchemy.orm import joinedload
-
 from shared.database_gen.sqlacodegen_models import Gtfsdataset
 from shared.dataset_service.main import (
     DatasetTraceService,
@@ -22,10 +18,14 @@
 from shared.helpers.database import Database
 from shared.helpers.logger import Logger
 from shared.helpers.parser import jsonify_pubsub
+from sqlalchemy import or_
+from sqlalchemy.orm import joinedload
+
 from bounding_box.bounding_box_extractor import (
     create_polygon_wkt_element,
     update_dataset_bounding_box,
 )
+from shared.helpers.pub_sub import publish_messages
 from reverse_geolocation.location_extractor import update_location, reverse_coords
 from stops_utils import get_gtfs_feed_bounds_and_points
 
@@ -232,11 +232,5 @@ def extract_location_batch(_):
         pass
 
     # Trigger update location for each dataset by publishing to Pub/Sub
-    publisher = pubsub_v1.PublisherClient()
-    topic_path = publisher.topic_path(os.getenv("PROJECT_ID"), pubsub_topic_name)
-    for data in datasets_data:
-        message_data = json.dumps(data).encode("utf-8")
-        future = publisher.publish(topic_path, message_data)
-        logging.info(f"Published message to Pub/Sub with ID: {future.result()}")
-
+    publish_messages(datasets_data, os.getenv("PROJECT_ID"), pubsub_topic_name)
     return f"Batch function triggered for {len(datasets_data)} datasets.", 200

functions-python/extract_location/tests/test_location_extraction.py

@@ -267,11 +267,11 @@ def test_extract_location_exception_2(
         },
     )
     @patch("main.Database")
-    @patch("main.pubsub_v1.PublisherClient")
+    @patch("main.publish_messages")
     @patch("main.Logger")
     @patch("uuid.uuid4")
     def test_extract_location_batch(
-        self, uuid_mock, logger_mock, publisher_client_mock, database_mock
+        self, uuid_mock, logger_mock, publish_messages_mock, database_mock
     ):
         mock_session = MagicMock()
         mock_dataset1 = Gtfsdataset(
@@ -302,37 +302,10 @@ def test_extract_location_batch(
             mock_session
         )
 
-        mock_publisher = MagicMock()
-        publisher_client_mock.return_value = mock_publisher
-        mock_future = MagicMock()
-        mock_future.result.return_value = "message_id"
-        mock_publisher.publish.return_value = mock_future
-
         response = extract_location_batch(None)
 
         logger_mock.init_logger.assert_called_once()
-        mock_publisher.publish.assert_any_call(
-            mock.ANY,
-            json.dumps(
-                {
-                    "stable_id": "1",
-                    "dataset_id": "stable_1",
-                    "url": "http://example.com/1",
-                    "execution_id": "batch-uuid",
-                }
-            ).encode("utf-8"),
-        )
-        mock_publisher.publish.assert_any_call(
-            mock.ANY,
-            json.dumps(
-                {
-                    "stable_id": "2",
-                    "dataset_id": "stable_2",
-                    "url": "http://example.com/2",
-                    "execution_id": "batch-uuid",
-                }
-            ).encode("utf-8"),
-        )
+        publish_messages_mock.assert_called_once()
         self.assertEqual(response, ("Batch function triggered for 2 datasets.", 200))
 
     @mock.patch.dict(

functions-python/helpers/database.py

@@ -18,6 +18,7 @@
 import logging
 import os
 import threading
+from functools import wraps
 from typing import Optional, ContextManager
 
 from sqlalchemy import create_engine, text, event, Engine
@@ -85,7 +86,7 @@ def mapper_configure_listener(mapper, class_):
 event.listen(mapper, "mapper_configured", mapper_configure_listener)
 
 
-def with_db_session(func):
+def with_db_session(_func=None, *, echo=True):
     """
     Decorator to handle the session management for the decorated function.
 
@@ -103,16 +104,23 @@ def with_db_session(func):
         - If 'db_session' is already provided, it simply calls the decorated function with the existing session.
     """
 
-    def wrapper(*args, **kwargs):
-        db_session = kwargs.get("db_session")
-        if db_session is None:
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            if "db_session" in kwargs:
+                return func(*args, **kwargs)
+
             db = Database()
-            with db.start_db_session() as session:
+            with db.start_db_session(echo=echo) as session:
                 kwargs["db_session"] = session
                 return func(*args, **kwargs)
-        return func(*args, **kwargs)
 
-    return wrapper
+        return wrapper
+
+    if _func is None:
+        return decorator
+    else:
+        return decorator(_func)
 
 
 class Database:

functions-python/helpers/logger.py

@@ -13,6 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import os
 
 import google.cloud.logging
 from google.cloud.logging_v2 import Client
@@ -47,6 +48,8 @@ def init_logger() -> Client:
         """
         Initializes the logger
         """
+        if os.getenv("DEBUG", "False") == "True":
+            return None
         client = google.cloud.logging.Client()
         client.get_default_handler()
         client.setup_logging()

functions-python/helpers/pub_sub.py

@@ -13,7 +13,10 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import json
+import logging
 import uuid
+from typing import Dict, List
 
 from google.cloud import pubsub_v1
 from google.cloud.pubsub_v1 import PublisherClient
@@ -43,3 +46,15 @@ def get_execution_id(request, prefix: str) -> str:
     trace_id = request.headers.get("X-Cloud-Trace-Context")
     execution_id = f"{prefix}-{trace_id}" if trace_id else f"{prefix}-{uuid.uuid4()}"
     return execution_id
+
+
+def publish_messages(data: List[Dict], project_id, topic_name) -> None:
+    """
+    Publishes the given data to the Pub/Sub topic.
+    """
+    publisher = get_pubsub_client()
+    topic_path = publisher.topic_path(project_id, topic_name)
+    for element in data:
+        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()}")

functions-python/helpers/requirements.txt

@@ -25,5 +25,5 @@ google-api-core
 google-cloud-firestore
 google-cloud-bigquery
 
-#Additional package
-pycountry
+# Additional package
+pycountry

functions-python/reverse_geolocation/.coveragerc

@@ -0,0 +1,9 @@
+[run]
+omit =
+    */test*/*
+    */helpers/*
+    */database_gen/*
+
+[report]
+exclude_lines =
+    if __name__ == .__main__.: