implement gbfs geolocation adjustment

Author: davidgamez

functions-python/tasks_executor/README.md

@@ -56,6 +56,7 @@ To update the geolocation files precision:
    "task": "update_geojson_files_precision",
    "payload": {
       "dry_run": true,
+      "data_type": "gtfs",
       "precision": 5,
       "limit": 10
    }

functions-python/tasks_executor/src/tasks/geojson/README.md

@@ -18,7 +18,8 @@ The function accepts the following payload:
 {
   "dry_run": true,                // [optional] If true, do not upload or modify the database (default: true)
   "precision": 5,                 // [optional] Number of decimal places to keep in coordinates (default: 5)
-  "limit": 10                     // [optional] Limit the number of feeds to process (default: no limit)
+  "limit": 10,                    // [optional] Limit the number of feeds to process (default: no limit)
+  "data_type": "gtfs"             // [optional] Type of data to process, either "gtfs" or "gbfs" (default: "gtfs")
 }
 ```
 
@@ -27,6 +28,7 @@ The function accepts the following payload:
 ```json
 {
   "dry_run": true,
+  "data_type": "gtfs",
   "limit": 10
 }
 ```
@@ -42,9 +44,10 @@ Also updates the `geolocation_file_created_date` and `geolocation_file_dataset_i
 
 The function requires the following environment variables:
 
-| Variable               | Description                                                                  |
-| ---------------------- | ---------------------------------------------------------------------------- |
-| `DATASETS_BUCKET_NAME` | The name of the GCS bucket used to store extracted GTFS files                |
+| Variable                       | Description                                                             |
+|--------------------------------|-------------------------------------------------------------------------|
+| `DATASETS_BUCKET_NAME`         | The name of the GCS bucket used to store extracted GTFS files           | 
+| `GBFS_SNAPSHOTS_BUCKET_NAME`   | The name of the GCS bucket used to store extracted GBFS snapshots files | 
 
 ---
 

functions-python/tasks_executor/src/tasks/geojson/update_geojson_files_precision.py

@@ -18,12 +18,13 @@
 import logging
 import os
 import re
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Literal
 
 from sqlalchemy import select, func
+from sqlalchemy.orm import Session
 
 from shared.database.database import with_db_session
-from shared.database_gen.sqlacodegen_models import Gtfsfeed
+from shared.database_gen.sqlacodegen_models import Gtfsfeed, Gbfsfeed
 from shared.helpers.locations import round_geojson_coords
 from shared.helpers.runtime_metrics import track_metrics
 
@@ -48,13 +49,16 @@ def _remove_osm_ids_from_properties(props: Dict[str, Any] | None):
         props.pop(k, None)
 
 
-def query_unprocessed_feeds(limit, db_session):
+def query_unprocessed_feeds(
+    limit: int, feed_type: Literal["gtfs", "gbfs"], db_session: Session
+) -> List[Gtfsfeed] | List[Gbfsfeed]:
     """
-    Query Gtfsfeed entries that have not been processed yet, geolocation_file_created_date is null.
+    Query eed entries that have not been processed yet, geolocation_file_created_date is null.
     """
+    model: Gtfsfeed | Gbfsfeed = Gtfsfeed if feed_type == "gtfs" else Gbfsfeed
     feeds = (
-        db_session.query(Gtfsfeed)
-        .filter(Gtfsfeed.geolocation_file_created_date.is_(None))
+        db_session.query(model)
+        .filter(model.geolocation_file_created_date.is_(None))
         .limit(limit)
         .all()
     )
@@ -72,8 +76,10 @@ def _upload_file(bucket, file_path, geojson):
 
 
 @track_metrics(metrics=("time", "memory", "cpu"))
-def _update_feed_info(feed: Gtfsfeed, timestamp):
+def _update_feed_info(feed: Gtfsfeed | Gbfsfeed, timestamp):
     feed.geolocation_file_created_date = timestamp
+    if isinstance(feed, Gbfsfeed):
+        return
     # find the most recent dataset with bounding box and set the id
     if feed.gtfsdatasets and any(d.bounding_box for d in feed.gtfsdatasets):
         latest_with_bbox = max(
@@ -141,29 +147,37 @@ def update_geojson_files_precision_handler(
         from google.cloud import storage
     except Exception as e:
         raise RuntimeError("google-cloud-storage is required at runtime: %s" % e)
-    bucket_name = payload.get("bucket_name") or os.getenv("DATASETS_BUCKET_NAME")
-    if not bucket_name:
-        raise ValueError(
-            "bucket_name must be provided in payload or set in GEOJSON_BUCKET env"
-        )
 
     dry_run = payload.get("dry_run", True)
+    data_type: Literal["gtfs", "gbfs"] = payload.get("data_type", "gtfs")
     precision = int(payload.get("precision", 5))
     limit = int(payload.get("limit", None))
+    bucket_name = payload.get("bucket_name") or (
+        os.getenv("DATASETS_BUCKET_NAME")
+        if data_type == "gtfs"
+        else os.getenv("GBFS_SNAPSHOTS_BUCKET_NAME")
+    )
+    if not bucket_name:
+        raise ValueError(
+            "bucket_name must be provided in payload or set in GEOJSON_BUCKET env"
+        )
     client = storage.Client()
     bucket = client.bucket(bucket_name)
 
     errors: List[Dict[str, str]] = []
     processed = 0
 
-    feeds: [Gtfsfeed] = query_unprocessed_feeds(limit, db_session)
+    feeds: List[Gtfsfeed] | List[Gbfsfeed] = query_unprocessed_feeds(
+        limit, data_type, db_session
+    )
     logging.info("Found %s feeds", len(feeds))
     timestamp = db_session.execute(select(func.current_timestamp())).scalar()
     for feed in feeds:
         try:
             if processed % 100 == 0:
                 logging.info("Processed %s/%s", processed, len(feeds))
-                db_session.commit()
+                if not dry_run and processed > 0:
+                    db_session.commit()
             file_path = f"{feed.stable_id}/{GEOLOCATION_FILENAME}"
             file = storage.Blob(bucket=bucket, name=file_path)
             if not file.exists():
@@ -189,6 +203,7 @@ def update_geojson_files_precision_handler(
         except Exception as e:
             logging.exception("Error processing feed %s: %s", feed.stable_id, e)
             errors.append(feed.stable_id)
+    logging.info("Processed %s/%s", processed, len(feeds))
     if not dry_run and processed > 0:
         db_session.commit()
     summary = {

functions-python/tasks_executor/tests/conftest.py

@@ -13,7 +13,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 #
-
+import uuid
 from datetime import datetime, UTC, timedelta
 
 from sqlalchemy.orm import Session
@@ -22,6 +22,7 @@
 from shared.database_gen.sqlacodegen_models import (
     Gtfsfeed,
     Gtfsdataset,
+    Gbfsfeed,
 )
 from test_shared.test_utils.database_utils import clean_testing_db, default_db_url
 
@@ -47,6 +48,14 @@ def populate_database(db_session: Session | None = None):
         )
         db_session.add(feed)
         feeds.append(feed)
+    gbfs_feed = Gbfsfeed(
+        id=f"feed_{uuid.uuid4()}",
+        stable_id=f"stable_feed_gbfs_{uuid.uuid4()}",
+        data_type="gbfs",
+        status="active",
+        created_at=now,
+    )
+    db_session.add(gbfs_feed)
     db_session.flush()
 
     datasets = []

functions-python/tasks_executor/tests/tasks/geojson/test_update_geojson_files_precision.py

@@ -8,6 +8,7 @@
 from sqlalchemy.orm import Session
 
 from shared.database.database import with_db_session
+from shared.database_gen.sqlacodegen_models import Gbfsfeed
 from shared.helpers.src.shared.database_gen.sqlacodegen_models import Gtfsfeed
 from tasks.geojson.update_geojson_files_precision import (
     process_geojson,
@@ -137,7 +138,7 @@ def test_process_geojson_single_feature_and_list_variants(self):
         )
 
     @with_db_session(db_url=default_db_url)
-    def test_handler_uploads_and_updates_feed_info(self, db_session: Session):
+    def test_handler_uploads_and_updates_gtfs_feed_info(self, db_session: Session):
         geo = {
             "type": "FeatureCollection",
             "features": [
@@ -209,6 +210,81 @@ def test_handler_uploads_and_updates_feed_info(self, db_session: Session):
         self.assertIsNotNone(reloaded_testing_feed.geolocation_file_dataset_id)
         self.assertIsNotNone(reloaded_testing_feed.geolocation_file_created_date)
 
+    @with_db_session(db_url=default_db_url)
+    def test_handler_uploads_and_updates_gbfs_feed_info(self, db_session: Session):
+        geo = {
+            "type": "FeatureCollection",
+            "features": [
+                {
+                    "type": "Feature",
+                    "geometry": {
+                        "type": "Point",
+                        "coordinates": [100.1234567, 0.9876543],
+                    },
+                    "properties": {"id": "node/1", "keep": "x"},
+                }
+            ],
+        }
+        testing_gbfs_feed = db_session.query(Gbfsfeed).limit(1).first()
+        self.assertIsNotNone(testing_gbfs_feed)
+        feed_stable_id = testing_gbfs_feed.stable_id
+        blob_name = f"{feed_stable_id}/{GEOLOCATION_FILENAME}"
+
+        fake_bucket = FakeBucket(initial_blobs={blob_name: json.dumps(geo)})
+        fake_storage = FakeStorageModule(fake_bucket, blob_exists=True)
+
+        # create module objects for google and google.cloud and inject via sys.modules
+        cloud_mod = types.ModuleType("google.cloud")
+        # 'from google.cloud import storage' in handler will bind 'storage' to this attribute
+        cloud_mod.storage = fake_storage
+        google_mod = types.ModuleType("google")
+        google_mod.cloud = cloud_mod
+
+        payload = {
+            "bucket_name": "any-bucket",
+            "dry_run": False,
+            "data_type": "gbfs",
+            "precision": 5,
+            "limit": 1,
+        }
+
+        # Inject modules into sys.modules for the duration of the handler call
+        with patch.dict(sys.modules, {"google.cloud": cloud_mod, "google": google_mod}):
+            # call wrapped handler to provide fake db_session
+            result = update_geojson_files_precision_handler(
+                payload, db_session=db_session
+            )
+
+        # verify upload happened
+        self.assertIn(blob_name, fake_bucket.uploaded)
+        uploaded_text = fake_bucket.uploaded[blob_name]
+        uploaded_geo = json.loads(uploaded_text)
+        coords = uploaded_geo.get("features")[0]["geometry"]["coordinates"]
+        self.assertEqual(coords, [round(100.1234567, 5), round(0.9876543, 5)])
+
+        self.assertEqual(
+            {
+                "total_processed_files": 1,
+                "errors": [],
+                "not_found_file": 0,
+                "params": {
+                    "dry_run": False,
+                    "precision": 5,
+                    "limit": 1,
+                },
+            },
+            result,
+        )
+        # feed updated
+        reloaded_testing_feed = (
+            db_session.query(Gbfsfeed)
+            .filter(Gbfsfeed.id.__eq__(testing_gbfs_feed.id))
+            .limit(1)
+            .first()
+        )
+        self.assertIsNone(reloaded_testing_feed.geolocation_file_dataset_id)
+        self.assertIsNotNone(reloaded_testing_feed.geolocation_file_created_date)
+
     @with_db_session(db_url=default_db_url)
     def test_handler_file_dont_exists(self, db_session: Session):
         fake_bucket = FakeBucket(initial_blobs={})

infra/functions-python/main.tf

@@ -1154,6 +1154,7 @@ resource "google_cloudfunctions2_function" "tasks_executor" {
       DATASET_PROCESSING_TOPIC_NAME = "datasets-batch-topic-${var.environment}"
       MATERIALIZED_VIEW_QUEUE = google_cloud_tasks_queue.refresh_materialized_view_task_queue.name
       DATASETS_BUCKET_NAME = "${var.datasets_bucket_name}-${var.environment}"
+      GBFS_SNAPSHOTS_BUCKET_NAME = google_storage_bucket.gbfs_snapshots_bucket.name
     }
     available_memory                 = local.function_tasks_executor_config.memory
     timeout_seconds                  = local.function_tasks_executor_config.timeout