feat: update feed geolocation information (#1361)

Author: davidgamez

api/src/shared/common/gcp_utils.py

@@ -1,15 +1,15 @@
 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():
+    from google.cloud import tasks_v2
+
     """
     Asynchronously refresh a materialized view.
     Ensures deduplication by generating a unique task name.
@@ -70,18 +70,19 @@ def create_refresh_materialized_view_task():
 
 
 def create_http_task_with_name(
-    client: "tasks_v2.CloudTasksClient",
+    client: any,  # tasks_v2.CloudTasksClient
     body: bytes,
     url: str,
     project_id: str,
     gcp_region: str,
     queue_name: str,
     task_name: str,
-    task_time: Timestamp,
-    http_method: "tasks_v2.HttpMethod",
+    task_time,
+    http_method: any,  # tasks_v2.HttpMethod
     timeout_s: int = 1800,  # 30 minutes
 ):
     """Creates a GCP Cloud Task."""
+    from google.cloud import tasks_v2
     from google.protobuf import duration_pb2
 
     token = tasks_v2.OidcToken(service_account_email=os.getenv("SERVICE_ACCOUNT_EMAIL"))

api/src/shared/database/database.py

@@ -5,7 +5,7 @@
 import uuid
 from typing import Type, Callable
 from dotenv import load_dotenv
-from sqlalchemy import create_engine, text, event
+from sqlalchemy import create_engine, text, event, func, select
 from sqlalchemy.orm import load_only, Query, class_mapper, Session, mapper
 from shared.database_gen.sqlacodegen_models import (
     Base,
@@ -33,6 +33,11 @@ def generate_unique_id() -> str:
     return str(uuid.uuid4())
 
 
+def get_db_timestamp(db_session: Session) -> func.current_timestamp():
+    """Get the current time from the database."""
+    return db_session.execute(select(func.current_timestamp())).scalar()
+
+
 def configure_polymorphic_mappers():
     """
     Configure the polymorphic mappers allowing polymorphic values on relationships.

functions-python/reverse_geolocation/src/reverse_geolocation_processor.py

@@ -26,13 +26,14 @@
 )
 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.database import with_db_session, get_db_timestamp
 from shared.database_gen.sqlacodegen_models import (
     Feed,
     Feedlocationgrouppoint,
     Osmlocationgroup,
     Gtfsdataset,
     Gtfsfeed,
+    Gbfsfeed,
 )
 from shared.dataset_service.dataset_service_commons import Status
 
@@ -134,15 +135,18 @@ def clean_stop_cache(db_session, feed, geometries_to_delete, logger):
     db_session.commit()
 
 
+@with_db_session
 def create_geojson_aggregate(
     location_groups: List[GeopolygonAggregate],
     total_stops: int,
-    stable_id: str,
     bounding_box: shapely.Polygon,
     data_type: str,
     logger,
+    feed: Gtfsfeed | Gbfsfeed,
+    gtfs_dataset: Gtfsdataset = None,
     extraction_urls: List[str] = None,
     public: bool = True,
+    db_session: Session = None,
 ) -> None:
     """Create a GeoJSON file with the aggregated locations. This file will be uploaded to GCS and used for
     visualization."""
@@ -197,10 +201,13 @@ def create_geojson_aggregate(
     else:
         raise ValueError("The data type must be either 'gtfs' or 'gbfs'.")
     bucket = storage_client.bucket(bucket_name)
-    blob = bucket.blob(f"{stable_id}/geolocation.geojson")
+    blob = bucket.blob(f"{feed.stable_id}/geolocation.geojson")
     blob.upload_from_string(json.dumps(json_data))
     if public:
         blob.make_public()
+    feed.geolocation_file_created_date = get_db_timestamp(db_session)
+    if gtfs_dataset:
+        feed.geolocation_file_dataset = gtfs_dataset
     logger.info("GeoJSON data saved to %s", blob.name)
 
 
@@ -210,10 +217,9 @@ def get_storage_client():
     return storage.Client()
 
 
-@with_db_session
 @track_metrics(metrics=("time", "memory", "cpu"))
 def update_dataset_bounding_box(
-    dataset_id: str, stops_df: pd.DataFrame, db_session: Session
+    gtfs_dataset: Gtfsdataset, stops_df: pd.DataFrame, db_session: Session
 ) -> shapely.Polygon:
     """
     Update the bounding box of the dataset using the stops DataFrame.
@@ -231,19 +237,12 @@ def update_dataset_bounding_box(
         f")",
         srid=4326,
     )
-    if not dataset_id:
-        return to_shape(bounding_box)
-    gtfs_dataset = (
-        db_session.query(Gtfsdataset)
-        .filter(Gtfsdataset.stable_id == dataset_id)
-        .one_or_none()
-    )
     if not gtfs_dataset:
-        raise ValueError(f"Dataset {dataset_id} does not exist in the database.")
+        return to_shape(bounding_box)
     gtfs_feed = db_session.get(Gtfsfeed, gtfs_dataset.feed_id)
     if not gtfs_feed:
         raise ValueError(
-            f"GTFS feed for dataset {dataset_id} does not exist in the database."
+            f"GTFS feed for dataset {gtfs_dataset.stable_id} does not exist in the database."
         )
     gtfs_feed.bounding_box = bounding_box
     gtfs_feed.bounding_box_dataset = gtfs_dataset
@@ -252,8 +251,22 @@ def update_dataset_bounding_box(
     return to_shape(bounding_box)
 
 
+def load_dataset(dataset_id: str, db_session: Session) -> Gtfsdataset:
+    gtfs_dataset = (
+        db_session.query(Gtfsdataset)
+        .filter(Gtfsdataset.stable_id == dataset_id)
+        .one_or_none()
+    )
+    if not gtfs_dataset:
+        raise ValueError(
+            f"Dataset with ID {dataset_id} does not exist in the database."
+        )
+    return gtfs_dataset
+
+
+@with_db_session()
 def reverse_geolocation_process(
-    request: flask.Request,
+    request: flask.Request, db_session: Session = None
 ) -> Tuple[str, int] | Tuple[Dict, int]:
     """
     Main function to handle reverse geolocation processing.
@@ -331,14 +344,21 @@ def reverse_geolocation_process(
 
     try:
         # Update the bounding box of the dataset
-        bounding_box = update_dataset_bounding_box(dataset_id, stops_df)
+        gtfs_dataset: Gtfsdataset = None
+        if dataset_id:
+            gtfs_dataset = load_dataset(dataset_id, db_session)
+        feed = load_feed(stable_id, data_type, logger, db_session)
+
+        bounding_box = update_dataset_bounding_box(gtfs_dataset, stops_df, db_session)
 
         location_groups = reverse_geolocation(
             strategy=strategy,
             stable_id=stable_id,
             stops_df=stops_df,
+            data_type=data_type,
             logger=logger,
             use_cache=use_cache,
+            db_session=db_session,
         )
 
         if not location_groups:
@@ -358,14 +378,19 @@ def reverse_geolocation_process(
         create_geojson_aggregate(
             list(location_groups.values()),
             total_stops=total_stops,
-            stable_id=stable_id,
             bounding_box=bounding_box,
             data_type=data_type,
             extraction_urls=extraction_urls,
             logger=logger,
             public=public,
+            feed=feed,
+            gtfs_dataset=gtfs_dataset,
+            db_session=db_session,
         )
 
+        # Commit the changes to the database
+        db_session.commit()
+        create_refresh_materialized_view_task()
         logger.info(
             "COMPLETED. Processed %s stops for stable ID %s with strategy. "
             "Retrieved %s locations.",
@@ -408,6 +433,7 @@ def reverse_geolocation(
     strategy,
     stable_id,
     stops_df,
+    data_type,
     logger,
     use_cache,
     db_session: Session = None,
@@ -417,7 +443,7 @@ def reverse_geolocation(
     """
     logger.info("Processing geopolygons with strategy: %s.", strategy)
 
-    feed = load_feed(stable_id, logger, db_session)
+    feed = load_feed(stable_id, data_type, logger, db_session)
 
     # Get Geopolygons with Geometry and cached location groups
     cache_location_groups, unmatched_stops_df = get_geopolygons_with_geometry(
@@ -453,13 +479,13 @@ def reverse_geolocation(
         logger=logger,
         db_session=db_session,
     )
-    create_refresh_materialized_view_task()
     return cache_location_groups
 
 
-def load_feed(stable_id, logger, db_session):
+def load_feed(stable_id, data_type, logger, db_session) -> Gtfsfeed | Gbfsfeed:
+    """Load feed from the database using the stable ID and data type."""
     feed = (
-        db_session.query(Feed)
+        db_session.query(Gbfsfeed if data_type == "gbfs" else Gtfsfeed)
         .options(joinedload(Feed.feedlocationgrouppoints))
         .filter(Feed.stable_id == stable_id)
         .one_or_none()
@@ -508,5 +534,3 @@ def update_feed_location(
                 gtfs_rt_feed.locations = feed_locations
     if feed_locations:
         feed.locations = feed_locations
-    # Commit the changes to the database
-    db_session.commit()