add strategy decision for extracting location from feeds

Author: davidgamez

functions-python/helpers/locations.py

@@ -1,10 +1,19 @@
+from enum import Enum
 from typing import Dict, Optional
 from sqlalchemy.orm import Session
 import pycountry
 from shared.database_gen.sqlacodegen_models import Feed, Location
 import logging
 
 
+class ReverseGeocodingStrategy(str, Enum):
+    """
+    Enum for reverse geocoding strategies.
+    """
+
+    PER_POINT = "per-point"
+
+
 def get_country_code(country_name: str) -> Optional[str]:
     """
     Get ISO 3166 country code from country name

functions-python/helpers/runtime_metrics.py

@@ -0,0 +1,70 @@
+import time
+import tracemalloc
+import psutil
+import logging
+
+
+def track_metrics(metrics=("time", "memory", "cpu"), logger=None):
+    """Decorator to track specified metrics (time, memory, cpu) during function execution.
+    The decorator logs the metrics using the provided logger or a default logger if none is provided.
+    Args:
+        metrics (tuple): Metrics to track. Options are "time", "memory", "cpu
+        logger (logging.Logger): Logger instance to log the metrics. If None, uses a default logger.
+    Usage:
+        @track_metrics(metrics=("time", "memory", "cpu"), logger=dynamic_logger)
+        def example_function():
+            data = [i for i in range(10**6)]  # Simulate work
+            time.sleep(1)  # Simulate delay
+            return sum(data)
+    """
+
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            if not logger:
+                # Use a default logger if none is provided
+                logger_instance = logging.getLogger(func.__name__)
+                # handler = logging.StreamHandler()
+                # formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+                # handler.setFormatter(formatter)
+                # logger_instance.addHandler(handler)
+                # logger_instance.setLevel(logging.DEBUG)
+            else:
+                logger_instance = logger
+
+            process = psutil.Process()
+            tracemalloc.start() if "memory" in metrics else None
+            start_time = time.time() if "time" in metrics else None
+            cpu_before = (
+                process.cpu_percent(interval=None) if "cpu" in metrics else None
+            )
+
+            try:
+                result = func(*args, **kwargs)
+            except Exception as e:
+                logger_instance.error(
+                    f"Function '{func.__name__}' raised an exception: {e}"
+                )
+                raise
+            finally:
+                if "time" in metrics:
+                    duration = time.time() - start_time
+                    logger_instance.debug(
+                        f"Function '{func.__name__}' executed in {duration:.2f} seconds."
+                    )
+                if "memory" in metrics:
+                    current, peak = tracemalloc.get_traced_memory()
+                    tracemalloc.stop()
+                    logger_instance.debug(
+                        f"Function '{func.__name__}' peak memory usage: {peak / (1024 ** 2):.2f} MB."
+                    )
+                if "cpu" in metrics:
+                    cpu_after = process.cpu_percent(interval=None)
+                    logger_instance.debug(
+                        f"Function '{func.__name__}' CPU usage: {cpu_after - cpu_before:.2f}%."
+                    )
+
+            return result
+
+        return wrapper
+
+    return decorator

functions-python/helpers/transform.py

@@ -16,15 +16,17 @@
 from typing import List, Optional
 
 
-def to_boolean(value):
+def to_boolean(value, default_value: Optional[bool] = False) -> bool:
     """
     Convert a value to a boolean.
     """
     if isinstance(value, bool):
         return value
+    if isinstance(value, (int, float)):
+        return value != 0
     if isinstance(value, str):
-        return value.lower() in ["true", "1", "yes", "y"]
-    return False
+        return value.strip().lower() in {"true", "1", "yes", "y"}
+    return default_value
 
 
 def get_nested_value(
@@ -53,3 +55,23 @@ def get_nested_value(
         result = current_data.strip()
         return result if result else default_value
     return current_data
+
+
+def to_enum(value, enum_class=None, default_value=None):
+    """
+    Convert a value to an enum member of the specified enum class.
+
+    Args:
+        value: The value to convert.
+        enum_class: The enum class to convert the value to.
+        default_value: The default value to return if conversion fails.
+
+    Returns:
+        An enum member if conversion is successful, otherwise the default value.
+    """
+    if enum_class and isinstance(value, enum_class):
+        return value
+    try:
+        return enum_class(str(value))
+    except (ValueError, TypeError):
+        return default_value

functions-python/reverse_geolocation/.coveragerc

@@ -4,6 +4,7 @@ omit =
     */helpers/*
     */database_gen/*
     */shared/*
+    */scripts/*
 
 [report]
 exclude_lines =

functions-python/reverse_geolocation/README.md

@@ -62,6 +62,8 @@ This function performs the core reverse geolocation logic. It processes location
 - `vehicle_status_url`: Required if `data_type` is `gbfs` and `station_information_url` and `free_bike_status_url` are omitted. URL of the GBFS `vehicle_status.json` file.
 - `free_bike_status_url`: Required if `data_type` is `gbfs` and `station_information_url` and `vehicle_status_url` are omitted. URL of the GBFS `free_bike_status.json` file.
 - `data_type`: Optional. Specifies the type of data being processed. Can be `gtfs` or `gbfs`. If not provided, the function will attempt to determine the type based on the URLs provided.
+- `strategy`: Optional. Specifies the reverse geolocation strategy to use. Defaults to `per-point`.
+- `public`: Optional. Indicates whether the resulting geojson files will be public or private. Defaults to `true`.
 
 ### Processing Steps:
 

functions-python/reverse_geolocation/requirements.txt

@@ -29,6 +29,7 @@ shapely
 gtfs-kit
 matplotlib
 jsonpath_ng
+psutil
 
 # Configuration
 python-dotenv==1.0.0

functions-python/reverse_geolocation/requirements_dev.txt

@@ -1,4 +1,6 @@
 Faker
 pytest~=7.4.3
 urllib3-mock
-requests-mock
+requests-mock
+gcp-storage-emulator
+geopandas

functions-python/reverse_geolocation/src/parse_request.py

@@ -7,6 +7,9 @@
 import requests
 from jsonpath_ng import parse
 
+from shared.helpers.locations import ReverseGeocodingStrategy
+from shared.helpers.transform import to_boolean, to_enum
+
 
 def parse_request_parameters(
     request: flask.Request,
@@ -45,7 +48,19 @@ def parse_request_parameters(
         raise ValueError(
             f"Invalid data_type '{data_type}'. Supported types are 'gtfs' and 'gbfs'."
         )
-    return df, stable_id, dataset_id, data_type, urls
+    public = True
+    if "public" in request_json:
+        public = to_boolean(request_json["public"], default_value=True)
+    strategy = ReverseGeocodingStrategy.PER_POINT
+    if "strategy" in request_json:
+        strategy = to_enum(
+            value=request_json["strategy"],
+            default_value=ReverseGeocodingStrategy.PER_POINT,
+        )
+    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
 
 
 def parse_request_parameters_gtfs(

functions-python/reverse_geolocation/src/reverse_geolocation.py

@@ -41,6 +41,13 @@ def parse_resource_data(data: dict) -> Tuple[str, str, str]:
 def reverse_geolocation_pubsub(request: CloudEvent) -> None:
     """
     Reverse geolocation function triggered by a Pub/Sub message.
+    @:request: CloudEvent containing the Pub/Sub message data. Example data:
+    {
+        "stable_id": "example_stable_id",
+        "dataset_id": "example_dataset_id",
+        "url": "https://example.com/path/to/feed.zip"
+    }
+
     """
     try:
         init(request)

functions-python/reverse_geolocation/src/reverse_geolocation_processor.py

@@ -15,6 +15,7 @@
 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,
@@ -37,6 +38,7 @@
     Gtfsfeed,
 )
 from shared.helpers.logger import get_logger
+from shared.helpers.runtime_metrics import track_metrics
 
 
 @with_db_session
@@ -231,6 +233,7 @@ def create_geojson_aggregate(
     data_type: str,
     logger,
     extraction_urls: List[str] = None,
+    public: bool = True,
 ) -> None:
     """Create a GeoJSON file with the aggregated locations. This file will be uploaded to GCS and used for
     visualization."""
@@ -286,7 +289,8 @@ def create_geojson_aggregate(
     bucket = storage_client.bucket(bucket_name)
     blob = bucket.blob(f"{stable_id}/geolocation.geojson")
     blob.upload_from_string(json.dumps(json_data))
-    blob.make_public()
+    if public:
+        blob.make_public()
     logger.info("GeoJSON data saved to %s", blob.name)
 
 
@@ -316,7 +320,7 @@ def get_or_create_location(
 
 
 @with_db_session
-def extract_location_aggregates(
+def extract_location_aggregates_per_point(
     feed_id: str,
     stops_df: pd.DataFrame,
     location_aggregates: Dict[str, GeopolygonAggregate],
@@ -381,6 +385,9 @@ def extract_location_aggregates(
 
 
 @with_db_session
+@track_metrics(
+    metrics=("time", "memory", "cpu"), logger=get_logger(__name__, "stable_id")
+)
 def update_dataset_bounding_box(
     dataset_id: str, stops_df: pd.DataFrame, db_session: Session
 ) -> shapely.Polygon:
@@ -418,9 +425,27 @@ def reverse_geolocation_process(
 ) -> Tuple[str, int] | Tuple[Dict, int]:
     """
     Main function to handle reverse geolocation processing.
-    """
-    overall_start = datetime.now()
+    @:request: Flask request object containing the parameters for the reverse geolocation process.
+    Example request JSON(GTFS):
+    {
+        "stable_id": "example_stable_id",
+        "dataset_id": "example_dataset_id",
+        "stops_url": "https://example.com/path/to/stops.csv",
+        "data_type": "gtfs",
+        "strategy": "per-point"
+    }
+    Example request JSON(GBFS):
+    {
+        "stable_id": "example_stable_id",
+        "dataset_id": "example_dataset_id",
+        "station_information_url": "https://example.com/path/to/station_information.json",
+        "vehicle_status_url": "https://example.com/path/to/vehicle_status.json",
+        "free_bike_status_url": "https://example.com/path/to/free_bike_status.json",
+        "strategy": "per-point"
+    }
+    @:returns: A tuple containing a message and the HTTP status code.
 
+    """
     try:
         # Parse request parameters
         (
@@ -429,6 +454,8 @@ def reverse_geolocation_process(
             dataset_id,
             data_type,
             extraction_urls,
+            public,
+            strategy,
         ) = parse_request_parameters(request)
 
         # Remove duplicate lat/lon points
@@ -447,19 +474,16 @@ def reverse_geolocation_process(
         return str(e), ERROR_STATUS_CODE
 
     logger = get_logger(__name__, stable_id)
-
     try:
         # Update the bounding box of the dataset
         bounding_box = update_dataset_bounding_box(dataset_id, stops_df)
 
-        # Get Cached Geopolygons
-        feed_id, location_groups, stops_df = get_cached_geopolygons(
-            stable_id, stops_df, logger
+        location_groups = reverse_geolocation(
+            strategy,
+            stable_id,
+            stops_df,
+            logger,
         )
-        logger.info("Number of location groups extracted: %s", len(location_groups))
-
-        # Extract Location Groups
-        extract_location_aggregates(feed_id, stops_df, location_groups, logger)
 
         # Create GeoJSON Aggregate
         create_geojson_aggregate(
@@ -470,15 +494,12 @@ def reverse_geolocation_process(
             data_type=data_type,
             extraction_urls=extraction_urls,
             logger=logger,
+            public=public,
         )
-
-        # Overall Time
-        overall_duration = (datetime.now() - overall_start).total_seconds()
-        logger.info(f"Total time taken for the process: {overall_duration:.2f} seconds")
         logger.info(
-            "COMPLETED. Processed %s stops for stable ID %s. "
+            "COMPLETED. Processed %s stops for stable ID %s with strategy. "
             "Retrieved %s locations.",
-            total_stops,
+            len(stops_df),
             stable_id,
             len(location_groups),
         )
@@ -493,3 +514,34 @@ def reverse_geolocation_process(
         logger.error("Error processing geopolygons: %s", e)
         logger.error(traceback.format_exc())  # Log full traceback
         return str(e), ERROR_STATUS_CODE
+
+
+@track_metrics(
+    metrics=("time", "memory", "cpu"), logger=get_logger(__name__, "stable_id")
+)
+def reverse_geolocation(
+    strategy,
+    stable_id,
+    stops_df,
+    logger,
+):
+    """
+    Reverse geolocation processing based on the specified strategy.
+    """
+    logger.info("Processing geopolygons with per-point strategy.")
+    # Get Cached Geopolygons
+    feed_id, location_groups, stops_df = get_cached_geopolygons(
+        stable_id, stops_df, 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, stops_df, location_groups, logger
+            )
+        case _:
+            logger.error("Invalid strategy: %s", strategy)
+            return f"Invalid strategy: {strategy}", ERROR_STATUS_CODE
+
+    return location_groups