feat: reverse geolocation per polygon strategy (#1318)

Author: davidgamez

api/src/shared/common/gcp_utils.py

@@ -1,8 +1,13 @@
+import json
 import logging
 import os
 from google.cloud import tasks_v2
 from google.protobuf.timestamp_pb2 import Timestamp
 
+REFRESH_VIEW_TASK_EXECUTOR_BODY = json.dumps(
+    {"task": "refresh_materialized_view", "payload": {"dry_run": False}}
+).encode()
+
 
 def create_refresh_materialized_view_task():
     """
@@ -39,20 +44,19 @@ def create_refresh_materialized_view_task():
         logging.debug("Queue name from env: %s", queue)
         gcp_region = os.getenv("GCP_REGION")
         environment_name = os.getenv("ENVIRONMENT")
-        url = f"https://{gcp_region}-" f"{project}.cloudfunctions.net/" f"tasks-executor-{environment_name}"
-
+        url = f"https://{gcp_region}-" f"{project}.cloudfunctions.net/" f"tasks_executor-{environment_name}"
         # Enqueue the task
         try:
             create_http_task_with_name(
                 client=tasks_v2.CloudTasksClient(),
-                body=b"",
+                body=REFRESH_VIEW_TASK_EXECUTOR_BODY,
                 url=url,
                 project_id=project,
                 gcp_region=gcp_region,
                 queue_name=queue,
                 task_name=task_name,
                 task_time=proto_time,
-                http_method=tasks_v2.HttpMethod.GET,
+                http_method=tasks_v2.HttpMethod.POST,
             )
             logging.info("Scheduled refresh materialized view task for %s", task_name)
             return {"message": "Refresh task for %s scheduled." % task_name}, 200
@@ -95,10 +99,12 @@ def create_http_task_with_name(
             headers={"Content-Type": "application/json"},
         ),
     )
-    logging.info("Task created with task_name: %s", task_name)
     try:
         response = client.create_task(parent=parent, task=task)
+        logging.info("Task created with task_name: %s", task_name)
     except Exception as e:
-        logging.error("Error creating task: %s", e)
-        logging.error("response: %s", response)
-    logging.info("Successfully created task in create_http_task_with_name")
+        if "Requested entity already exists" in str(e):
+            logging.info("Task already exists for %s, skipping.", task_name)
+        else:
+            logging.error("Error creating task: %s", e)
+            logging.error("response: %s", response)

functions-python/batch_datasets/src/main.py

@@ -29,7 +29,8 @@
 from sqlalchemy.orm import Session
 
 from shared.database_gen.sqlacodegen_models import Gtfsfeed, Gtfsdataset
-from shared.dataset_service.main import BatchExecutionService, BatchExecution
+from shared.dataset_service.dataset_service_commons import BatchExecution
+from shared.dataset_service.main import BatchExecutionService
 from shared.database.database import with_db_session
 from shared.helpers.logger import init_logger
 

functions-python/dataset_service/dataset_service_commons.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+from datetime import datetime
+from enum import Enum
+from typing import Optional
+
+
+# Status of the dataset trace
+class Status(Enum):
+    FAILED = "FAILED"
+    SUCCESS = "SUCCESS"
+    PUBLISHED = "PUBLISHED"
+    NOT_PUBLISHED = "NOT_PUBLISHED"
+    PROCESSING = "PROCESSING"
+
+
+# Stage of the pipeline
+class PipelineStage(Enum):
+    DATASET_PROCESSING = "DATASET_PROCESSING"
+    LOCATION_EXTRACTION = "LOCATION_EXTRACTION"
+    GBFS_VALIDATION = "GBFS_VALIDATION"
+
+
+# Dataset trace class to store the trace of a dataset
+@dataclass
+class DatasetTrace:
+    stable_id: str
+    status: Status
+    timestamp: datetime
+    dataset_id: Optional[str] = None
+    trace_id: Optional[str] = None
+    execution_id: Optional[str] = None
+    file_sha256_hash: Optional[str] = None
+    hosted_url: Optional[str] = None
+    pipeline_stage: PipelineStage = PipelineStage.DATASET_PROCESSING
+    error_message: Optional[str] = None
+
+
+# Batch execution class to store the trace of a batch execution
+@dataclass
+class BatchExecution:
+    execution_id: str
+    timestamp: datetime
+    feeds_total: int

functions-python/dataset_service/main.py

@@ -13,15 +13,30 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
+import importlib
 import logging
 import uuid
-from datetime import datetime
-from enum import Enum
-from dataclasses import dataclass, asdict
-from typing import Optional, Final
+from dataclasses import asdict
+from typing import Final
 from google.cloud import datastore
 from google.cloud.datastore import Client
 
+# This allows the module to be run as a script or imported as a module
+if __package__ is None or __package__ == "":
+    import os
+    import sys
+
+    sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+    import dataset_service_commons
+else:
+    dataset_service_commons = importlib.import_module(
+        ".dataset_service_commons", package=__package__
+    )
+
+Status = dataset_service_commons.Status
+PipelineStage = dataset_service_commons.PipelineStage
+BatchExecution = dataset_service_commons.BatchExecution
+DatasetTrace = dataset_service_commons.DatasetTrace
 
 # This files contains the dataset trace and batch execution models and services.
 # The dataset trace is used to store the trace of a dataset and the batch execution
@@ -30,45 +45,6 @@
 # The persistent layer used is Google Cloud Datastore.
 
 
-# Status of the dataset trace
-class Status(Enum):
-    FAILED = "FAILED"
-    SUCCESS = "SUCCESS"
-    PUBLISHED = "PUBLISHED"
-    NOT_PUBLISHED = "NOT_PUBLISHED"
-    PROCESSING = "PROCESSING"
-
-
-# Stage of the pipeline
-class PipelineStage(Enum):
-    DATASET_PROCESSING = "DATASET_PROCESSING"
-    LOCATION_EXTRACTION = "LOCATION_EXTRACTION"
-    GBFS_VALIDATION = "GBFS_VALIDATION"
-
-
-# Dataset trace class to store the trace of a dataset
-@dataclass
-class DatasetTrace:
-    stable_id: str
-    status: Status
-    timestamp: datetime
-    dataset_id: Optional[str] = None
-    trace_id: Optional[str] = None
-    execution_id: Optional[str] = None
-    file_sha256_hash: Optional[str] = None
-    hosted_url: Optional[str] = None
-    pipeline_stage: PipelineStage = PipelineStage.DATASET_PROCESSING
-    error_message: Optional[str] = None
-
-
-# Batch execution class to store the trace of a batch execution
-@dataclass
-class BatchExecution:
-    execution_id: str
-    timestamp: datetime
-    feeds_total: int
-
-
 dataset_trace_collection: Final[str] = "dataset_trace"
 batch_execution_collection: Final[str] = "batch_execution"
 

functions-python/dataset_service/tests/__init__.py

@@ -2,13 +2,8 @@
 from datetime import datetime
 from unittest.mock import patch, MagicMock
 
-from main import (
-    DatasetTrace,
-    DatasetTraceService,
-    Status,
-    BatchExecutionService,
-    BatchExecution,
-)
+from dataset_service_commons import DatasetTrace, Status, BatchExecution
+from main import DatasetTraceService, BatchExecutionService
 
 
 class TestDatasetService(unittest.TestCase):

functions-python/dataset_service/tests/test_dataset_service.py

@@ -4,6 +4,7 @@ omit =
     database.py
     */database_gen/*
     */dataset_service/*
+    */shared/common/*
 
 [report]
 exclude_lines =

functions-python/helpers/.coveragerc

@@ -1,8 +1,13 @@
 from enum import Enum
 from typing import Dict, Optional
+
+from geoalchemy2 import WKTElement
 from sqlalchemy.orm import Session
+from sqlalchemy import func, cast
+from geoalchemy2.types import Geography
+
 import pycountry
-from shared.database_gen.sqlacodegen_models import Feed, Location
+from shared.database_gen.sqlacodegen_models import Feed, Location, Geopolygon
 import logging
 
 
@@ -11,8 +16,14 @@ class ReverseGeocodingStrategy(str, Enum):
     Enum for reverse geocoding strategies.
     """
 
+    # Per point strategy uses point-in-polygon to find the location for each point
+    # It queries the database for each point, which can be slow for large datasets
     PER_POINT = "per-point"
 
+    # Per polygon strategy uses point-in-polygon to find the location for each point
+    # It queries the database for each polygon, which can be faster for large datasets
+    PER_POLYGON = "per-polygon"
+
 
 def get_country_code(country_name: str) -> Optional[str]:
     """
@@ -133,3 +144,90 @@ def translate_feed_locations(feed: Feed, location_translations: Dict):
                 if location_translation["country_translation"]
                 else location.country
             )
+
+
+def to_shapely(g):
+    """
+    Convert a GeoAlchemy WKB/WKT element or WKT string into a Shapely geometry.
+    If it's already a Shapely geometry, return it as-is.
+    """
+    # Import here to avoid adding unnecessary dependencies if not used to GCP functions
+    from shapely import wkt as shapely_wkt
+    from geoalchemy2 import WKTElement, WKBElement
+    from geoalchemy2.shape import to_shape
+
+    if isinstance(g, WKBElement):
+        return to_shape(g)
+    if isinstance(g, WKTElement):
+        return shapely_wkt.loads(g.data)
+    if isinstance(g, str):
+        # assume WKT
+        return shapely_wkt.loads(g)
+    return g  # assume already shapely
+
+
+def select_highest_level_polygon(geopolygons: list[Geopolygon]) -> Optional[Geopolygon]:
+    """
+    Select the geopolygon with the highest admin_level from a list of geopolygons.
+    Admin levels are compared, with NULL treated as the lowest priority.
+    """
+    if not geopolygons:
+        return None
+    # Treat NULL admin_level as the lowest priority
+    return max(
+        geopolygons, key=lambda g: (-1 if g.admin_level is None else g.admin_level)
+    )
+
+
+def select_lowest_level_polygon(geopolygons: list[Geopolygon]) -> Optional[Geopolygon]:
+    """
+    Select the geopolygon with the lowest admin_level from a list of geopolygons.
+    Admin levels are compared, with NULL treated as the lowest priority.
+    """
+    if not geopolygons:
+        return None
+    # Treat NULL admin_level as the lowest priority
+    return min(
+        geopolygons, key=lambda g: (100 if g.admin_level is None else g.admin_level)
+    )
+
+
+def get_country_code_from_polygons(geopolygons: list[Geopolygon]) -> Optional[str]:
+    """
+    Given a list of polygon GeoJSON-like features (each with 'properties'),
+    return the country code (ISO 3166-1 alpha-2) from the most likely polygon.
+
+    Args:
+        polygons: List of dicts, each must have 'properties' with
+                  'admin_level' and 'iso_3166_1_code'
+
+    Returns:
+        A two-letter country code string or None if not found
+    """
+    country_polygons = [g for g in geopolygons if g.iso_3166_1_code]
+    if not country_polygons:
+        return None
+
+    # Prefer the one with the lowest admin_level (most local)
+    lowest_admin_level_polygon = select_lowest_level_polygon(country_polygons)
+    return lowest_admin_level_polygon.iso_3166_1_code
+
+
+def get_geopolygons_covers(stop_point: WKTElement, db_session: Session):
+    """
+    Get all geopolygons that cover a given point using BigQuery-compatible semantics.
+    """
+    # BigQuery-compatible point-in-polygon (geodesic + border-inclusive)
+    geopolygons = (
+        db_session.query(Geopolygon)
+        # optional prefilter to use your GiST index on geometry (fast)
+        .filter(func.ST_Intersects(Geopolygon.geometry, stop_point))
+        # exact check matching BigQuery's GEOGRAPHY semantics
+        .filter(
+            func.ST_Covers(
+                cast(Geopolygon.geometry, Geography(srid=4326)),
+                cast(stop_point, Geography(srid=4326)),
+            )
+        ).all()
+    )
+    return geopolygons

functions-python/helpers/locations.py

@@ -17,8 +17,6 @@
 import logging
 import threading
 
-import google.cloud.logging
-
 from shared.common.logging_utils import get_env_logging_level
 
 
@@ -61,6 +59,8 @@ def init_logger():
             if _logging_initialized:
                 return
             try:
+                import google.cloud.logging
+
                 client = google.cloud.logging.Client()
                 client.setup_logging()
             except Exception as error:

functions-python/helpers/logger.py

@@ -27,4 +27,5 @@ google-cloud-firestore
 google-cloud-bigquery
 
 # Additional package
-pycountry
+pycountry
+shapely