diff --git a/api/src/shared/common/gcp_utils.py b/api/src/shared/common/gcp_utils.py index 52d26fe79..c7922263e 100644 --- a/api/src/shared/common/gcp_utils.py +++ b/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) diff --git a/functions-python/batch_datasets/src/main.py b/functions-python/batch_datasets/src/main.py index 7798f5eac..137efc632 100644 --- a/functions-python/batch_datasets/src/main.py +++ b/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 diff --git a/functions-python/dataset_service/dataset_service_commons.py b/functions-python/dataset_service/dataset_service_commons.py new file mode 100644 index 000000000..dea026d6e --- /dev/null +++ b/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 diff --git a/functions-python/dataset_service/main.py b/functions-python/dataset_service/main.py index e4429a612..a397ffa9e 100644 --- a/functions-python/dataset_service/main.py +++ b/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" diff --git a/functions-python/dataset_service/tests/__init__.py b/functions-python/dataset_service/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/functions-python/dataset_service/tests/test_dataset_service.py b/functions-python/dataset_service/tests/test_dataset_service.py index d87e1a7e9..c4ce9ad3e 100644 --- a/functions-python/dataset_service/tests/test_dataset_service.py +++ b/functions-python/dataset_service/tests/test_dataset_service.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): diff --git a/functions-python/helpers/.coveragerc b/functions-python/helpers/.coveragerc index b941555d1..705e38d1d 100644 --- a/functions-python/helpers/.coveragerc +++ b/functions-python/helpers/.coveragerc @@ -4,6 +4,7 @@ omit = database.py */database_gen/* */dataset_service/* + */shared/common/* [report] exclude_lines = diff --git a/functions-python/helpers/locations.py b/functions-python/helpers/locations.py index 14e60bb98..7d0274abd 100644 --- a/functions-python/helpers/locations.py +++ b/functions-python/helpers/locations.py @@ -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 diff --git a/functions-python/helpers/logger.py b/functions-python/helpers/logger.py index 3cb9385d7..dbbec6611 100644 --- a/functions-python/helpers/logger.py +++ b/functions-python/helpers/logger.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: diff --git a/functions-python/helpers/requirements.txt b/functions-python/helpers/requirements.txt index 283688f78..76e892a28 100644 --- a/functions-python/helpers/requirements.txt +++ b/functions-python/helpers/requirements.txt @@ -27,4 +27,5 @@ google-cloud-firestore google-cloud-bigquery # Additional package -pycountry \ No newline at end of file +pycountry +shapely \ No newline at end of file diff --git a/functions-python/helpers/tests/test_locations.py b/functions-python/helpers/tests/test_locations.py index 5ef805ac8..b1ade8e6b 100644 --- a/functions-python/helpers/tests/test_locations.py +++ b/functions-python/helpers/tests/test_locations.py @@ -2,14 +2,25 @@ import unittest from unittest.mock import MagicMock -from shared.database_gen.sqlacodegen_models import Feed, Location + +from geoalchemy2 import WKTElement +from geoalchemy2.shape import from_shape +from shapely.geometry.point import Point +from shapely.geometry.polygon import Polygon + from locations import ( translate_feed_locations, get_country_code, create_or_get_location, + to_shapely, + select_highest_level_polygon, + select_lowest_level_polygon, + get_country_code_from_polygons, ) from unittest.mock import patch +from shared.database_gen.sqlacodegen_models import Location, Feed, Geopolygon + class TestLocations(unittest.TestCase): """Test cases for location-related functionality.""" @@ -277,3 +288,170 @@ def test_translate_feed_locations_multiple_locations(self): self.assertEqual(mock_location2.subdivision_name, "Translated State 2") self.assertEqual(mock_location2.municipality, "Translated City 2") self.assertEqual(mock_location2.country, "Translated Country 2") + + def test_to_shapely_wkt_element(self): + wkt_element = WKTElement("POINT (1 2)", srid=4326) + result = to_shapely(wkt_element) + self.assertIsInstance(result, Point) + self.assertEqual(result.x, 1) + self.assertEqual(result.y, 2) + + def test_to_shapely_wkb_element(self): + shapely_point = Point(1, 2) + wkb_element = from_shape(shapely_point, srid=4326) + result = to_shapely(wkb_element) + self.assertIsInstance(result, Point) + self.assertEqual(result.x, 1) + self.assertEqual(result.y, 2) + + def test_to_shapely_wkt_string(self): + wkt_string = "POINT (1 2)" + result = to_shapely(wkt_string) + self.assertIsInstance(result, Point) + self.assertEqual(result.x, 1) + self.assertEqual(result.y, 2) + + def test_to_shapely_shapely_geometry(self): + shapely_polygon = Polygon([(0, 0), (1, 0), (1, 1), (0, 1), (0, 0)]) + result = to_shapely(shapely_polygon) + self.assertIs(result, shapely_polygon) + + def test_to_shapely_invalid_input(self): + invalid_input = 12345 + result = to_shapely(invalid_input) + self.assertEqual(result, invalid_input) + + def test_select_highest_level_polygon_mpty_list_returns_none(self): + result = select_highest_level_polygon([]) + assert result is None + + def test_select_highest_level_polygon_single_polygon(self): + g1 = Geopolygon(osm_id=1, admin_level=5) + result = select_highest_level_polygon([g1]) + assert result == g1 + + def test_select_highest_level_polygon_multiple_polygons_selects_highest(self): + g1 = Geopolygon(osm_id=1, admin_level=3) + g2 = Geopolygon(osm_id=2, admin_level=7) + g3 = Geopolygon(osm_id=3, admin_level=5) + result = select_highest_level_polygon([g1, g2, g3]) + assert result == g2 + + def test_select_highest_level_polygon_null_admin_level_treated_lowest(self): + g1 = Geopolygon(osm_id=1, admin_level=None) + g2 = Geopolygon(osm_id=2, admin_level=4) + result = select_highest_level_polygon([g1, g2]) + assert result == g2 + + def test_select_highest_level_polygon_all_null_admin_levels(self): + g1 = Geopolygon(osm_id=1, admin_level=None) + g2 = Geopolygon(osm_id=2, admin_level=None) + result = select_highest_level_polygon([g1, g2]) + # Should return one of them, but not None + assert result in [g1, g2] + + def test_select_highest_level_polygon_ties_highest_admin_level(self): + g1 = Geopolygon(osm_id=1, admin_level=6) + g2 = Geopolygon(osm_id=2, admin_level=6) + result = select_highest_level_polygon([g1, g2]) + # Either polygon with admin_level=6 is valid + assert result.admin_level == 6 + + def test_select_lowest_level_polygon_empty_list_returns_none(self): + self.assertIsNone(select_lowest_level_polygon([])) + + def test_select_lowest_level_polygon_single_polygon_is_returned(self): + g1 = Geopolygon(osm_id=1, admin_level=5) + self.assertEqual(g1, select_lowest_level_polygon([g1])) + + def test_select_lowest_level_polygon_chooses_smallest_numeric_level(self): + g1 = Geopolygon(osm_id=1, admin_level=7) + g2 = Geopolygon(osm_id=2, admin_level=3) + g3 = Geopolygon(osm_id=3, admin_level=5) + result = select_lowest_level_polygon([g1, g2, g3]) + self.assertEqual(g2, result) + self.assertEqual(3, result.admin_level) + + def test_select_lowest_level_polygon_ignores_none_when_numbers_exist(self): + g1 = Geopolygon(osm_id=1, admin_level=None) + g2 = Geopolygon(osm_id=2, admin_level=4) + result = select_lowest_level_polygon([g1, g2]) + self.assertEqual(g2, result) + self.assertEqual(4, result.admin_level) + + def test_select_lowest_level_polygon_all_none_returns_one_of_inputs(self): + g1 = Geopolygon(osm_id=1, admin_level=None) + g2 = Geopolygon(osm_id=2, admin_level=None) + result = select_lowest_level_polygon([g1, g2]) + self.assertIn(result, (g1, g2)) + self.assertIsNone(result.admin_level) + + def test_select_lowest_level_polygon_ties_return_one_with_that_level(self): + g1 = Geopolygon(osm_id=1, admin_level=2) + g2 = Geopolygon(osm_id=2, admin_level=2) + result = select_lowest_level_polygon([g1, g2]) + self.assertEqual(2, result.admin_level) + + def test_get_country_code_from_polygons_returns_none_for_empty_list(self): + self.assertIsNone(get_country_code_from_polygons([])) + + def test_get_country_code_from_polygons_ignores_polygons_without_country_code(self): + # Only one polygon has an ISO code -> it should be chosen regardless of admin_level + polys = [ + Geopolygon(osm_id=1, admin_level=5, iso_3166_1_code=None), + Geopolygon(osm_id=2, admin_level=3, iso_3166_1_code=""), # falsy -> ignored + Geopolygon(osm_id=3, admin_level=7, iso_3166_1_code="CA"), + ] + self.assertEqual("CA", get_country_code_from_polygons(polys)) + + def test_get_country_code_from_polygons_returns_none_when_no_iso_codes_present( + self, + ): + polys = [ + Geopolygon(osm_id=1, admin_level=3, iso_3166_1_code=None), + Geopolygon(osm_id=2, admin_level=2, iso_3166_1_code=""), + ] + self.assertIsNone(get_country_code_from_polygons(polys)) + + def test_get_country_code_from_polygons_picks_lowest_admin_level(self): + # Among those with ISO codes, choose the one with the smallest admin_level + polys = [ + Geopolygon(osm_id=1, admin_level=7, iso_3166_1_code="US"), + Geopolygon(osm_id=2, admin_level=3, iso_3166_1_code="CA"), + Geopolygon(osm_id=3, admin_level=5, iso_3166_1_code="MX"), + ] + self.assertEqual( + "CA", get_country_code_from_polygons(polys) + ) # admin_level=3 is lowest + + def test_get_country_code_from_polygons_tie_returns_any_with_that_level(self): + # If two have the same lowest admin_level, either is fine. + polys = [ + Geopolygon(osm_id=1, admin_level=2, iso_3166_1_code="US"), + Geopolygon(osm_id=2, admin_level=2, iso_3166_1_code="CA"), + Geopolygon(osm_id=3, admin_level=4, iso_3166_1_code="MX"), + ] + result = get_country_code_from_polygons(polys) + self.assertIn(result, {"US", "CA"}) + + def test_get_country_code_from_polygons_none_admin_levels_are_low_priority_when_numbers_exist( + self, + ): + # If select_lowest_level_polygon treats None as "lowest priority", + # polygons with numeric admin_level should win over None. + polys = [ + Geopolygon(osm_id=1, admin_level=None, iso_3166_1_code="US"), + Geopolygon(osm_id=2, admin_level=4, iso_3166_1_code="CA"), + ] + self.assertEqual("CA", get_country_code_from_polygons(polys)) + + def test_get_country_code_from_polygons_all_none_admin_levels_returns_one_with_code( + self, + ): + # When all eligible have admin_level=None, any with an ISO code is acceptable. + polys = [ + Geopolygon(osm_id=1, admin_level=None, iso_3166_1_code="US"), + Geopolygon(osm_id=2, admin_level=None, iso_3166_1_code="CA"), + ] + result = get_country_code_from_polygons(polys) + self.assertIn(result, {"US", "CA"}) diff --git a/functions-python/helpers/transform.py b/functions-python/helpers/transform.py index f0df6537e..284bf9782 100644 --- a/functions-python/helpers/transform.py +++ b/functions-python/helpers/transform.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import logging from typing import List, Optional @@ -73,5 +74,6 @@ def to_enum(value, enum_class=None, default_value=None): return value try: return enum_class(str(value)) - except (ValueError, TypeError): + except (ValueError, TypeError) as e: + logging.warning("Failed to convert value to enum member: %s", e) return default_value diff --git a/functions-python/helpers/utils.py b/functions-python/helpers/utils.py index 0dffa43fd..b2c13dcbf 100644 --- a/functions-python/helpers/utils.py +++ b/functions-python/helpers/utils.py @@ -17,10 +17,12 @@ import logging import os import ssl +from datetime import date, datetime +from logging import Logger +from typing import Optional import requests import urllib3 -from google.cloud import storage from requests.adapters import HTTPAdapter from urllib3.util.retry import Retry from urllib3.util.ssl_ import create_urllib3_context @@ -32,6 +34,8 @@ def create_bucket(bucket_name): Creates GCP storage bucket if it doesn't exist :param bucket_name: name of the bucket to create """ + from google.cloud import storage + storage_client = storage.Client() bucket = storage_client.lookup_bucket(bucket_name) if bucket is None: @@ -53,6 +57,8 @@ def download_from_gcs(bucket_name: str, blob_path: str, local_path: str) -> str: Returns: The absolute path to the downloaded file. """ + from google.cloud import storage + storage_client = storage.Client() bucket = storage_client.bucket(bucket_name) blob = bucket.blob(blob_path) @@ -211,3 +217,85 @@ def create_http_task( task_time=proto_time, http_method=tasks_v2.HttpMethod.POST, ) + + +def get_execution_id(json_payload: dict, stable_id: Optional[str]) -> str: + """ + Extracts the execution_id from the JSON payload. + If not present, defaults to today's date in YYYY-MM-DD format followed by a hyphen and the stable_id if provided. + """ + execution_id = json_payload.get("execution_id") + if not execution_id: + execution_id = f"{str(date.today())}" + if stable_id: + execution_id += f"-{stable_id}" + else: + # Even this should not happen, but just in case we are defaulting it to the current time + execution_id += f"-{datetime.now().strftime('%H:%M:%S')}" + return execution_id + + +def check_maximum_executions( + execution_id: str, stable_id: str, logger: Logger, maximum_executions: int = 1 +) -> str: + """ + Checks if the dataset has been executed more than the maximum allowed times. + If it has, returns an error message; otherwise, returns None. + :param execution_id: The ID of the execution. + :param stable_id: The stable ID of the dataset. + :param logger: Logger instance to log messages. + :param maximum_executions: The maximum number of allowed executions. + :return: Error message if the maximum executions are exceeded, otherwise None. + """ + from shared.dataset_service.main import DatasetTraceService + + trace_service = DatasetTraceService() + trace = trace_service.get_by_execution_and_stable_ids(execution_id, stable_id) + executions = len(trace) if trace else 0 + logger.info( + f"Function executed times={executions}/{maximum_executions} " + f"in execution=[{execution_id}] " + ) + + if executions > 0: + if executions >= maximum_executions: + message = ( + f"Function already executed maximum times " + f"in execution: [{execution_id}]" + ) + logger.warning(message) + return message + return None + + +def record_execution_trace( + execution_id, + stable_id, + status, + logger=None, + dataset_file=None, + error_message=None, +): + """ + Record the trace in the datastore + """ + from shared.dataset_service.main import DatasetTraceService + from shared.dataset_service.dataset_service_commons import DatasetTrace + from shared.helpers.logger import get_logger + + trace_service = DatasetTraceService() + + (logger if logger else get_logger()).info( + f"Recording trace in execution: [{execution_id}] with status: [{status}]" + ) + trace = DatasetTrace( + trace_id=None, + stable_id=stable_id, + status=status, + execution_id=execution_id, + file_sha256_hash=dataset_file.file_sha256_hash if dataset_file else None, + hosted_url=dataset_file.hosted_url if dataset_file else None, + error_message=error_message, + timestamp=datetime.now(), + ) + trace_service.save(trace) diff --git a/functions-python/reverse_geolocation/function_config.json b/functions-python/reverse_geolocation/function_config.json index c071a7d6f..c649d4092 100644 --- a/functions-python/reverse_geolocation/function_config.json +++ b/functions-python/reverse_geolocation/function_config.json @@ -5,7 +5,7 @@ "timeout": 540, "available_memory": "4Gi", "trigger_http": true, - "include_folders": ["helpers"], + "include_folders": ["helpers", "dataset_service"], "include_api_folders": ["database_gen", "database", "common"], "environment_variables": [], "secret_environment_variables": [ diff --git a/functions-python/reverse_geolocation/src/location_group_utils.py b/functions-python/reverse_geolocation/src/location_group_utils.py index 83b6e413b..2e827a652 100644 --- a/functions-python/reverse_geolocation/src/location_group_utils.py +++ b/functions-python/reverse_geolocation/src/location_group_utils.py @@ -1,12 +1,26 @@ +from logging import Logger from typing import List, Optional import matplotlib.pyplot as plt import pycountry +from geoalchemy2 import WKTElement from geoalchemy2.shape import to_shape +from pyproj import Geod +from shapely.geometry.geo import mapping +from sqlalchemy.orm import Session -from shared.database_gen.sqlacodegen_models import Geopolygon, Osmlocationgroup +from shared.database_gen.sqlacodegen_models import ( + Geopolygon, + Osmlocationgroup, + Feedosmlocationgroup, + Location, + Feed, + Feedlocationgrouppoint, +) +from shared.helpers.locations import get_geopolygons_covers ERROR_STATUS_CODE = 299 # Custom error code for the function to avoid retries +GEOD = Geod(ellps="WGS84") # Geod object for geodesic calculations def generate_color( @@ -29,6 +43,7 @@ class GeopolygonAggregate: """ def __init__(self, location_group: Osmlocationgroup, stops_count: int): + self.location_group = location_group self.group_id = location_group.group_id self.group_name = location_group.group_name self.geopolygons = [ @@ -130,3 +145,208 @@ def geopolygons_as_string(geopolygons: List[Geopolygon]) -> str: return ", ".join( [str(geopolygon) for geopolygon in detach_from_session(geopolygons)] ) + + +def extract_location_aggregate( + stop_point: WKTElement, logger: Logger, db_session: Session +) -> Optional[GeopolygonAggregate]: + """ + Extract the location group for a given stop point. + """ + geopolygons = get_geopolygons_covers(stop_point, db_session) + + if len(geopolygons) <= 1: + logger.warning( + "Invalid number of geopolygons for point: %s -> %s", stop_point, geopolygons + ) + return None + return extract_location_aggregate_geopolygons( + stop_point=stop_point, + geopolygons=geopolygons, + logger=logger, + db_session=db_session, + ) + + +def geodesic_area_m2(geom) -> float: + """Return absolute geodesic area in m² for a Shapely geometry (Polygon/MultiPolygon).""" + # pyproj can handle GeoJSON-like mappings, including MultiPolygons + area, _perim = GEOD.geometry_area_perimeter(mapping(geom)) + return abs(area) + + +def resolve_same_level_conflicts(items: List[Geopolygon]) -> Geopolygon: + # 1) Prefer iso_3166_2_code if present + candidates = [g for g in items if g.iso_3166_2_code] or items + + if len(candidates) > 1: + # 2) Prefer the smallest geodesic area (m²) + def area_for(g: Geopolygon) -> float: + try: + geom = to_shape(g.geometry) # -> Shapely geometry (in EPSG:4326) + return geodesic_area_m2(geom) + except Exception: + # If geometry missing/bad, push to the end + return float("inf") + + candidates.sort(key=area_for) + + # 3) Prefer with name, then lowest OSM id + candidates.sort(key=lambda g: (0 if (g.name and g.name.strip()) else 1, g.osm_id)) + return candidates[0] + + +def dedupe_by_admin_level(geopolygons: List[Geopolygon], logger) -> List[Geopolygon]: + """Keep one polygon per admin_level using tie-break rules above.""" + by_level: dict[int, List[Geopolygon]] = {} + for g in geopolygons: + by_level.setdefault(g.admin_level, []).append(g) + + winners: List[Geopolygon] = [] + for level, items in sorted( + by_level.items(), key=lambda kv: kv[0] + ): # keep order by level + if len(items) > 1: + logger.debug( + "Tie at admin_level=%s among osm_ids=%s", + level, + [i.osm_id for i in items], + ) + winners.append(resolve_same_level_conflicts(items)) + return winners + + +def extract_location_aggregate_geopolygons( + stop_point: WKTElement, geopolygons, logger, db_session: Session +) -> Optional[GeopolygonAggregate]: + admin_levels = {g.admin_level for g in geopolygons} + # If duplicates per admin_level exist, resolve instead of returning None + if len(admin_levels) != len(geopolygons): + logger.warning( + "Duplicate admin levels for point: %s -> %s", + stop_point, + geopolygons_as_string(geopolygons), + ) + geopolygons = dedupe_by_admin_level(geopolygons, logger) + logger.warning( + "Deduplicated admin levels for point: %s -> %s", + stop_point, + geopolygons_as_string(geopolygons), + ) + + valid_iso_3166_1 = any(g.iso_3166_1_code for g in geopolygons) + valid_iso_3166_2 = any(g.iso_3166_2_code for g in geopolygons) + if not valid_iso_3166_1 or not valid_iso_3166_2: + logger.warning( + "Invalid ISO codes for point: %s -> %s", + stop_point, + geopolygons_as_string(geopolygons), + ) + return + + # Sort the polygons by admin level so that lower levels come first + geopolygons.sort(key=lambda x: x.admin_level) + + group_id = ".".join([str(g.osm_id) for g in geopolygons]) + group = ( + db_session.query(Osmlocationgroup) + .filter(Osmlocationgroup.group_id == group_id) + .one_or_none() + ) + if not group: + group = Osmlocationgroup( + group_id=group_id, + group_name=", ".join([g.name for g in geopolygons]), + osms=geopolygons, + ) + db_session.add(group) + db_session.flush() + logger.debug( + "Point %s matched to %s", stop_point, ", ".join([g.name for g in geopolygons]) + ) + return GeopolygonAggregate(group, 1) + + +def create_or_update_stop_group( + feed: Feed, + stop_point: WKTElement, + group: Osmlocationgroup, + logger: Logger, + db_session: Session, +): + """ + Create or update the stop group for a given stop point. + This function ensures that each stop point is associated with the correct location group. + At the end of the function, the group and stop entities are flushed to ensure they are in sync with the database. + """ + stop = ( + db_session.query(Feedlocationgrouppoint) + .filter( + Feedlocationgrouppoint.feed_id == feed.id, + Feedlocationgrouppoint.geometry == stop_point, + ) + .one_or_none() + ) + if not stop: + stop = Feedlocationgrouppoint( + feed_id=feed.id, + geometry=stop_point, + ) + stop.group = group + db_session.add(stop) + else: + if stop.group_id != group.group_id: + logger.info( + "Updating stop point %s from group %s to %s", + stop_point, + stop.group_id, + group.group_id, + ) + stop.group = group + db_session.flush() # Ensure the group and stop entity is in sync with the DB + + +def get_or_create_feed_osm_location_group( + feed_id: str, location_aggregate: GeopolygonAggregate, db_session: Session +) -> Feedosmlocationgroup: + """Get or create the feed osm location group.""" + feed_osm_location = ( + db_session.query(Feedosmlocationgroup) + .filter( + Feedosmlocationgroup.feed_id == feed_id, + Feedosmlocationgroup.group_id == location_aggregate.group_id, + ) + .one_or_none() + ) + if not feed_osm_location: + feed_osm_location = Feedosmlocationgroup( + feed_id=feed_id, + group_id=location_aggregate.group_id, + ) + feed_osm_location.stops_count = location_aggregate.stop_count + return feed_osm_location + + +def get_or_create_location( + location_group: GeopolygonAggregate, logger, db_session: Session +) -> Optional[Location]: + """Get or create the Location entity.""" + try: + logger.debug("Location ID : %s", location_group.location_id()) + location = ( + db_session.query(Location) + .filter(Location.id == location_group.location_id()) + .one_or_none() + ) + if not location: + location = Location( + id=location_group.location_id(), + country_code=location_group.iso_3166_1_code, + country=location_group.country(), + subdivision_name=location_group.subdivision_name(), + municipality=location_group.municipality(), + ) + return location + except Exception as e: + logger.error("Error creating location: %s", e) + return None diff --git a/functions-python/reverse_geolocation/src/parse_request.py b/functions-python/reverse_geolocation/src/parse_request.py index e47e7f33a..e74d4d08f 100644 --- a/functions-python/reverse_geolocation/src/parse_request.py +++ b/functions-python/reverse_geolocation/src/parse_request.py @@ -8,16 +8,27 @@ from jsonpath_ng import parse from shared.helpers.locations import ReverseGeocodingStrategy +from shared.helpers.runtime_metrics import track_metrics from shared.helpers.transform import to_boolean, to_enum +@track_metrics(metrics=("time", "memory", "cpu")) def parse_request_parameters( request: flask.Request, ) -> Tuple[pd.DataFrame, str, Optional[str], str, List[str]]: """ Parse the request parameters and return a DataFrame with the stops data. - @:returns Tuple: A tuple containing the stops DataFrame, stable ID, dataset ID, data type, and a list of URLs that - were used to fetch the data. + @:returns Tuple: A tuple containing: + - df: DataFrame + - feed_stable_id: str + - dataset_id: str (only for GTFS) + - data_type: str, either 'gtfs' or 'gbfs'. + - urls: List, a list of URLs that were used to fetch the data. + - public: bool(Optional), whether the data should be public or not. Default is True. + - strategy: ReverseGeocodingStrategy, the strategy to use for reverse geocoding. Default is PER_POINT. + - use_cache: bool(Optional), whether to use cache or not. Default is True for GBFS, false otherwise. + @:raises ValueError: If the request mandatory parameters are invalid or missing. + """ logging.info("Parsing request parameters.") request_json = request.get_json(silent=True) @@ -51,16 +62,39 @@ def parse_request_parameters( public = True if "public" in request_json: public = to_boolean(request_json["public"], default_value=True) - strategy = ReverseGeocodingStrategy.PER_POINT + strategy = ReverseGeocodingStrategy.PER_POLYGON if "strategy" in request_json: strategy = to_enum( + enum_class=ReverseGeocodingStrategy, value=request_json["strategy"], - default_value=ReverseGeocodingStrategy.PER_POINT, + default_value=ReverseGeocodingStrategy.PER_POLYGON, ) else: logging.info("No strategy provided, using default") logging.info("Strategy set to: %s.", strategy) - return df, stable_id, dataset_id, data_type, urls, public, strategy + if "use_cache" in request_json: + use_cache = to_boolean( + request_json["use_cache"], default_value=(data_type == "gtfs") + ) + logging.info("Use cache: %s", use_cache) + else: + use_cache = data_type == "gtfs" + logging.info("No use_cache provided, using(%s): %s", data_type, use_cache) + if "maximum_executions" in request_json: + maximum_executions = int(request_json["maximum_executions"]) + else: + maximum_executions = 1 + return ( + df, + stable_id, + dataset_id, + data_type, + urls, + public, + strategy, + use_cache, + maximum_executions, + ) def parse_request_parameters_gtfs( @@ -141,6 +175,7 @@ def parse_free_bike_status_url(free_bike_status_url): return pd.DataFrame(bikes_info) +@track_metrics(metrics=("time", "memory", "cpu")) def parse_request_parameters_gbfs( request_json: dict, ) -> Tuple[pd.DataFrame, str, Optional[str], List[str]]: diff --git a/functions-python/reverse_geolocation/src/reverse_geolocation_processor.py b/functions-python/reverse_geolocation/src/reverse_geolocation_processor.py index 446dfa332..fd7984ed0 100644 --- a/functions-python/reverse_geolocation/src/reverse_geolocation_processor.py +++ b/functions-python/reverse_geolocation/src/reverse_geolocation_processor.py @@ -3,109 +3,102 @@ import os import traceback from datetime import datetime -from typing import Dict, Tuple, Optional, List +from logging import Logger +from typing import Dict, Tuple, List import flask import pandas as pd import shapely.geometry from geoalchemy2 import WKTElement from geoalchemy2.shape import to_shape -from google.cloud import storage + from shapely.geometry import mapping from sqlalchemy import func from sqlalchemy.orm import Session from sqlalchemy.orm import joinedload -from shared.helpers.locations import ReverseGeocodingStrategy from location_group_utils import ( ERROR_STATUS_CODE, GeopolygonAggregate, generate_color, - geopolygons_as_string, + get_or_create_feed_osm_location_group, + get_or_create_location, ) from parse_request import parse_request_parameters from shared.common.gcp_utils import create_refresh_materialized_view_task from shared.database.database import with_db_session - from shared.database_gen.sqlacodegen_models import ( - Geopolygon, Feed, Feedlocationgrouppoint, Osmlocationgroup, - Feedosmlocationgroup, - Location, Gtfsdataset, Gtfsfeed, ) +from shared.dataset_service.dataset_service_commons import Status + +from shared.helpers.locations import ReverseGeocodingStrategy from shared.helpers.logger import get_logger from shared.helpers.runtime_metrics import track_metrics +from shared.helpers.utils import ( + check_maximum_executions, + get_execution_id, + record_execution_trace, +) +from strategy_extraction_per_point import extract_location_aggregates_per_point +from strategy_extraction_per_polygon import extract_location_aggregates_per_polygon @with_db_session -def get_cached_geopolygons( - stable_id: str, stops_df: pd.DataFrame, logger, db_session: Session +def get_geopolygons_with_geometry( + feed: Feed, + stops_df: pd.DataFrame, + use_cache: bool, + logger: Logger, + db_session: Session, ) -> Tuple[str, Dict[str, GeopolygonAggregate], pd.DataFrame]: """ - Get the geopolygons from the database cache. + @:returns a tuple containing: - feed_id: The ID of the feed. - location_groups: A dictionary of location groups with the group ID as the key. - - unmatched_stop_df: DataFrame of unmatched stops for further processing. + - unmatched_stop_df: DataFrame of unmatched stops with geo """ logger.info("Getting cached geopolygons for stable ID.") - - if stops_df.empty: - logger.warning("The provided stops DataFrame is empty.") - raise ValueError("The provided stops DataFrame is empty.") - stops_df["geometry"] = stops_df.apply( lambda x: WKTElement(f"POINT ({x['stop_lon']} {x['stop_lat']})", srid=4326), axis=1, ) stops_df["geometry_str"] = stops_df["geometry"].apply(str) - feed = ( - db_session.query(Feed) - .options(joinedload(Feed.feedlocationgrouppoints)) - .filter(Feed.stable_id == stable_id) - .one_or_none() - ) - if not feed: - logger.warning("No feed found for stable ID.") - raise ValueError(f"No feed found for stable ID {stable_id}.") - - feed_id = feed.id - cached_geometries = { to_shape(stop.geometry).wkt for stop in feed.feedlocationgrouppoints } - matched_stops_df = stops_df[stops_df["geometry_str"].isin(cached_geometries)] - unmatched_stop_df = stops_df[~stops_df["geometry_str"].isin(cached_geometries)] - + matched_stops_df = ( + stops_df[stops_df["geometry_str"].isin(cached_geometries)] + if use_cache + else pd.DataFrame(columns=stops_df.columns) + ) + unmatched_stop_df = ( + stops_df[~stops_df["geometry_str"].isin(cached_geometries)] + if use_cache + else stops_df + ) logger.info( "Matched stops: %s | Unmatched stops: %s", len(matched_stops_df), len(unmatched_stop_df), ) - - df_geometry_set = set(matched_stops_df["geometry_str"].tolist()) - geometries_to_delete = cached_geometries - df_geometry_set - - if geometries_to_delete: - logger.info("Deleting %s outdated cached stops.", len(geometries_to_delete)) - db_session.query(Feedlocationgrouppoint).filter( - Feedlocationgrouppoint.feed_id == feed_id, - func.ST_AsText(Feedlocationgrouppoint.geometry).in_( - list(geometries_to_delete) - ), - ).delete(synchronize_session=False) - db_session.flush() + if use_cache: + df_geometry_set = set(matched_stops_df["geometry_str"].tolist()) + geometries_to_delete = cached_geometries - df_geometry_set + if geometries_to_delete: + clean_stop_cache(db_session, feed, geometries_to_delete, logger) matched_geometries = matched_stops_df["geometry"].tolist() if not matched_geometries: logger.info("No matched geometries found.") - return feed_id, dict(), unmatched_stop_df + return dict(), unmatched_stop_df location_group_counts = ( db_session.query( Osmlocationgroup, @@ -113,7 +106,7 @@ def get_cached_geopolygons( ) .join(Feedlocationgrouppoint, Osmlocationgroup.feedlocationgrouppoints) .filter( - Feedlocationgrouppoint.feed_id == feed_id, + Feedlocationgrouppoint.feed_id == feed.id, Feedlocationgrouppoint.geometry.in_(matched_geometries), ) .group_by(Osmlocationgroup.group_id) @@ -127,102 +120,18 @@ def get_cached_geopolygons( } logger.info("Total location groups retrieved: %s", len(location_groups)) - return feed_id, location_groups, unmatched_stop_df - - -def extract_location_aggregate( - feed_id: str, stop_point: WKTElement, logger, db_session: Session -) -> Optional[GeopolygonAggregate]: - """ - Extract the location group for a given stop point. - """ - geopolygons = ( - db_session.query(Geopolygon) - .filter(Geopolygon.geometry.ST_Contains(stop_point)) - .all() - ) - - if len(geopolygons) <= 1: - logger.warning( - "Invalid number of geopolygons for point: %s -> %s", stop_point, geopolygons - ) - return None - admin_levels = {g.admin_level for g in geopolygons} - if len(admin_levels) != len(geopolygons): - logger.warning( - "Duplicate admin levels for point: %s -> %s", - stop_point, - geopolygons_as_string(geopolygons), - ) - return None - - valid_iso_3166_1 = any(g.iso_3166_1_code for g in geopolygons) - valid_iso_3166_2 = any(g.iso_3166_2_code for g in geopolygons) - if not valid_iso_3166_1 or not valid_iso_3166_2: - logger.warning( - "Invalid ISO codes for point: %s -> %s", - stop_point, - geopolygons_as_string(geopolygons), - ) - return - - # Sort the polygons by admin level so that lower levels come first - geopolygons.sort(key=lambda x: x.admin_level) - - group_id = ".".join([str(g.osm_id) for g in geopolygons]) - group = ( - db_session.query(Osmlocationgroup) - .filter(Osmlocationgroup.group_id == group_id) - .one_or_none() - ) - if not group: - group = Osmlocationgroup( - group_id=group_id, - group_name=", ".join([g.name for g in geopolygons]), - osms=geopolygons, - ) - db_session.add(group) - db_session.flush() # Ensure the group is added before using it - stop = ( - db_session.query(Feedlocationgrouppoint) - .filter( - Feedlocationgrouppoint.feed_id == feed_id, - Feedlocationgrouppoint.geometry == stop_point, - ) - .one_or_none() - ) - if not stop: - stop = Feedlocationgrouppoint( - feed_id=feed_id, - geometry=stop_point, - ) - db_session.add(stop) - stop.group = group - logger.info( - "Point %s matched to %s", stop_point, ", ".join([g.name for g in geopolygons]) - ) - return GeopolygonAggregate(group, 1) + return location_groups, unmatched_stop_df -def get_or_create_feed_osm_location_group( - feed_id: str, location_aggregate: GeopolygonAggregate, db_session: Session -) -> Feedosmlocationgroup: - """Get or create the feed osm location group.""" - feed_osm_location = ( - db_session.query(Feedosmlocationgroup) - .filter( - Feedosmlocationgroup.feed_id == feed_id, - Feedosmlocationgroup.group_id == location_aggregate.group_id, - ) - .one_or_none() - ) - if not feed_osm_location: - feed_osm_location = Feedosmlocationgroup( - feed_id=feed_id, - group_id=location_aggregate.group_id, - ) - feed_osm_location.stops_count = location_aggregate.stop_count - return feed_osm_location +@track_metrics(metrics=("time", "memory", "cpu")) +def clean_stop_cache(db_session, feed, geometries_to_delete, logger): + """Clean the stop cache by deleting outdated cached stops.""" + logger.info("Deleting %s outdated cached stops.", len(geometries_to_delete)) + db_session.query(Feedlocationgrouppoint).filter( + Feedlocationgrouppoint.feed_id == feed.id, + func.ST_AsText(Feedlocationgrouppoint.geometry).in_(list(geometries_to_delete)), + ).delete(synchronize_session=False) + db_session.commit() def create_geojson_aggregate( @@ -279,7 +188,7 @@ def create_geojson_aggregate( for osm_id in geo_polygon_count ], } - storage_client = storage.Client() + storage_client = get_storage_client() if data_type == "gtfs": bucket_name = os.getenv("DATASETS_BUCKET_NAME_GTFS") elif data_type == "gbfs": @@ -294,94 +203,10 @@ def create_geojson_aggregate( logger.info("GeoJSON data saved to %s", blob.name) -def get_or_create_location( - location_group: GeopolygonAggregate, logger, db_session: Session -) -> Optional[Location]: - """Get or create the Location entity.""" - try: - logger.info("Location ID : %s", location_group.location_id()) - location = ( - db_session.query(Location) - .filter(Location.id == location_group.location_id()) - .one_or_none() - ) - if not location: - location = Location( - id=location_group.location_id(), - country_code=location_group.iso_3166_1_code, - country=location_group.country(), - subdivision_name=location_group.subdivision_name(), - municipality=location_group.municipality(), - ) - return location - except Exception as e: - logger.error("Error creating location: %s", e) - return None +def get_storage_client(): + from google.cloud import storage - -@with_db_session -def extract_location_aggregates_per_point( - feed_id: str, - stops_df: pd.DataFrame, - location_aggregates: Dict[str, GeopolygonAggregate], - logger: logging.Logger, - db_session: Session, -) -> None: - """Extract the location aggregates for the stops. The location_aggregates dictionary will be updated with the new - location groups, keeping track of the stop count for each aggregate.""" - i = 0 - total_stop_count = len(stops_df) - for _, stop in stops_df.iterrows(): - i += 1 - logger.info("Processing stop %s/%s", i, total_stop_count) - location_aggregate = extract_location_aggregate( - feed_id, stop["geometry"], logger, db_session - ) - if not location_aggregate: - continue - if location_aggregate.group_id in location_aggregates: - location_aggregates[location_aggregate.group_id].merge(location_aggregate) - else: - location_aggregates[location_aggregate.group_id] = location_aggregate - if ( - i % 100 == 0 - ): # Commit every 100 stops to avoid reprocessing all stops in case of failure - db_session.commit() - - feed = db_session.query(Feed).filter(Feed.id == feed_id).one_or_none() - osm_location_groups = [ - get_or_create_feed_osm_location_group( - feed_id, location_aggregates[location_group.group_id], db_session - ) - for location_group in location_aggregates.values() - ] - feed.feedosmlocationgroups.clear() - feed.feedosmlocationgroups.extend(osm_location_groups) - feed_locations = [] - for location_aggregate in location_aggregates.values(): - location = get_or_create_location(location_aggregate, logger, db_session) - if location: - feed_locations.append(location) - - if feed.data_type == "gtfs": - gtfs_feed = db_session.query(Gtfsfeed).filter(Feed.id == feed_id).one_or_none() - for gtfs_rt_feed in gtfs_feed.gtfs_rt_feeds: - logger.info( - "Updating GTFS-RT feed with stable ID %s", gtfs_rt_feed.stable_id - ) - gtfs_rt_feed.feedosmlocationgroups.clear() - gtfs_rt_feed.feedosmlocationgroups.extend(osm_location_groups) - if feed_locations: - gtfs_rt_feed.locations.clear() - gtfs_rt_feed.locations = feed_locations - - if feed_locations: - feed.locations = feed_locations - - # Commit the changes to the database before refreshing the materialized view - db_session.commit() - - create_refresh_materialized_view_task() + return storage.Client() @with_db_session @@ -454,8 +279,30 @@ def reverse_geolocation_process( extraction_urls, public, strategy, + use_cache, + maximum_executions, ) = parse_request_parameters(request) + logger = get_logger(__name__, stable_id) + + # Check for maximum executions to avoid repeated processing during the same day + request_json = request.get_json(silent=True) + execution_id = get_execution_id(request_json, stable_id) + max_execution_error = check_maximum_executions( + execution_id, stable_id, logger, maximum_executions + ) + if max_execution_error: + logger.warning(max_execution_error) + return max_execution_error, ERROR_STATUS_CODE + + record_execution_trace( + execution_id=execution_id, + stable_id=stable_id, + status=Status.PROCESSING, + logger=logger, + dataset_file=None, + error_message=None, + ) # Remove duplicate lat/lon points stops_df["stop_lat"] = pd.to_numeric(stops_df["stop_lat"], errors="coerce") stops_df["stop_lon"] = pd.to_numeric(stops_df["stop_lon"], errors="coerce") @@ -463,26 +310,41 @@ def reverse_geolocation_process( stops_df["stop_lat"].notnull() & stops_df["stop_lon"].notnull() ] stops_df = stops_df.drop_duplicates(subset=["stop_lat", "stop_lon"]) - if stops_df.empty: - logging.warning("All stops have null lat/lon values.") - return "All stops have null lat/lon values", ERROR_STATUS_CODE total_stops = len(stops_df) except ValueError as e: - logging.error(f"Error parsing request parameters: {e}") + logging.error("Error parsing request parameters: %s", e) return str(e), ERROR_STATUS_CODE - logger = get_logger(__name__, stable_id) + if stops_df.empty: + no_stops_message = "No stops found in the feed." + logger.warning(no_stops_message) + return str(no_stops_message), ERROR_STATUS_CODE + try: # Update the bounding box of the dataset bounding_box = update_dataset_bounding_box(dataset_id, stops_df, logger) location_groups = reverse_geolocation( - strategy, - stable_id, - stops_df, - logger, + strategy=strategy, + stable_id=stable_id, + stops_df=stops_df, + logger=logger, + use_cache=use_cache, ) + if not location_groups: + no_locations_message = "No locations found for the provided stops." + logger.warning(no_locations_message) + record_execution_trace( + execution_id=execution_id, + stable_id=stable_id, + status=Status.FAILED, + logger=logger, + dataset_file=None, + error_message=no_locations_message, + ) + return no_locations_message, ERROR_STATUS_CODE + # Create GeoJSON Aggregate create_geojson_aggregate( list(location_groups.values()), @@ -494,6 +356,7 @@ def reverse_geolocation_process( logger=logger, public=public, ) + logger.info( "COMPLETED. Processed %s stops for stable ID %s with strategy. " "Retrieved %s locations.", @@ -501,6 +364,14 @@ def reverse_geolocation_process( stable_id, len(location_groups), ) + record_execution_trace( + execution_id=execution_id, + stable_id=stable_id, + status=Status.SUCCESS, + logger=logger, + dataset_file=None, + error_message=None, + ) return ( f"Processed {total_stops} stops for stable ID {stable_id}. " f"Retrieved {len(location_groups)} locations.", @@ -511,33 +382,122 @@ def reverse_geolocation_process( logger = logger if logger else logging logger.error("Error processing geopolygons: %s", e) logger.error(traceback.format_exc()) # Log full traceback + record_execution_trace( + execution_id=execution_id, + stable_id=stable_id, + status=Status.FAILED, + logger=logger, + dataset_file=None, + error_message=str(e), + ) return str(e), ERROR_STATUS_CODE +@with_db_session @track_metrics(metrics=("time", "memory", "cpu")) def reverse_geolocation( strategy, stable_id, stops_df, logger, + use_cache, + db_session: Session = None, ): """ Reverse geolocation processing based on the specified strategy. """ - logger.info("Processing geopolygons with per-point strategy.") - # Get Cached Geopolygons - feed_id, location_groups, unmatched_stops_df = get_cached_geopolygons( - stable_id, stops_df, logger + logger.info("Processing geopolygons with strategy: %s.", strategy) + + feed = load_feed(stable_id, logger, db_session) + + # Get Geopolygons with Geometry and cached location groups + cache_location_groups, unmatched_stops_df = get_geopolygons_with_geometry( + feed=feed, stops_df=stops_df, use_cache=use_cache, logger=logger ) - logger.info("Number of location groups extracted: %s", len(location_groups)) - # Extract Location Groups - match strategy: - case ReverseGeocodingStrategy.PER_POINT: - extract_location_aggregates_per_point( - feed_id, unmatched_stops_df, location_groups, logger - ) - case _: - logger.error("Invalid strategy: %s", strategy) - return f"Invalid strategy: {strategy}", ERROR_STATUS_CODE + logger.info("Number of location groups cached: %s", len(cache_location_groups)) + if len(unmatched_stops_df) > 0: + # Extract Location Groups + match strategy: + case ReverseGeocodingStrategy.PER_POINT: + extract_location_aggregates_per_point( + feed=feed, + stops_df=unmatched_stops_df, + location_aggregates=cache_location_groups, + use_cache=use_cache, + logger=logger, + ) + case ReverseGeocodingStrategy.PER_POLYGON: + extract_location_aggregates_per_polygon( + feed=feed, + stops_df=unmatched_stops_df, + location_aggregates=cache_location_groups, + use_cache=use_cache, + logger=logger, + ) + case _: + logger.error("Invalid strategy: %s", strategy) + return f"Invalid strategy: {strategy}", ERROR_STATUS_CODE + + update_feed_location( + cache_location_groups=cache_location_groups, + feed=feed, + logger=logger, + db_session=db_session, + ) + create_refresh_materialized_view_task() + return cache_location_groups + + +def load_feed(stable_id, logger, db_session): + feed = ( + db_session.query(Feed) + .options(joinedload(Feed.feedlocationgrouppoints)) + .filter(Feed.stable_id == stable_id) + .one_or_none() + ) + if not feed: + logger.warning("No feed found for stable ID.") + raise ValueError(f"No feed found for stable ID {stable_id}.") + return feed - return location_groups + +@track_metrics(metrics=("time", "memory", "cpu")) +def update_feed_location( + cache_location_groups: Dict[str, GeopolygonAggregate], + feed: Feed, + logger: Logger, + db_session: Session, +): + osm_location_groups = [ + get_or_create_feed_osm_location_group( + feed.id, cache_location_groups[location_group.group_id], db_session + ) + for location_group in cache_location_groups.values() + ] + feed.feedosmlocationgroups.clear() + feed.feedosmlocationgroups.extend(osm_location_groups) + feed_locations = [] + # The location_ids set is used to avoid duplicates when creating locations. + # Fixes: https://github.com/MobilityData/mobility-feed-api/issues/1289 + location_ids = set() + for location_aggregate in cache_location_groups.values(): + location = get_or_create_location(location_aggregate, logger, db_session) + if location: + if location.id not in location_ids: + feed_locations.append(location) + location_ids.add(location.id) + if feed.data_type == "gtfs": + gtfs_feed = db_session.query(Gtfsfeed).filter(Feed.id == feed.id).one_or_none() + for gtfs_rt_feed in gtfs_feed.gtfs_rt_feeds: + logger.info( + "Updating GTFS-RT feed with stable ID %s", gtfs_rt_feed.stable_id + ) + gtfs_rt_feed.feedosmlocationgroups.clear() + gtfs_rt_feed.feedosmlocationgroups.extend(osm_location_groups) + if feed_locations: + gtfs_rt_feed.locations.clear() + gtfs_rt_feed.locations = feed_locations + if feed_locations: + feed.locations = feed_locations + # Commit the changes to the database + db_session.commit() diff --git a/functions-python/reverse_geolocation/src/scripts/reverse_geolocation_process_verifier.py b/functions-python/reverse_geolocation/src/scripts/reverse_geolocation_process_verifier.py index 7194fa92c..00bbb884c 100644 --- a/functions-python/reverse_geolocation/src/scripts/reverse_geolocation_process_verifier.py +++ b/functions-python/reverse_geolocation/src/scripts/reverse_geolocation_process_verifier.py @@ -1,26 +1,88 @@ +# This script is used to verify the reverse geolocation process +# Before running this script, ensure you have the necessary environment set up: +# 1. Google DataStore emulator running on localhost:8081 by running: +# gcloud beta emulators datastore start --project=your-project-id + import json import logging import os +import uuid from io import BytesIO +from typing import Dict import folium import requests from dotenv import load_dotenv from google.cloud import storage +from sqlalchemy.orm import Session from reverse_geolocation_processor import reverse_geolocation_process +from shared.database.database import with_db_session +from shared.database_gen.sqlacodegen_models import Gtfsfeed, Gbfsfeed +from shared.helpers.locations import ReverseGeocodingStrategy from shared.helpers.logger import init_logger +from shared.helpers.runtime_metrics import track_metrics HOST = "localhost" PORT = 9023 BUCKET_NAME = "verifier" -feed_stable_id = "mdb-1" -station_information_url = ( - "https://data.rideflamingo.com/gbfs/3/auckland/station_information.json" +feeds = [ + { + # 0. 1539 stops, NZ, 1 location + "stable_id": "local-test-gbfs-flamingo_auckland", + "station_information_url": "https://data.rideflamingo.com/gbfs/3/auckland/station_information.json", + "vehicle_status_url": "https://data.rideflamingo.com/gbfs/3/auckland/vehicle_status.json", + "data_type": "gbfs", + }, + { + # 1. 11777 stops, JP, 241 locations + "stable_id": "local-test-gbfs-hellocycling", + "station_information_url": "https://api-public.odpt.org/api/v4/gbfs/hellocycling/station_information.json", + "data_type": "gbfs", + }, + { + # 2. 308611, UK aggregated, 225 locations + "stable_id": "local-test-2014", + "stops_url": "https://storage.googleapis.com/mobilitydata-datasets-prod/mdb-2014/" + "mdb-2014-202508120303/extracted/stops.txt", + "data_type": "gtfs", + }, + { + # 3. 663 stops, Europe, 334 locations + "stable_id": "local-test-1139", + "stops_url": "https://storage.googleapis.com/mobilitydata-datasets-prod/mdb-1139/" + "mdb-1139-202406071559/stops.txt", + "data_type": "gtfs", + }, + { + # 4. 10985 stops, Spain, duplicate key error(https://github.com/MobilityData/mobility-feed-api/issues/1289) + "stable_id": "local-test-gtfs-mdb-2825", + "stops_url": "https://storage.googleapis.com/mobilitydata-datasets-prod/mdb-2825/" + "mdb-2825-202508181628/extracted/stops.txt", + "data_type": "gtfs", + }, + { + # 5. 1355 stops, Canada, Saskatchewan , + # Duplicate admin level https://github.com/MobilityData/mobility-feed-api/issues/965 + "stable_id": "local-test-gtfs-mdb-716", + "stops_url": "https://storage.googleapis.com/mobilitydata-datasets-dev/mdb-716/mdb-716-202507082001/" + "extracted/stops.txt", + "data_type": "gtfs", + }, + { + # 6. 667 stops, Canada, New Brunswick, + # Duplicate admin level https://github.com/MobilityData/mobility-feed-api/issues/965 + "stable_id": "local-test-gtfs-mdb-1111", + "stops_url": "https://storage.googleapis.com/mobilitydata-datasets-dev/mdb-1111/mdb-1111-202507082012/" + "extracted/stops.txt", + "data_type": "gtfs", + }, +] +run_with_feed_index = ( + 1 # Set to an integer index to run with a specific feed from the list above ) -vehicle_status_url = "https://data.rideflamingo.com/gbfs/3/auckland/vehicle_status.json" -free_bike_status_url = None + # Load environment variables from .env.local load_dotenv(dotenv_path=".env.local") @@ -28,7 +90,10 @@ init_logger() -def download_to_local(url: str, filename: str): +@track_metrics(metrics=("time", "memory", "cpu")) +def download_to_local( + feed_stable_id: str, url: str, filename: str, force_download: bool = False +): """ Download a file from a URL and upload it to the Google Cloud Storage emulator. If the file already exists, it will not be downloaded again. @@ -36,13 +101,15 @@ def download_to_local(url: str, filename: str): url (str): The URL to download the file from. filename (str): The name of the file to save in the emulator. """ + if not url: + return blob_path = f"{feed_stable_id}/{filename}" client = storage.Client() bucket = client.bucket(BUCKET_NAME) blob = bucket.blob(blob_path) # Check if the blob already exists in the emulator - if not blob.exists(): + if not blob.exists() or force_download: logging.info(f"Downloading and uploading: {blob_path}") with requests.get(url, stream=True) as response: response.raise_for_status() @@ -55,7 +122,13 @@ def download_to_local(url: str, filename: str): logging.info(f"Blob already exists: gs://{BUCKET_NAME}/{blob_path}") -def verify_reverse_geolocation_process(): +def verify_reverse_geolocation_process( + feed_stable_id: str, + feed_dict: Dict, + data: Dict, + strategy: ReverseGeocodingStrategy, + force_download: bool = True, +): """ Verify the reverse geolocation process by downloading the necessary files, triggering the reverse geolocation process, and visualizing the resulting GeoJSON file. @@ -64,8 +137,36 @@ def verify_reverse_geolocation_process(): location, which can be viewed in a web browser. """ app = Flask(__name__) - download_to_local(url=station_information_url, filename="station_information.json") - download_to_local(url=vehicle_status_url, filename="vehicle_status.json") + + if feed_dict["data_type"] == "gbfs": + if "station_information_url" in feed_dict: + download_to_local( + feed_stable_id=feed_stable_id, + url=feed_dict["station_information_url"], + filename="station_information.json", + force_download=True, + ) + if "vehicle_status_url" in feed_dict: + download_to_local( + feed_stable_id=feed_stable_id, + url=feed_dict["vehicle_status_url"], + filename="vehicle_status.json", + force_download=True, + ) + if "free_bike_status_url" in feed_dict: + download_to_local( + feed_stable_id=feed_stable_id, + url=feed_dict["free_bike_status_url"], + filename="free_bike_status.json", + force_download=True, + ) + else: + download_to_local( + feed_stable_id=feed_stable_id, + url=feed_dict["stops_url"], + filename="stops.txt", + force_download=force_download, + ) with app.test_request_context( path="/reverse_geolocation", @@ -97,7 +198,37 @@ def verify_reverse_geolocation_process(): m.fit_bounds([[miny, minx], [maxy, maxx]]) # [[south, west], [north, east]] # Save the map - m.save(f".cloudstorage/verifier/{feed_stable_id}/geojson_map.html") + m.save( + f".cloudstorage/verifier/{feed_stable_id}/geojson_map_{strategy.value}.html" + ) + + +@with_db_session +def create_test_data(feed_stable_id: str, feed_dict: Dict, db_session: Session = None): + """ + Create test data in the database if it does not exist. + This function is used to ensure that the reverse geolocation process has the necessary data to work with. + """ + # Here you would typically interact with your database to create the necessary test data + # For this example, we will just log the action + logging.info(f"Creating test data for {feed_stable_id} with data: {feed_dict}") + model = Gtfsfeed if feed_dict["data_type"] == "gtfs" else Gbfsfeed + local_feed = ( + db_session.query(model).filter(model.stable_id == feed_stable_id).one_or_none() + ) + if not local_feed: + local_feed = model( + id=uuid.uuid4(), + stable_id=feed_stable_id, + data_type=feed_dict["data_type"], + feed_name="Test Feed", + note="This is a test feed created for reverse geolocation verification.", + producer_url="https://files.mobilitydatabase.org/mdb-2014/mdb-2014-202508120303/mdb-2014-202508120303.zip", + authentication_type="0", + status="active", + ) + db_session.add(local_feed) + db_session.commit() if __name__ == "__main__": @@ -105,26 +236,47 @@ def verify_reverse_geolocation_process(): from gcp_storage_emulator.server import create_server from flask import Flask, Request + strategy = ReverseGeocodingStrategy.PER_POINT + + feed_dict = feeds[run_with_feed_index] + feed_stable_id = feed_dict["stable_id"] + # create test data in the database if does not exist + create_test_data(feed_stable_id=feed_stable_id, feed_dict=feed_dict) data = { "stable_id": feed_stable_id, - "dataset_id": "example_dataset_id", - "station_information_url": f"http://{HOST}:{PORT}/{BUCKET_NAME}/{feed_stable_id}/station_information.json", - "vehicle_status_url": f"http://{HOST}:{PORT}/{BUCKET_NAME}/{feed_stable_id}/vehicle_status.json", - "free_bike_status_url": free_bike_status_url, - "strategy": "per-point", - "data_type": "gbfs", - "public": "False", + "dataset_id": feed_dict["dataset_id"] if "dataset_id" in feed_dict else None, + "station_information_url": f"http://{HOST}:{PORT}/{BUCKET_NAME}/{feed_stable_id}/station_information.json" + if "station_information_url" in feed_dict + else None, + "vehicle_status_url": f"http://{HOST}:{PORT}/{BUCKET_NAME}/{feed_stable_id}/vehicle_status.json" + if "vehicle_status_url" in feed_dict + else None, + "free_bike_status_url": f"http://{HOST}:{PORT}/{BUCKET_NAME}/{feed_stable_id}/free_bike_status.json" + if "free_bike_status_url" in feed_dict + else None, + "stops_url": f"http://{HOST}:{PORT}/{BUCKET_NAME}/{feed_stable_id}/stops.txt", + "strategy": str(strategy.value), + "data_type": feed_dict["data_type"], + # "use_cache": False, + "public": False, + "maximum_executions": 1000, } try: os.environ["STORAGE_EMULATOR_HOST"] = f"http://{HOST}:{PORT}" os.environ["DATASETS_BUCKET_NAME_GBFS"] = BUCKET_NAME os.environ["DATASETS_BUCKET_NAME_GTFS"] = BUCKET_NAME + os.environ["DATASTORE_EMULATOR_HOST"] = "localhost:8081" server = create_server( host=HOST, port=PORT, in_memory=False, default_bucket=BUCKET_NAME ) server.start() - verify_reverse_geolocation_process() + verify_reverse_geolocation_process( + feed_stable_id=feed_stable_id, + feed_dict=feed_dict, + strategy=strategy, + data=data, + ) except Exception as e: logging.error(f"Error verifying download content: {e}") finally: diff --git a/functions-python/reverse_geolocation/src/strategy_extraction_per_point.py b/functions-python/reverse_geolocation/src/strategy_extraction_per_point.py new file mode 100644 index 000000000..6a81e4412 --- /dev/null +++ b/functions-python/reverse_geolocation/src/strategy_extraction_per_point.py @@ -0,0 +1,68 @@ +import logging +from typing import Dict + +import pandas as pd +from sqlalchemy.orm import Session + +from location_group_utils import ( + GeopolygonAggregate, + extract_location_aggregate, + create_or_update_stop_group, +) +from shared.database.database import with_db_session +from shared.database_gen.sqlacodegen_models import Feed +from shared.helpers.runtime_metrics import track_metrics + + +@with_db_session +@track_metrics(metrics=("time", "memory", "cpu")) +def extract_location_aggregates_per_point( + feed: Feed, + stops_df: pd.DataFrame, + location_aggregates: Dict[str, GeopolygonAggregate], + use_cache: bool, + logger: logging.Logger, + db_session: Session, +) -> None: + """Extract the location aggregates for the stops. The location_aggregates dictionary will be updated with the new + location groups, keeping track of the stop count for each aggregate.""" + i = 0 + total_stop_count = len(stops_df) + if total_stop_count == 0: + logger.warning("No stops to process") + return + batch_size = max( + 1, int(total_stop_count * 0.05) + ) # Process ~5% of the total stops in each batch + for _, stop in stops_df.iterrows(): + if i % batch_size == 0: + remaining_stops_count = total_stop_count - i + logger.info( + "Progress %.2f%% (%d/%d)", + 100 - (remaining_stops_count / total_stop_count) * 100, + remaining_stops_count, + total_stop_count, + ) + i += 1 + location_aggregate = extract_location_aggregate( + stop["geometry"], logger, db_session + ) + if not location_aggregate: + continue + if use_cache: + create_or_update_stop_group( + feed, + stop["geometry"], + location_aggregate.location_group, + logger, + db_session, + ) + + if location_aggregate.group_id in location_aggregates: + location_aggregates[location_aggregate.group_id].merge(location_aggregate) + else: + location_aggregates[location_aggregate.group_id] = location_aggregate + if ( + i % batch_size == 0 + ): # Commit every batch stops to avoid reprocessing all stops in case of failure + db_session.commit() diff --git a/functions-python/reverse_geolocation/src/strategy_extraction_per_polygon.py b/functions-python/reverse_geolocation/src/strategy_extraction_per_polygon.py new file mode 100644 index 000000000..79cf71a32 --- /dev/null +++ b/functions-python/reverse_geolocation/src/strategy_extraction_per_polygon.py @@ -0,0 +1,175 @@ +import json +import logging +import os +from functools import lru_cache +from typing import Dict + +import pandas as pd +from sqlalchemy.orm import Session + +from location_group_utils import ( + GeopolygonAggregate, + extract_location_aggregate_geopolygons, + create_or_update_stop_group, +) + +from shared.database.database import with_db_session +from shared.database_gen.sqlacodegen_models import Feed +from shared.helpers.locations import ( + select_highest_level_polygon, + to_shapely, + get_country_code_from_polygons, + get_geopolygons_covers, +) +from shared.helpers.runtime_metrics import track_metrics + + +# TODO: Move this to a configuration service or a config file +# TODO: review admin_level threshold per country/region +@lru_cache(maxsize=1) +def get_country_locality_admin_levels() -> Dict[str, int]: + """ + Lazily load and cache the country locality admin levels from the environment variable. + """ + return json.loads(os.getenv("COUNTRY_LOCALITY_ADMIN_LEVELS", '{"JP": 7}')) + + +@lru_cache(maxsize=1) +def get_country_locality_admin_level_default() -> Dict[str, int]: + """ + Lazily load and cache the country locality admin levels from the environment variable. + """ + return os.getenv("COUNTRY_LOCALITY_ADMIN_LEVEL_DEFAULT", 7) + + +# TODO: Move this to a configuration service or a config file +def get_country_locality_admin_level(country_code: str) -> Dict[str, int]: + """ + Get the country locality admin levels from the environment variable. + If the variable is not set, return a default mapping. + The default mapping is: + {"JP": 7} + """ + return get_country_locality_admin_levels().get( + country_code, get_country_locality_admin_level_default() + ) + + +@with_db_session +@track_metrics(metrics=("time", "memory", "cpu")) +def extract_location_aggregates_per_polygon( + feed: Feed, + stops_df: pd.DataFrame, + location_aggregates: Dict[str, GeopolygonAggregate], + use_cache: bool, + logger: logging.Logger, + db_session: Session, +) -> None: + """ + Batch points by their containing geopolygon and compute one location aggregate per group. + """ + remaining_stops_df = stops_df.copy() + processed_groups: set[str] = set() + + total_stop_count = len(remaining_stops_df) + last_seen_count = total_stop_count + batch_size = max( + 1, int(total_stop_count * 0.05) + ) # Process ~5% of the total stops in each batch + stop_clustered_total = 0 + while not remaining_stops_df.empty: + if (last_seen_count - len(remaining_stops_df)) >= batch_size or len( + remaining_stops_df + ) == total_stop_count: + logger.info( + "Progress %.2f%% (%d/%d)", + 100 - (len(remaining_stops_df) / total_stop_count) * 100, + len(remaining_stops_df), + total_stop_count, + ) + last_seen_count = len(remaining_stops_df) + # Commit the changes to the database after processing the batch + db_session.commit() + + stop_point = remaining_stops_df.iloc[0][ + "geometry" + ] # GeoAlchemy WKT/WKB element or WKT string + stops_in_polygon = remaining_stops_df.iloc[[0]] + # remove the first point from the remaining stops + remaining_stops_df = remaining_stops_df.iloc[1:] + + # Get all polygons containing this point (SQL, uses DB index on geopolygon.geometry) + geopolygons = get_geopolygons_covers(stop_point, db_session) + + highest = select_highest_level_polygon(geopolygons) + if highest is None or highest.geometry is None: + logger.warning("No geopolygons found for point: %s", stop_point) + continue + + country_code = get_country_code_from_polygons(geopolygons) + if highest.admin_level >= get_country_locality_admin_level(country_code): + # If admin_level >= locality_admin_level, we can filter points inside this polygon + # Convert to Shapely geometry for spatial operations + poly_shp = to_shapely(highest.geometry) + in_poly_mask = remaining_stops_df["geometry"].apply( + lambda g: poly_shp.contains(to_shapely(g)) + ) + count_before = len(remaining_stops_df) + stops_mask = remaining_stops_df.loc[in_poly_mask] + # concat the points that are inside the polygon + # This will include the point that is being processed in this iteration + stops_in_polygon = pd.concat([stops_in_polygon, stops_mask]) + # Remove them from the remaining pool + remaining_stops_df = remaining_stops_df.drop(stops_mask.index) + logger.debug( + "Points clustered in polygon %s: %d", + highest.admin_level, + count_before - len(remaining_stops_df), + ) + stop_clustered_total += count_before - len(remaining_stops_df) + else: + # If admin_level < locality_admin_level, we assume the polygon is too large to filter points + # directly, so we just use the first point as a representative + logger.debug( + "Point cannot be clustered in polygon. " + "osm_id: %s, name: %s, " + "iso_3166_1_code: %s, iso_3166_2_code: %s, admin_level: %s, ", + highest.osm_id, + highest.name, + highest.iso_3166_1_code, + highest.iso_3166_2_code, + highest.admin_level, + ) + + # Process ONLY ONE representative point for this stop "cluster" + location_aggregate = extract_location_aggregate_geopolygons( + stop_point=stop_point, + geopolygons=geopolygons, + logger=logger, + db_session=db_session, + ) + if not location_aggregate or location_aggregate.group_id in processed_groups: + continue + location_aggregate.stop_count = len(stops_in_polygon) + if use_cache: + for _, stop in stops_in_polygon.iterrows(): + # Create or update the stop group for each stop in the cluster + # This will also update the location_aggregate with the stop count + create_or_update_stop_group( + feed, + stop["geometry"], + location_aggregate.location_group, + logger, + db_session, + ) + processed_groups.add(location_aggregate.group_id) + + if location_aggregate.group_id in location_aggregates: + location_aggregates[location_aggregate.group_id].merge(location_aggregate) + else: + location_aggregates[location_aggregate.group_id] = location_aggregate + # Make sure to commit the changes after processing all points + db_session.commit() + logger.info( + "Completed processing all points with clustered total %d", stop_clustered_total + ) diff --git a/functions-python/reverse_geolocation/tests/test_extraction_location_per_point.py b/functions-python/reverse_geolocation/tests/test_extraction_location_per_point.py new file mode 100644 index 000000000..898af9071 --- /dev/null +++ b/functions-python/reverse_geolocation/tests/test_extraction_location_per_point.py @@ -0,0 +1,108 @@ +import logging +import unittest +from unittest.mock import patch + + +import pandas as pd +from faker import Faker +from geoalchemy2 import WKTElement + +from location_group_utils import GeopolygonAggregate +from shared.database.database import with_db_session +from shared.database_gen.sqlacodegen_models import ( + Gtfsfeed, + Gtfsrealtimefeed, + Geopolygon, + Osmlocationgroup, +) +from test_shared.test_utils.database_utils import clean_testing_db, default_db_url + + +faker = Faker() +logger = logging.getLogger(__name__) + + +class TestReverseGeolocationProcessor(unittest.TestCase): + @patch("location_group_utils.extract_location_aggregate") + @with_db_session(db_url=default_db_url) + def test_extract_location_aggregates_per_point( + self, + mock_extract_location_aggregate, + db_session, + ): + from strategy_extraction_per_point import extract_location_aggregates_per_point + + clean_testing_db() + + # Create sample feed + feed_id = faker.uuid4(cast_to=str) + feed = Gtfsfeed(id=feed_id, stable_id=faker.uuid4(cast_to=str)) + gtfs_rt_feed = Gtfsrealtimefeed( + id=faker.uuid4(cast_to=str), stable_id=faker.uuid4(cast_to=str) + ) + feed.gtfs_rt_feeds = [gtfs_rt_feed] + db_session.add(feed) + db_session.commit() + + # Prepare stops DataFrame + stops_df = pd.DataFrame( + { + "stop_id": [1, 2, 3], + "stop_lat": [2.0, 3.0, 10.0], # Two inside polygon, one unmatched + "stop_lon": [2.0, 3.0, 10.0], + } + ) + + stops_df["geometry"] = stops_df.apply( + lambda x: WKTElement(f"POINT ({x['stop_lon']} {x['stop_lat']})", srid=4326), + axis=1, + ) + + # Prepare mock GeopolygonAggregate for matched stops + geopolygon = Geopolygon( + osm_id=faker.random_int(), + name=faker.city(), + admin_level=3, + geometry=WKTElement("POLYGON((0 0, 5 0, 5 5, 0 5, 0 0))", srid=4326), + iso_3166_1_code=faker.country_code(), + ) + + group = Osmlocationgroup( + group_id=faker.uuid4(cast_to=str), + group_name=geopolygon.name, + osms=[geopolygon], + ) + + mock_aggregate = GeopolygonAggregate(group, stops_count=1) + db_session.add(group) + db_session.commit() + + # Mock extract_location_aggregate behavior + def side_effect(stop_geometry, _, __): + if stop_geometry.data == "POINT (10.0 10.0)": # Simulate unmatched stop + return None + return mock_aggregate + + mock_extract_location_aggregate.side_effect = side_effect + + # Prepare location_aggregates dict (empty initially) + location_aggregates = {} + + # Call the function + extract_location_aggregates_per_point( + feed=feed, + stops_df=stops_df, + location_aggregates=location_aggregates, + use_cache=False, + logger=logger, + db_session=db_session, + ) + + # Assertions + # Ensure only matched stops are aggregated + self.assertEqual(len(location_aggregates), 1) + first_aggregate = list(location_aggregates.values())[0] + self.assertIsInstance(first_aggregate, GeopolygonAggregate) + self.assertEqual(first_aggregate.stop_count, 2) # Two matched stops + + db_session.close_all() diff --git a/functions-python/reverse_geolocation/tests/test_extraction_location_per_polygon.py b/functions-python/reverse_geolocation/tests/test_extraction_location_per_polygon.py new file mode 100644 index 000000000..1c216cdc7 --- /dev/null +++ b/functions-python/reverse_geolocation/tests/test_extraction_location_per_polygon.py @@ -0,0 +1,116 @@ +import logging +import unittest +from unittest.mock import patch +import pandas as pd +from geoalchemy2 import WKTElement +from shapely import wkt +from shared.database.database import with_db_session +from shared.database_gen.sqlacodegen_models import Gtfsfeed, Geopolygon +from test_shared.test_utils.database_utils import clean_testing_db, default_db_url + +logger = logging.getLogger(__name__) + + +class TestExtractLocationAggregatesPerPolygon(unittest.TestCase): + @with_db_session(db_url=default_db_url) + @patch("strategy_extraction_per_polygon.get_geopolygons_covers") + def test_valid_stops_processing(self, mock_get_geopolygons_covers, db_session): + from strategy_extraction_per_polygon import ( + extract_location_aggregates_per_polygon, + ) + + # Clean the database + clean_testing_db() + + # Create a test feed + feed = Gtfsfeed(id="test_feed", stable_id="test_feed", status="active") + db_session.add(feed) + db_session.commit() + + # Prepare stops DataFrame + stops_df = pd.DataFrame( + { + "stop_id": [1, 2], + "stop_lat": [2.0, 3.0], + "stop_lon": [2.0, 3.0], + } + ) + stops_df["geometry"] = stops_df.apply( + lambda x: WKTElement(f"POINT ({x['stop_lon']} {x['stop_lat']})", srid=4326), + axis=1, + ) + + qc_geopolygons = [ + Geopolygon( + osm_id=1, + admin_level=2, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_1_code="CA", + name="Canada", + ), + Geopolygon( + osm_id=2, + admin_level=4, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_2_code="CA-QC", + name="Quebec", + ), + ] + mtl_geopolygons = [ + Geopolygon( + osm_id=3, + admin_level=2, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_1_code="CA", + name="Canada", + ), + Geopolygon( + osm_id=4, + admin_level=4, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_2_code="CA-QC", + name="Quebec", + ), + Geopolygon( + osm_id=5, + admin_level=7, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_2_code="CA-QC", + name="Montreal", + ), + ] + + def mock_geopolygons_covers(stop: WKTElement, *args, **kwargs): + point = wkt.loads(stop.desc) + # Extract latitude + longitude = point.x + if longitude == 2.0: + return mtl_geopolygons + else: + return qc_geopolygons + + mock_get_geopolygons_covers.side_effect = mock_geopolygons_covers + + # Prepare location_aggregates dict + location_aggregates = {} + + # Call the function + extract_location_aggregates_per_polygon( + feed=feed, + stops_df=stops_df, + location_aggregates=location_aggregates, + use_cache=False, + logger=logger, + db_session=db_session, + ) + + # Assertions + self.assertEqual(2, len(location_aggregates)) + self.assertTrue("3.4.5" in location_aggregates) + self.assertTrue("1.2" in location_aggregates) + location_aggregate = location_aggregates["1.2"] + self.assertEqual("1.2", location_aggregate.group_id) + self.assertEqual(1, location_aggregate.stop_count) + self.assertEqual("CA", location_aggregate.iso_3166_1_code) + + clean_testing_db() diff --git a/functions-python/reverse_geolocation/tests/test_location_group_utils.py b/functions-python/reverse_geolocation/tests/test_location_group_utils.py index 4bbb29728..5966d1c4d 100644 --- a/functions-python/reverse_geolocation/tests/test_location_group_utils.py +++ b/functions-python/reverse_geolocation/tests/test_location_group_utils.py @@ -1,13 +1,23 @@ +import logging import unittest +from unittest.mock import MagicMock import pytest from faker import Faker from geoalchemy2 import WKTElement from location_group_utils import GeopolygonAggregate -from shared.database_gen.sqlacodegen_models import Geopolygon, Osmlocationgroup +from shared.database.database import with_db_session +from shared.database_gen.sqlacodegen_models import ( + Geopolygon, + Osmlocationgroup, + Gtfsfeed, + Location, +) +from test_shared.test_utils.database_utils import default_db_url, clean_testing_db faker = Faker() +logger = logging.getLogger(__name__) class TestLocationGroupUtils(unittest.TestCase): @@ -76,6 +86,259 @@ def test_geopolygon_aggregate(self): geopolygon_aggregate.merge(geopolygon_aggregate_2) self.assertEqual(geopolygon_aggregate.stop_count, 2) + @with_db_session(db_url=default_db_url) + def test_extract_location_group(self, db_session): + from location_group_utils import extract_location_aggregate + + clean_testing_db() + geopolygon_country_lvl = Geopolygon( + osm_id=faker.random_int(), + name=faker.country(), + admin_level=2, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_1_code=faker.country_code(), + ) + geopolygon_subdivision_lvl = Geopolygon( + osm_id=faker.random_int(), + name=faker.city(), + admin_level=3, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_2_code=faker.country_code(), + ) + feed_id = faker.uuid4(cast_to=str) + feed = Gtfsfeed(id=feed_id, stable_id=faker.uuid4(cast_to=str)) + db_session.add(geopolygon_country_lvl) + db_session.add(geopolygon_subdivision_lvl) + db_session.add(feed) + db_session.commit() + stop_wkt = WKTElement("POINT (0.5 0.5)", srid=4326) + aggregate = extract_location_aggregate(stop_wkt, logger, db_session) + self.assertTrue( + aggregate.iso_3166_1_code == geopolygon_country_lvl.iso_3166_1_code + ) + self.assertTrue( + aggregate.iso_3166_2_code == geopolygon_subdivision_lvl.iso_3166_2_code + ) + + @with_db_session(db_url=default_db_url) + def test_extract_location_duplicate_admin_level(self, db_session): + from location_group_utils import extract_location_aggregate + + print("test extract location duplicate admin level") + clean_testing_db() + print("Done cleaning the db") + geopolygon_country_lvl = Geopolygon( + osm_id=faker.random_int(), + name=faker.country(), + admin_level=2, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_1_code=faker.country_code(), + ) + geopolygon_subdivision_lvl = Geopolygon( + osm_id=faker.random_int(), + name=faker.city(), + admin_level=2, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_2_code=faker.country_code(), + ) + feed_id = faker.uuid4(cast_to=str) + feed = Gtfsfeed(id=feed_id, stable_id=faker.uuid4(cast_to=str)) + db_session.add(geopolygon_country_lvl) + db_session.add(geopolygon_subdivision_lvl) + db_session.add(feed) + db_session.commit() + stop_wkt = WKTElement("POINT (0.5 0.5)", srid=4326) + aggregate = extract_location_aggregate(stop_wkt, logger, db_session) + self.assertIsNone(aggregate) + + @with_db_session(db_url=default_db_url) + def test_extract_location_not_enough_geopolygons(self, db_session): + from location_group_utils import extract_location_aggregate + + clean_testing_db() + geopolygon_country_lvl = Geopolygon( + osm_id=faker.random_int(), + name=faker.country(), + admin_level=2, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_1_code=faker.country_code(), + ) + feed_id = faker.uuid4(cast_to=str) + feed = Gtfsfeed(id=feed_id, stable_id=faker.uuid4(cast_to=str)) + db_session.add(geopolygon_country_lvl) + db_session.add(feed) + db_session.commit() + stop_wkt = WKTElement("POINT (0.5 0.5)", srid=4326) + aggregate = extract_location_aggregate(stop_wkt, logger, db_session) + self.assertIsNone(aggregate) + + @with_db_session(db_url=default_db_url) + def test_extract_location_missing_iso_codes(self, db_session): + from location_group_utils import extract_location_aggregate + + clean_testing_db() + geopolygon_country_lvl = Geopolygon( + osm_id=faker.random_int(), + name=faker.country(), + admin_level=2, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_1_code=faker.country_code(), + ) + geopolygon_subdivision_lvl = Geopolygon( + osm_id=faker.random_int(), + name=faker.city(), + admin_level=3, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + ) + feed_id = faker.uuid4(cast_to=str) + feed = Gtfsfeed(id=feed_id, stable_id=faker.uuid4(cast_to=str)) + db_session.add(geopolygon_country_lvl) + db_session.add(geopolygon_subdivision_lvl) + db_session.add(feed) + db_session.commit() + stop_wkt = WKTElement("POINT (0.5 0.5)", srid=4326) + aggregate = extract_location_aggregate(stop_wkt, logger, db_session) + self.assertIsNone(aggregate) + + @with_db_session(db_url=default_db_url) + def test_create_feed_osm_location(self, db_session): + from location_group_utils import get_or_create_feed_osm_location_group + + clean_testing_db() + feed_id = faker.uuid4(cast_to=str) + feed = Gtfsfeed(id=feed_id, stable_id=faker.uuid4(cast_to=str)) + db_session.add(feed) + db_session.commit() + stops_count = faker.random_int() + aggregate = GeopolygonAggregate( + Osmlocationgroup( + group_id=faker.uuid4(cast_to=str), + group_name=faker.country(), + osms=[ + Geopolygon( + osm_id=faker.random_int(), + admin_level=2, + geometry=WKTElement( + "POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326 + ), + iso_3166_1_code=faker.country_code(), + ) + ], + ), + stops_count=stops_count, + ) + feed_osm_location = get_or_create_feed_osm_location_group( + feed_id, aggregate, db_session + ) + self.assertIsNotNone(feed_osm_location) + self.assertEqual(feed_osm_location.stops_count, stops_count) + + @with_db_session(db_url=default_db_url) + def test_create_new_location(self, db_session): + from reverse_geolocation_processor import get_or_create_location + + # Clean DB before test + clean_testing_db() + + # Create sample Geopolygon and GeopolygonAggregate + geopolygon_country = Geopolygon( + osm_id=faker.random_int(), + name=faker.city(), + admin_level=3, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_1_code=faker.country_code(), + ) + geopolygon_subdivision = Geopolygon( + osm_id=faker.random_int(), + name=faker.city(), + admin_level=3, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_2_code=faker.country_code(), + ) + + group = Osmlocationgroup( + group_id=faker.uuid4(cast_to=str), + group_name=f"{geopolygon_country.name}, {geopolygon_subdivision.name}", + osms=[geopolygon_country, geopolygon_subdivision], + ) + + location_aggregate = GeopolygonAggregate(group, stops_count=5) + + # Call the function + location = get_or_create_location(location_aggregate, logger, db_session) + + # Assert location is created + self.assertIsNotNone(location) + self.assertEqual(location.id, location_aggregate.location_id()) + self.assertEqual(location.country_code, location_aggregate.iso_3166_1_code) + self.assertEqual(location.country, location_aggregate.country()) + self.assertEqual( + location.subdivision_name, location_aggregate.subdivision_name() + ) + self.assertEqual(location.municipality, location_aggregate.municipality()) + + @with_db_session(db_url=default_db_url) + def test_retrieve_existing_location(self, db_session): + from reverse_geolocation_processor import get_or_create_location + + # Clean DB before test + clean_testing_db() + + # Create sample Location + existing_location = Location( + id=faker.uuid4(cast_to=str), + country_code=faker.country_code(), + country=faker.country(), + subdivision_name=faker.state(), + municipality=faker.city(), + ) + + db_session.add(existing_location) + db_session.commit() + + # Create GeopolygonAggregate with same location_id + geopolygon_country = Geopolygon( + osm_id=faker.random_int(), + name=faker.city(), + admin_level=3, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_1_code=faker.country_code(), + ) + geopolygon_subdivision = Geopolygon( + osm_id=faker.random_int(), + name=faker.city(), + admin_level=3, + geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), + iso_3166_2_code=faker.country_code(), + ) + + group = Osmlocationgroup( + group_id=faker.uuid4(cast_to=str), + group_name=f"{geopolygon_country.name}, {geopolygon_subdivision.name}", + osms=[geopolygon_country, geopolygon_subdivision], + ) + + location_aggregate = GeopolygonAggregate(group, stops_count=5) + location_aggregate.location_id = ( + lambda: existing_location.id + ) # Mocking location_id method + + # Call the function + location = get_or_create_location(location_aggregate, logger, db_session) + + # Assert the existing location was returned + self.assertIsNotNone(location) + self.assertEqual(location.id, existing_location.id) + self.assertEqual(location.country, existing_location.country) + + def test_retrieve_location_exception(self): + from reverse_geolocation_processor import get_or_create_location + + mock_session = MagicMock() + mock_session.query.side_effect = Exception("Test exception") + location = get_or_create_location(None, logger, mock_session) + self.assertIsNone(location) + @pytest.mark.parametrize( "values", diff --git a/functions-python/reverse_geolocation/tests/test_reverse_geolocation_processor.py b/functions-python/reverse_geolocation/tests/test_reverse_geolocation_processor.py index 5d2ec1262..29eb61bfb 100644 --- a/functions-python/reverse_geolocation/tests/test_reverse_geolocation_processor.py +++ b/functions-python/reverse_geolocation/tests/test_reverse_geolocation_processor.py @@ -6,13 +6,14 @@ import pandas as pd import shapely from faker import Faker +from flask import Request from geoalchemy2 import WKTElement +from sqlalchemy.orm import Session -from location_group_utils import GeopolygonAggregate +from location_group_utils import GeopolygonAggregate, ERROR_STATUS_CODE from shared.database.database import with_db_session from shared.database_gen.sqlacodegen_models import ( Geopolygon, - Location, Gtfsdataset, Gtfsrealtimefeed, ) @@ -21,6 +22,7 @@ Feedlocationgrouppoint, Osmlocationgroup, ) +from shared.helpers.locations import ReverseGeocodingStrategy from test_shared.test_utils.database_utils import ( default_db_url, clean_testing_db, @@ -56,6 +58,8 @@ def test_parse_request_parameters(self, requests_mock): urls, public, strategy, + use_cache, + maximum_executions, ) = parse_request_parameters(request) self.assertEqual("test_stable_id", stable_id) self.assertEqual("test_dataset_id", dataset_id) @@ -64,6 +68,8 @@ def test_parse_request_parameters(self, requests_mock): self.assertEqual(["test_url"], urls) self.assertEqual(True, public) self.assertEqual("per-point", strategy) + self.assertEqual(True, use_cache) + self.assertEqual(1, maximum_executions) # Exception should be raised requests_mock.get.return_value.content = None @@ -102,6 +108,8 @@ def test_parse_request_parameters_gbfs_station_information(self, requests_mock): urls, public, strategy, + use_cache, + maximum_executions, ) = parse_request_parameters(request) self.assertEqual("stable123", stable_id) @@ -109,8 +117,11 @@ def test_parse_request_parameters_gbfs_station_information(self, requests_mock): self.assertEqual("gbfs", data_type) self.assertEqual("http://dummy.url", urls[0]) self.assertEqual((2, 2), df.shape) - self.assertEqual("per-point", strategy) + self.assertEqual("per-polygon", strategy) self.assertEqual(True, public) + # Cache is disabled for GBFS data by default + self.assertEqual(False, use_cache) + self.assertEqual(1, maximum_executions) @patch("parse_request.requests") def test_parse_request_parameters_gbfs_vehicle_status(self, requests_mock): @@ -132,6 +143,7 @@ def test_parse_request_parameters_gbfs_vehicle_status(self, requests_mock): "vehicle_status_url": "http://dummy.vehicle", "data_type": "gbfs", "public": "False", + "maximum_executions": 10, } ( @@ -142,6 +154,8 @@ def test_parse_request_parameters_gbfs_vehicle_status(self, requests_mock): urls, public, strategy, + use_cache, + maximum_executions, ) = parse_request_parameters(request) self.assertEqual("stable456", stable_id) @@ -149,8 +163,11 @@ def test_parse_request_parameters_gbfs_vehicle_status(self, requests_mock): self.assertEqual("gbfs", data_type) self.assertEqual("http://dummy.vehicle", urls[0]) self.assertEqual((2, 2), df.shape) - self.assertEqual("per-point", strategy) + self.assertEqual("per-polygon", strategy) self.assertEqual(False, public) + # Cache is disabled for GBFS data by default + self.assertEqual(False, use_cache) + self.assertEqual(10, maximum_executions) @patch("parse_request.requests") def test_parse_request_parameters_invalid_request(self, requests_mock): @@ -174,29 +191,16 @@ def test_parse_request_parameters_invalid_request(self, requests_mock): @with_db_session(db_url=default_db_url) def test_get_cached_geopolygons_empty_df(self, db_session): - from reverse_geolocation_processor import get_cached_geopolygons - - with self.assertRaises(ValueError): - get_cached_geopolygons("test-stable-id", pd.DataFrame(), logger) + from reverse_geolocation_processor import get_geopolygons_with_geometry - @with_db_session(db_url=default_db_url) - def test_get_cached_geopolygons_no_feed(self, db_session): - from reverse_geolocation_processor import get_cached_geopolygons - - stops_df = pd.DataFrame( - { - "stop_id": [1, 2], - "stop_name": ["stop1", "stop2"], - "stop_lat": [1.0, 2.0], - "stop_lon": [1.0, 2.0], - } - ) with self.assertRaises(ValueError): - get_cached_geopolygons("test-stable-id", stops_df, logger) + get_geopolygons_with_geometry( + "test-stable-id", pd.DataFrame(), False, logger + ) @with_db_session(db_url=default_db_url) def test_get_cached_geopolygons_no_cached_stop(self, db_session): - from reverse_geolocation_processor import get_cached_geopolygons + from reverse_geolocation_processor import get_geopolygons_with_geometry stops_df = pd.DataFrame( { @@ -215,16 +219,15 @@ def test_get_cached_geopolygons_no_cached_stop(self, db_session): ) db_session.add(feed) db_session.commit() - result_feed_id, location_groups, results_df = get_cached_geopolygons( - stable_id, stops_df, logger + location_groups, results_df = get_geopolygons_with_geometry( + feed, stops_df, False, logger ) - self.assertEqual(result_feed_id, feed_id) self.assertDictEqual(location_groups, {}) self.assertEqual(results_df.shape, (2, 6)) # Added geometry columns @with_db_session(db_url=default_db_url) def test_get_cached_geopolygons_w_cached_stop(self, db_session): - from reverse_geolocation_processor import get_cached_geopolygons + from reverse_geolocation_processor import get_geopolygons_with_geometry stops_df = pd.DataFrame( { @@ -264,164 +267,23 @@ def test_get_cached_geopolygons_w_cached_stop(self, db_session): ) db_session.add(feed) db_session.commit() - result_feed_id, location_groups, results_df = get_cached_geopolygons( - stable_id, stops_df, logger + location_groups, results_df = get_geopolygons_with_geometry( + feed, stops_df, True, logger ) - self.assertEqual(result_feed_id, feed_id) self.assertEqual(len(location_groups), 1) self.assertEqual(results_df.shape, (1, 6)) # Added geometry columns - @with_db_session(db_url=default_db_url) - def test_extract_location_group(self, db_session): - from reverse_geolocation_processor import extract_location_aggregate - - clean_testing_db() - geopolygon_country_lvl = Geopolygon( - osm_id=faker.random_int(), - name=faker.country(), - admin_level=2, - geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), - iso_3166_1_code=faker.country_code(), - ) - geopolygon_subdivision_lvl = Geopolygon( - osm_id=faker.random_int(), - name=faker.city(), - admin_level=3, - geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), - iso_3166_2_code=faker.country_code(), - ) - feed_id = faker.uuid4(cast_to=str) - feed = Gtfsfeed(id=feed_id, stable_id=faker.uuid4(cast_to=str)) - db_session.add(geopolygon_country_lvl) - db_session.add(geopolygon_subdivision_lvl) - db_session.add(feed) - db_session.commit() - stop_wkt = WKTElement("POINT (0.5 0.5)", srid=4326) - aggregate = extract_location_aggregate(feed_id, stop_wkt, logger, db_session) - self.assertTrue( - aggregate.iso_3166_1_code == geopolygon_country_lvl.iso_3166_1_code - ) - self.assertTrue( - aggregate.iso_3166_2_code == geopolygon_subdivision_lvl.iso_3166_2_code - ) - - @with_db_session(db_url=default_db_url) - def test_extract_location_duplicate_admin_level(self, db_session): - from reverse_geolocation_processor import extract_location_aggregate - - print("test extract location duplicate admin level") - clean_testing_db() - print("Done cleaning the db") - geopolygon_country_lvl = Geopolygon( - osm_id=faker.random_int(), - name=faker.country(), - admin_level=2, - geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), - iso_3166_1_code=faker.country_code(), - ) - geopolygon_subdivision_lvl = Geopolygon( - osm_id=faker.random_int(), - name=faker.city(), - admin_level=2, - geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), - iso_3166_2_code=faker.country_code(), - ) - feed_id = faker.uuid4(cast_to=str) - feed = Gtfsfeed(id=feed_id, stable_id=faker.uuid4(cast_to=str)) - db_session.add(geopolygon_country_lvl) - db_session.add(geopolygon_subdivision_lvl) - db_session.add(feed) - db_session.commit() - stop_wkt = WKTElement("POINT (0.5 0.5)", srid=4326) - aggregate = extract_location_aggregate(feed_id, stop_wkt, logger, db_session) - self.assertIsNone(aggregate) - - @with_db_session(db_url=default_db_url) - def test_extract_location_not_enough_geopolygons(self, db_session): - from reverse_geolocation_processor import extract_location_aggregate - - clean_testing_db() - geopolygon_country_lvl = Geopolygon( - osm_id=faker.random_int(), - name=faker.country(), - admin_level=2, - geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), - iso_3166_1_code=faker.country_code(), - ) - feed_id = faker.uuid4(cast_to=str) - feed = Gtfsfeed(id=feed_id, stable_id=faker.uuid4(cast_to=str)) - db_session.add(geopolygon_country_lvl) - db_session.add(feed) - db_session.commit() - stop_wkt = WKTElement("POINT (0.5 0.5)", srid=4326) - aggregate = extract_location_aggregate(feed_id, stop_wkt, logger, db_session) - self.assertIsNone(aggregate) - - @with_db_session(db_url=default_db_url) - def test_extract_location_missing_iso_codes(self, db_session): - from reverse_geolocation_processor import extract_location_aggregate + @patch("reverse_geolocation_processor.get_storage_client") + @patch("os.getenv") + def test_create_geojson_aggregate(self, mock_getenv, mock_storage_client): + from reverse_geolocation_processor import create_geojson_aggregate - clean_testing_db() - geopolygon_country_lvl = Geopolygon( - osm_id=faker.random_int(), - name=faker.country(), - admin_level=2, - geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), - iso_3166_1_code=faker.country_code(), + # Mock the specific environment variable + mock_getenv.side_effect = ( + lambda var_name: "test_bucket" + if var_name in ["DATASETS_BUCKET_NAME_GTFS", "DATASETS_BUCKET_NAME_GBFS"] + else None ) - geopolygon_subdivision_lvl = Geopolygon( - osm_id=faker.random_int(), - name=faker.city(), - admin_level=3, - geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), - ) - feed_id = faker.uuid4(cast_to=str) - feed = Gtfsfeed(id=feed_id, stable_id=faker.uuid4(cast_to=str)) - db_session.add(geopolygon_country_lvl) - db_session.add(geopolygon_subdivision_lvl) - db_session.add(feed) - db_session.commit() - stop_wkt = WKTElement("POINT (0.5 0.5)", srid=4326) - aggregate = extract_location_aggregate(feed_id, stop_wkt, logger, db_session) - self.assertIsNone(aggregate) - - @with_db_session(db_url=default_db_url) - def test_create_feed_osm_location(self, db_session): - from reverse_geolocation_processor import get_or_create_feed_osm_location_group - - clean_testing_db() - feed_id = faker.uuid4(cast_to=str) - feed = Gtfsfeed(id=feed_id, stable_id=faker.uuid4(cast_to=str)) - db_session.add(feed) - db_session.commit() - stops_count = faker.random_int() - aggregate = GeopolygonAggregate( - Osmlocationgroup( - group_id=faker.uuid4(cast_to=str), - group_name=faker.country(), - osms=[ - Geopolygon( - osm_id=faker.random_int(), - admin_level=2, - geometry=WKTElement( - "POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326 - ), - iso_3166_1_code=faker.country_code(), - ) - ], - ), - stops_count=stops_count, - ) - feed_osm_location = get_or_create_feed_osm_location_group( - feed_id, aggregate, db_session - ) - self.assertIsNotNone(feed_osm_location) - self.assertEqual(feed_osm_location.stops_count, stops_count) - - @patch("reverse_geolocation_processor.storage.Client") - @patch("os.getenv", return_value="test_bucket") - def test_create_geojson_aggregate(self, _, mock_storage_client): - from reverse_geolocation_processor import create_geojson_aggregate # Mock the storage client and blob mock_bucket = MagicMock() @@ -480,112 +342,6 @@ def test_create_geojson_aggregate(self, _, mock_storage_client): # Check if the blob was made public mock_blob.make_public.assert_called_once() - @with_db_session(db_url=default_db_url) - def test_create_new_location(self, db_session): - from reverse_geolocation_processor import get_or_create_location - - # Clean DB before test - clean_testing_db() - - # Create sample Geopolygon and GeopolygonAggregate - geopolygon_country = Geopolygon( - osm_id=faker.random_int(), - name=faker.city(), - admin_level=3, - geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), - iso_3166_1_code=faker.country_code(), - ) - geopolygon_subdivision = Geopolygon( - osm_id=faker.random_int(), - name=faker.city(), - admin_level=3, - geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), - iso_3166_2_code=faker.country_code(), - ) - - group = Osmlocationgroup( - group_id=faker.uuid4(cast_to=str), - group_name=f"{geopolygon_country.name}, {geopolygon_subdivision.name}", - osms=[geopolygon_country, geopolygon_subdivision], - ) - - location_aggregate = GeopolygonAggregate(group, stops_count=5) - - # Call the function - location = get_or_create_location(location_aggregate, logger, db_session) - - # Assert location is created - self.assertIsNotNone(location) - self.assertEqual(location.id, location_aggregate.location_id()) - self.assertEqual(location.country_code, location_aggregate.iso_3166_1_code) - self.assertEqual(location.country, location_aggregate.country()) - self.assertEqual( - location.subdivision_name, location_aggregate.subdivision_name() - ) - self.assertEqual(location.municipality, location_aggregate.municipality()) - - @with_db_session(db_url=default_db_url) - def test_retrieve_existing_location(self, db_session): - from reverse_geolocation_processor import get_or_create_location - - # Clean DB before test - clean_testing_db() - - # Create sample Location - existing_location = Location( - id=faker.uuid4(cast_to=str), - country_code=faker.country_code(), - country=faker.country(), - subdivision_name=faker.state(), - municipality=faker.city(), - ) - - db_session.add(existing_location) - db_session.commit() - - # Create GeopolygonAggregate with same location_id - geopolygon_country = Geopolygon( - osm_id=faker.random_int(), - name=faker.city(), - admin_level=3, - geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), - iso_3166_1_code=faker.country_code(), - ) - geopolygon_subdivision = Geopolygon( - osm_id=faker.random_int(), - name=faker.city(), - admin_level=3, - geometry=WKTElement("POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326), - iso_3166_2_code=faker.country_code(), - ) - - group = Osmlocationgroup( - group_id=faker.uuid4(cast_to=str), - group_name=f"{geopolygon_country.name}, {geopolygon_subdivision.name}", - osms=[geopolygon_country, geopolygon_subdivision], - ) - - location_aggregate = GeopolygonAggregate(group, stops_count=5) - location_aggregate.location_id = ( - lambda: existing_location.id - ) # Mocking location_id method - - # Call the function - location = get_or_create_location(location_aggregate, logger, db_session) - - # Assert the existing location was returned - self.assertIsNotNone(location) - self.assertEqual(location.id, existing_location.id) - self.assertEqual(location.country, existing_location.country) - - def test_retrieve_location_exception(self): - from reverse_geolocation_processor import get_or_create_location - - mock_session = MagicMock() - mock_session.query.side_effect = Exception("Test exception") - location = get_or_create_location(None, logger, mock_session) - self.assertIsNone(location) - @with_db_session(db_url=default_db_url) def test_update_dataset_bounding_box_success(self, db_session): from reverse_geolocation_processor import update_dataset_bounding_box @@ -658,85 +414,348 @@ def test_update_dataset_bounding_box_exception(self, db_session): faker.uuid4(cast_to=str), stops_df, db_session=db_session ) + @patch("reverse_geolocation_processor.parse_request_parameters") + @patch("reverse_geolocation_processor.update_dataset_bounding_box") + @patch("reverse_geolocation_processor.reverse_geolocation") + @patch("reverse_geolocation_processor.create_geojson_aggregate") + @patch("reverse_geolocation_processor.check_maximum_executions") + @patch("reverse_geolocation_processor.get_execution_id") + @patch("reverse_geolocation_processor.record_execution_trace") + def test_valid_request( + self, + _, + mock_get_execution_id, + mock_check_maximum_executions, + mock_create_geojson_aggregate, + mock_reverse_geolocation, + mock_update_bounding_box, + mock_parse_request_parameters, + ): + from reverse_geolocation_processor import reverse_geolocation_process + + mock_get_execution_id.return_value = "test_execution_id" + mock_check_maximum_executions.return_value = None + # Mocking the parsed request parameters + mock_parse_request_parameters.return_value = ( + pd.DataFrame({"stop_lat": [1.0], "stop_lon": [1.0]}), + "test_stable_id", + "test_dataset_id", + "gtfs", + ["test_url"], + True, + "per-point", + False, + 1, + ) + mock_update_bounding_box.return_value = MagicMock() + mock_reverse_geolocation.return_value = {"group_id": MagicMock()} + + # Mocking a Flask request + request = MagicMock(spec=Request) + + # Call the function + response, status_code = reverse_geolocation_process(request) + + # Assertions + self.assertEqual(status_code, 200) + self.assertIn("Processed 1 stops", response) + mock_parse_request_parameters.assert_called_once() + mock_update_bounding_box.assert_called_once() + mock_reverse_geolocation.assert_called_once() + mock_create_geojson_aggregate.assert_called_once() + + @patch("reverse_geolocation_processor.parse_request_parameters") + def test_invalid_request(self, mock_parse_request_parameters): + from reverse_geolocation_processor import reverse_geolocation_process + + # Mocking parse_request_parameters to raise a ValueError + mock_parse_request_parameters.side_effect = ValueError("Invalid request") + + # Mocking a Flask request + request = MagicMock(spec=Request) + + # Call the function + response, status_code = reverse_geolocation_process(request) + + # Assertions + self.assertEqual(status_code, ERROR_STATUS_CODE) + self.assertEqual(response, "Invalid request") + mock_parse_request_parameters.assert_called_once() + + @patch("reverse_geolocation_processor.parse_request_parameters") + @patch("reverse_geolocation_processor.update_dataset_bounding_box") + @patch("reverse_geolocation_processor.reverse_geolocation") + @patch("reverse_geolocation_processor.check_maximum_executions") + @patch("reverse_geolocation_processor.get_execution_id") + @patch("reverse_geolocation_processor.record_execution_trace") + def test_exception_handling( + self, + _, + mock_check_get_execution_id, + mock_check_maximum_executions, + mock_reverse_geolocation, + mock_update_bounding_box, + mock_parse_request_parameters, + ): + from reverse_geolocation_processor import reverse_geolocation_process + + mock_check_get_execution_id.return_value = "test_execution_id" + mock_check_maximum_executions.return_value = None + # mock_dataset_service.get_by_execution_and_stable_ids.return_value = 0 + # Mocking the parsed request parameters + mock_parse_request_parameters.return_value = ( + pd.DataFrame({"stop_lat": [1.0], "stop_lon": [1.0]}), + "test_stable_id", + "test_dataset_id", + "gtfs", + ["test_url"], + True, + "per-point", + False, + 1, + ) + mock_update_bounding_box.side_effect = Exception("Unexpected error") + + # Mocking a Flask request + request = MagicMock(spec=Request) + + # Call the function + response, status_code = reverse_geolocation_process(request) + + # Assertions + self.assertEqual(status_code, ERROR_STATUS_CODE) + self.assertIn("Unexpected error", response) + mock_parse_request_parameters.assert_called_once() + mock_update_bounding_box.assert_called_once() + mock_reverse_geolocation.assert_not_called() + + @patch("reverse_geolocation_processor.parse_request_parameters") + @patch("reverse_geolocation_processor.check_maximum_executions") + @patch("reverse_geolocation_processor.get_execution_id") + @patch("reverse_geolocation_processor.record_execution_trace") + def test_valid_request_empty_stops( + self, + _, + mock_get_execution_id, + mock_check_maximum_executions, + mock_parse_request_parameters, + ): + from reverse_geolocation_processor import reverse_geolocation_process + + mock_get_execution_id.return_value = "test_execution_id" + mock_check_maximum_executions.return_value = None + # Mocking the parsed request parameters + mock_parse_request_parameters.return_value = ( + pd.DataFrame({"stop_lat": [], "stop_lon": []}), + "test_stable_id", + "test_dataset_id", + "gtfs", + ["test_url"], + True, + "per-point", + False, + 1, + ) + + # Mocking a Flask request + request = MagicMock(spec=Request) + + # Call the function + response, status_code = reverse_geolocation_process(request) + + # Assertions + self.assertEqual(status_code, ERROR_STATUS_CODE) + self.assertIn("No stops found in the feed", response) + mock_parse_request_parameters.assert_called_once() + + @patch("reverse_geolocation_processor.get_geopolygons_with_geometry") + @patch("reverse_geolocation_processor.extract_location_aggregates_per_point") + @patch("reverse_geolocation_processor.update_feed_location") + @patch("reverse_geolocation_processor.load_feed") + def test_valid_per_point_strategy( + self, + mock_load_feed, + mock_update_feed_location, + mock_extract_per_point, + mock_get_geopolygons, + ): + from reverse_geolocation_processor import reverse_geolocation + + # Mock feed and database session + mock_feed_instance = MagicMock() + mock_load_feed.return_value = mock_feed_instance + + # Mock geopolygons + mock_get_geopolygons.return_value = ( + {}, + pd.DataFrame({"stop_lat": [1.0], "stop_lon": [1.0]}), + ) + + # Call the function + result = reverse_geolocation( + strategy=ReverseGeocodingStrategy.PER_POINT, + stable_id="test_stable_id", + stops_df=pd.DataFrame({"stop_lat": [1.0], "stop_lon": [1.0]}), + logger=MagicMock(), + use_cache=True, + db_session=MagicMock(spec=Session), + ) + + # Assertions + self.assertEqual(result, {}) + mock_get_geopolygons.assert_called_once() + mock_extract_per_point.assert_called_once() + mock_update_feed_location.assert_called_once() + + @patch("reverse_geolocation_processor.get_geopolygons_with_geometry") + @patch("reverse_geolocation_processor.update_feed_location") + @patch("reverse_geolocation_processor.load_feed") + @patch("reverse_geolocation_processor.extract_location_aggregates_per_polygon") + def test_valid_per_polygon_strategy( + self, + mock_extract_per_polygon, + mock_load_feed, + mock_update_feed_location, + mock_get_geopolygons, + ): + from reverse_geolocation_processor import reverse_geolocation + + # Mock feed and database session + mock_feed_instance = MagicMock() + mock_load_feed.query.return_value = mock_feed_instance + + # Mock geopolygons + mock_get_geopolygons.return_value = ( + {}, + pd.DataFrame({"stop_lat": [1.0], "stop_lon": [1.0]}), + ) + mock_extract_per_polygon.return_value = () + # Call the function + result = reverse_geolocation( + strategy=ReverseGeocodingStrategy.PER_POLYGON, + stable_id="test_stable_id", + stops_df=pd.DataFrame({"stop_lat": [1.0], "stop_lon": [1.0]}), + logger=MagicMock(), + use_cache=True, + db_session=MagicMock(spec=Session), + ) + + # Assertions + self.assertEqual(result, {}) + mock_get_geopolygons.assert_called_once() + mock_update_feed_location.assert_called_once() + mock_extract_per_polygon.assert_called_once() + @with_db_session(db_url=default_db_url) - @patch("reverse_geolocation_processor.create_refresh_materialized_view_task") - @patch("reverse_geolocation_processor.extract_location_aggregate") - def test_extract_location_aggregates_per_point( + def test_invalid_strategy( self, - mock_extract_location_aggregate, - mock_create_refresh_materialized_view_task, - db_session, + db_session: Session, ): - from reverse_geolocation_processor import extract_location_aggregates_per_point + from reverse_geolocation_processor import reverse_geolocation clean_testing_db() - - # Create sample feed - feed_id = faker.uuid4(cast_to=str) - feed = Gtfsfeed(id=feed_id, stable_id=faker.uuid4(cast_to=str)) - gtfs_rt_feed = Gtfsrealtimefeed( - id=faker.uuid4(cast_to=str), stable_id=faker.uuid4(cast_to=str) + feed = Gtfsfeed( + id="test_feed", + stable_id="test_feed", + status="active", + gtfsdatasets=[], + locations=[], ) - feed.gtfs_rt_feeds = [gtfs_rt_feed] db_session.add(feed) db_session.commit() - # Prepare stops DataFrame - stops_df = pd.DataFrame( - { - "stop_id": [1, 2, 3], - "stop_lat": [2.0, 3.0, 10.0], # Two inside polygon, one unmatched - "stop_lon": [2.0, 3.0, 10.0], - } + # Call the function + result = reverse_geolocation( + strategy="invalid_strategy", + stable_id="test_feed", + stops_df=pd.DataFrame({"stop_lat": [1.0], "stop_lon": [1.0]}), + logger=MagicMock(), + use_cache=True, + db_session=db_session, ) - - stops_df["geometry"] = stops_df.apply( - lambda x: WKTElement(f"POINT ({x['stop_lon']} {x['stop_lat']})", srid=4326), - axis=1, + db_session.commit() + # Assertions + self.assertEqual( + ("Invalid strategy: invalid_strategy", ERROR_STATUS_CODE), result ) + clean_testing_db() - # Prepare mock GeopolygonAggregate for matched stops - geopolygon = Geopolygon( - osm_id=faker.random_int(), - name=faker.city(), - admin_level=3, - geometry=WKTElement("POLYGON((0 0, 5 0, 5 5, 0 5, 0 0))", srid=4326), - iso_3166_1_code=faker.country_code(), - ) + @with_db_session(db_url=default_db_url) + def test_load_feed_missing_feed(self, db_session): + from reverse_geolocation_processor import reverse_geolocation + + # Call the function and expect a ValueError + with self.assertRaises(ValueError) as context: + reverse_geolocation( + strategy=ReverseGeocodingStrategy.PER_POINT, + stable_id="missing_stable_id", + stops_df=pd.DataFrame({"stop_lat": [1.0], "stop_lon": [1.0]}), + logger=MagicMock(), + use_cache=True, + db_session=db_session, + ) + + self.assertIn("No feed found for stable ID", str(context.exception)) + + @with_db_session(db_url=default_db_url) + def test_update_feed_location_success(self, db_session): + from reverse_geolocation_processor import update_feed_location + clean_testing_db() + feed = Gtfsfeed( + id="test_feed", + data_type="gtfs", + stable_id="test_feed", + status="active", + gtfsdatasets=[], + locations=[], + gtfs_rt_feeds=[ + Gtfsrealtimefeed( + id="test_rt_feed", + ) + ], + ) group = Osmlocationgroup( - group_id=faker.uuid4(cast_to=str), - group_name=geopolygon.name, - osms=[geopolygon], + group_id="group_1", + group_name="group_name", + osms=[ + Geopolygon( + osm_id=1, + admin_level=2, + geometry=WKTElement( + "POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326 + ), + iso_3166_1_code="CA", + name="Canada", + ), + Geopolygon( + osm_id=2, + admin_level=4, + geometry=WKTElement( + "POLYGON((0 0, 1 0, 1 1, 0 1, 0 0))", srid=4326 + ), + iso_3166_2_code="CA-QC", + name="Quebec", + ), + ], ) - - mock_aggregate = GeopolygonAggregate(group, stops_count=1) + db_session.add(feed) db_session.add(group) db_session.commit() - # Mock extract_location_aggregate behavior - def side_effect(_, stop_geometry, __, ___): - if stop_geometry.data == "POINT (10.0 10.0)": # Simulate unmatched stop - return None - return mock_aggregate - - mock_extract_location_aggregate.side_effect = side_effect - - # Prepare location_aggregates dict (empty initially) - location_aggregates = {} + location_group = GeopolygonAggregate(group, stops_count=1) + cache_location_groups = {"group_1": location_group} # Call the function - extract_location_aggregates_per_point( - feed_id, stops_df, location_aggregates, logger, db_session=db_session + update_feed_location( + cache_location_groups=cache_location_groups, + feed=feed, + logger=logger, + db_session=db_session, ) - # Assertions - # Ensure only matched stops are aggregated - self.assertEqual(len(location_aggregates), 1) - first_aggregate = list(location_aggregates.values())[0] - self.assertIsInstance(first_aggregate, GeopolygonAggregate) - self.assertEqual(first_aggregate.stop_count, 2) # Two matched stops - - # Verify materialized view was refreshed - mock_create_refresh_materialized_view_task.assert_called_once() - db_session.close_all() + self.assertEqual(1, len(feed.locations)) + self.assertEqual("CA", feed.locations[0].country_code) + self.assertEqual(1, len(feed.gtfs_rt_feeds[0].locations)) + self.assertEqual("CA", feed.gtfs_rt_feeds[0].locations[0].country_code) + clean_testing_db() diff --git a/functions-python/reverse_geolocation_populate/src/scripts/populate_local_db_countries.py b/functions-python/reverse_geolocation_populate/src/scripts/populate_local_db_countries.py index 53c75198b..c75adf2ff 100644 --- a/functions-python/reverse_geolocation_populate/src/scripts/populate_local_db_countries.py +++ b/functions-python/reverse_geolocation_populate/src/scripts/populate_local_db_countries.py @@ -23,7 +23,7 @@ SCRIPT_DIR = Path(__file__).resolve().parent JSON_PATH = SCRIPT_DIR / iso_3166_1_codes_file -country_filter = ["US", "CA", "GB", "FR", "DE", "IT", "ES", "JP", "IN"] +country_filter = ["US", "CA", "GB", "FR", "DE", "IT", "ES", "JP", "IN", "NZ"] logging.basicConfig(level=logging.INFO) diff --git a/infra/functions-python/main.tf b/infra/functions-python/main.tf index 0a77e8fb5..99b83d42a 100644 --- a/infra/functions-python/main.tf +++ b/infra/functions-python/main.tf @@ -1004,6 +1004,11 @@ resource "google_cloudfunctions2_function" "reverse_geolocation_processor" { service_config { environment_variables = { PYTHONNODEBUGRANGES = 0 + PROJECT_ID = var.project_id + GCP_REGION = var.gcp_region + ENVIRONMENT = var.environment + MATERIALIZED_VIEW_QUEUE = google_cloud_tasks_queue.refresh_materialized_view_task_queue.name + SERVICE_ACCOUNT_EMAIL = google_service_account.functions_service_account.email DATASETS_BUCKET_NAME_GTFS = "${var.datasets_bucket_name}-${var.environment}" DATASETS_BUCKET_NAME_GBFS = "${var.gbfs_bucket_name}-${var.environment}" } @@ -1229,7 +1234,7 @@ resource "google_cloudfunctions2_function" "tasks_executor" { service_config { environment_variables = { PROJECT_ID = var.project_id - ENV = var.environment + ENVIRONMENT = var.environment BOUNDING_BOXES_PUBSUB_TOPIC_NAME = google_pubsub_topic.rebuild_missing_bounding_boxes.name DATASET_PROCESSING_TOPIC_NAME = "datasets-batch-topic-${var.environment}" MATERIALIZED_VIEW_QUEUE = google_cloud_tasks_queue.refresh_materialized_view_task_queue.name