[fun] location extraction process

Author: cka-y

api/.flake8

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

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
@@ -43,7 +43,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
@@ -129,42 +128,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),
                 *BasicFeedImpl.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(
@@ -389,8 +385,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/helpers/database.py

@@ -19,6 +19,7 @@
 import os
 import threading
 from typing import Optional
+from functools import wraps
 
 from sqlalchemy import create_engine, text, event, Engine
 from sqlalchemy.orm import sessionmaker, Session, mapper, class_mapper
@@ -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,26 @@ 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:
-            db = Database()
-            with db.start_db_session() as session:
-                kwargs["db_session"] = session
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            db_session = kwargs.get("db_session")
+
+            if db_session is None:
+                db = Database()
+                with db.start_db_session(echo=echo) as session:
+                    kwargs["db_session"] = session
+                    return func(*args, **kwargs)
+            else:
                 return func(*args, **kwargs)
-        return func(*args, **kwargs)
 
-    return wrapper
+        return wrapper
+
+    # Allow decorator to be used with or without parentheses
+    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,9 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import json
 import uuid
+from typing import Dict, List
 
 from google.cloud import pubsub_v1
 from google.cloud.pubsub_v1 import PublisherClient
@@ -43,3 +45,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)
+        print(f"Published message to Pub/Sub with ID: {future.result()}")

functions-python/helpers/utils.py

@@ -17,6 +17,7 @@
 import hashlib
 import logging
 import os
+from typing import List
 
 import requests
 import urllib3
@@ -40,6 +41,37 @@ def create_bucket(bucket_name):
         logging.info(f"Bucket {bucket_name} already exists.")
 
 
+def cors_configuration(bucket_name, origin=["*"], method=["GET"], max_age_seconds=3600):
+    """Set a bucket's CORS policies configuration."""
+    storage_client = storage.Client()
+    bucket = storage_client.get_bucket(bucket_name)
+    bucket.cors = [
+        {
+            "origin": origin,
+            "responseHeader": [
+                "Content-Type",
+                "Access-Control-Allow-Origin",
+                "x-goog-resumable",
+            ],
+            "method": method,
+            "maxAgeSeconds": max_age_seconds,
+        }
+    ]
+    bucket.patch()
+    logging.info(
+        f"CORS policy for {bucket_name} set to allow requests from any origin."
+    )
+
+
+def list_blobs(bucket_name: str, prefix: str = "", suffix="") -> List[storage.Blob]:
+    """
+    List all files in a GCP bucket with the given prefix and suffix.
+    """
+    storage_client = storage.Client()
+    blobs = list(storage_client.list_blobs(bucket_name, prefix=prefix))
+    return [blob for blob in blobs if blob.name.endswith(suffix)]
+
+
 def download_url_content(url, with_retry=False):
     """
     Downloads the content of a URL

functions-python/reverse_geolocation/.coveragerc

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

functions-python/reverse_geolocation/README.md

@@ -0,0 +1 @@
+# todo