diff --git a/api/tests/test_utils/db_utils.py b/api/tests/test_utils/db_utils.py index cbabafa35..736782b9c 100644 --- a/api/tests/test_utils/db_utils.py +++ b/api/tests/test_utils/db_utils.py @@ -192,23 +192,14 @@ def is_test_db(url): def empty_database(db, url): if is_test_db(url): - - metadata_tables = Base.metadata.tables - - # Get all table names excluding those in the excluded_tables list - all_table_names = [table_name for table_name in metadata_tables.keys() if table_name not in excluded_tables] - - # Sort the table names in reverse order of dependencies - tables_to_delete = sorted( - all_table_names, key=lambda name: len(metadata_tables[name].foreign_keys), reverse=True - ) - try: with db.start_db_session() as session: - for table_name in tables_to_delete: - table = Base.metadata.tables[table_name] - delete_stmt = delete(table) - session.execute(delete_stmt) - + # Using sorted_tables to respect foreign key constraints + for table in reversed(Base.metadata.sorted_tables): + if table.name not in excluded_tables: + table = Base.metadata.tables[table.name] + delete_stmt = delete(table) + session.execute(delete_stmt) + session.commit() except Exception as error: logging.error(f"Error while deleting from test db: {error}") diff --git a/functions-python/helpers/locations.py b/functions-python/helpers/locations.py index 0ec2d6da5..e0286b70b 100644 --- a/functions-python/helpers/locations.py +++ b/functions-python/helpers/locations.py @@ -6,7 +6,6 @@ from sqlalchemy import func, cast from geoalchemy2.types import Geography -import pycountry from shared.database_gen.sqlacodegen_models import Feed, Location, Geopolygon import logging @@ -35,6 +34,8 @@ def get_country_code(country_name: str) -> Optional[str]: Returns: Optional[str]: Two-letter ISO country code or None if not found """ + import pycountry + # Return None for empty or whitespace-only strings if not country_name or not country_name.strip(): logging.error("Could not find country code for: empty string") diff --git a/functions-python/process_validation_report/tests/conftest.py b/functions-python/process_validation_report/tests/conftest.py index 6b30c9db3..4fb9f3088 100644 --- a/functions-python/process_validation_report/tests/conftest.py +++ b/functions-python/process_validation_report/tests/conftest.py @@ -143,7 +143,7 @@ def pytest_sessionfinish(session, exitstatus): returning the exit status to the system. """ # Cleaned at the beginning instead of the end so we can examine the DB after the test. - # clean_testing_db() + clean_testing_db() def pytest_unconfigure(config): diff --git a/functions-python/process_validation_report/tests/test_validation_report.py b/functions-python/process_validation_report/tests/test_validation_report.py index 0b2d36827..8651c6d4f 100644 --- a/functions-python/process_validation_report/tests/test_validation_report.py +++ b/functions-python/process_validation_report/tests/test_validation_report.py @@ -82,6 +82,7 @@ def test_get_dataset(self, db_session): ) try: db_session.add(feed) + db_session.flush() db_session.add(dataset) db_session.flush() returned_dataset = get_dataset(dataset_stable_id, db_session) @@ -123,6 +124,7 @@ def test_create_validation_report_entities(self, mock_get, db_session): ) try: db_session.add(feed) + db_session.flush() db_session.add(dataset) db_session.commit() create_validation_report_entities(feed_stable_id, dataset_stable_id, "1.0") @@ -347,17 +349,18 @@ def test_create_validation_report_entities_missing_validator_version( ], }, ) - feed_stable_id = faker.word() - dataset_stable_id = faker.word() + feed_stable_id = faker.uuid4() + dataset_stable_id = faker.uuid4() # Create GTFS Feed - feed = Gtfsfeed(id=faker.word(), data_type="gtfs", stable_id=feed_stable_id) + feed = Gtfsfeed(id=faker.uuid4(), data_type="gtfs", stable_id=feed_stable_id) # Create a new dataset dataset = Gtfsdataset( - id=faker.word(), feed_id=feed.id, stable_id=dataset_stable_id, latest=True + id=faker.uuid4(), feed_id=feed.id, stable_id=dataset_stable_id, latest=True ) try: db_session.add(feed) + db_session.flush() db_session.add(dataset) db_session.commit() create_validation_report_entities(feed_stable_id, dataset_stable_id, "1.0") diff --git a/functions-python/tasks_executor/README.md b/functions-python/tasks_executor/README.md index 1c1057558..a4ef5613e 100644 --- a/functions-python/tasks_executor/README.md +++ b/functions-python/tasks_executor/README.md @@ -50,3 +50,15 @@ To get the list of supported tasks use: "payload": {} } ``` +To update the geolocation files precision: +```json +{ + "task": "update_geojson_files_precision", + "payload": { + "dry_run": true, + "data_type": "gtfs", + "precision": 5, + "limit": 10 + } +} +``` \ No newline at end of file diff --git a/functions-python/tasks_executor/requirements.txt b/functions-python/tasks_executor/requirements.txt index e312258c9..0dbf26655 100644 --- a/functions-python/tasks_executor/requirements.txt +++ b/functions-python/tasks_executor/requirements.txt @@ -11,6 +11,7 @@ pluggy~=1.3.0 certifi~=2025.8.3 fastapi uvicorn[standard] +psutil # SQL Alchemy and Geo Alchemy diff --git a/functions-python/tasks_executor/src/main.py b/functions-python/tasks_executor/src/main.py index 02a8a1ca9..56308c3cc 100644 --- a/functions-python/tasks_executor/src/main.py +++ b/functions-python/tasks_executor/src/main.py @@ -35,6 +35,9 @@ from tasks.visualization_files.rebuild_missing_visualization_files import ( rebuild_missing_visualization_files_handler, ) +from tasks.geojson.update_geojson_files_precision import ( + update_geojson_files_precision_handler, +) init_logger() LIST_COMMAND: Final[str] = "list" @@ -66,6 +69,10 @@ "description": "Rebuilds missing dataset files for GTFS datasets.", "handler": rebuild_missing_dataset_files_handler, }, + "update_geojson_files": { + "description": "Iterate over bucket looking for {feed_stable_id}/geolocation.geojson and update precision.", + "handler": update_geojson_files_precision_handler, + }, "rebuild_missing_visualization_files": { "description": "Rebuilds missing visualization files for GTFS datasets.", "handler": rebuild_missing_visualization_files_handler, diff --git a/functions-python/tasks_executor/src/tasks/geojson/README.md b/functions-python/tasks_executor/src/tasks/geojson/README.md new file mode 100644 index 000000000..0ed28521a --- /dev/null +++ b/functions-python/tasks_executor/src/tasks/geojson/README.md @@ -0,0 +1,58 @@ +# Update GeoJSON files + +This task adjust the GeoJSON files removing the map IDs and reducing the precision of the coordinates to 5 decimal places. ALso updates the geolocation_file_created_date and geolocation_file_dataset_id fields in the Feed table. + +--- + +## Task ID + +Use task ID: `update_geojson_files_precision` + +--- + +## Usage + +The function accepts the following payload: + +```json +{ + "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) + "data_type": "gtfs" // [optional] Type of data to process, either "gtfs" or "gbfs" (default: "gtfs") +} +``` + +### Example: + +```json +{ + "dry_run": true, + "data_type": "gtfs", + "limit": 10 +} +``` + +--- + +## What It Does + +List all feeds with GeoJSON files, download each file, remove map IDs, reduce coordinate precision to the specified number of decimal places, and re-upload the modified file. +Also updates the `geolocation_file_created_date` and `geolocation_file_dataset_id` fields in the `Feed` table. + +## GCP Environment Variables + +The function requires the following environment variables: + +| 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 | + +--- + +## Additional Notes + +* Commits to the database occur in batches of 100 feeds to improve performance and avoid large transaction blocks. +* If `dry_run` is enabled, files are uploads or DB modifications are performed. Only the number of affected feeds is logged. +* The function is safe to rerun. It will only affect feeds with missing geolocation_file_dataset_id. diff --git a/functions-python/tasks_executor/src/tasks/geojson/update_geojson_files_precision.py b/functions-python/tasks_executor/src/tasks/geojson/update_geojson_files_precision.py new file mode 100644 index 000000000..34665e799 --- /dev/null +++ b/functions-python/tasks_executor/src/tasks/geojson/update_geojson_files_precision.py @@ -0,0 +1,222 @@ +# +# MobilityData 2025 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import json +import logging +import os +import re +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, Gbfsfeed +from shared.helpers.locations import round_geojson_coords +from shared.helpers.runtime_metrics import track_metrics + +GEOLOCATION_FILENAME = "geolocation.geojson" + + +def _remove_osm_ids_from_properties(props: Dict[str, Any] | None): + """Remove OSM identifier-like properties from a feature's properties in-place.""" + if not props or not isinstance(props, dict): + return + keys_to_remove = [] + for k, v in list(props.items()): + lk = k.lower() + if "osm" in lk or lk == "@id": + keys_to_remove.append(k) + elif lk == "id": + if isinstance(v, str) and re.search( + r"\b(node|way|relation)\b|/", v, re.IGNORECASE + ): + keys_to_remove.append(k) + for k in keys_to_remove: + props.pop(k, None) + + +def query_unprocessed_feeds( + limit: int, feed_type: Literal["gtfs", "gbfs"], db_session: Session +) -> List[Gtfsfeed] | List[Gbfsfeed]: + """ + Query feed 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(model) + .filter(model.geolocation_file_created_date.is_(None)) + .limit(limit) + .all() + ) + return feeds + + +@track_metrics(metrics=("time", "memory", "cpu")) +def _upload_file(bucket, file_path, geojson): + processed_blob = bucket.blob(file_path) + processed_blob.upload_from_string( + json.dumps(geojson, ensure_ascii=False), + content_type="application/geo+json", + ) + processed_blob.make_public() + + +@track_metrics(metrics=("time", "memory", "cpu")) +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( + (d for d in feed.gtfsdatasets if d.bounding_box), + key=lambda d: d.downloaded_at or timestamp, + ) + feed.geolocation_file_dataset_id = latest_with_bbox.id + else: + logging.info( + "No GTFS datasets available with bounding box for feed %s", feed.id + ) + + +@track_metrics(metrics=("time", "memory", "cpu")) +def process_geojson(geopjson, precision): + # Normalize GeoJSON structure to FeatureCollection-like list of features + if isinstance(geopjson, dict) and geopjson.get("type") == "FeatureCollection": + features = geopjson.get("features", []) + elif isinstance(geopjson, dict) and geopjson.get("type") == "Feature": + features = [geopjson] + elif isinstance(geopjson, list): + features = geopjson + else: + # Unknown structure, skip + return + # Apply rounding via shared helper and remove osm ids + for f in features: + if not isinstance(f, dict): + continue + geom = f.get("geometry") + if geom: + # round_geojson_coords returns a new geometry object + try: + f["geometry"] = round_geojson_coords(geom, precision=precision) + except Exception as e: + logging.warning("Error processing feature %s: %s", f.get("name"), e) + return + props = f.get("properties") + _remove_osm_ids_from_properties(props) + # If original was a FeatureCollection, update it; if single Feature, keep as-is; if list, use list + if isinstance(geopjson, dict) and geopjson.get("type") == "FeatureCollection": + geopjson["features"] = features + elif isinstance(geopjson, dict) and geopjson.get("type") == "Feature": + geopjson = features[0] if features else geopjson + else: + geopjson = features + return geopjson + + +@with_db_session +def update_geojson_files_precision_handler( + payload: Dict[str, Any], db_session +) -> Dict[str, Any]: + """ + Update GeoJSON files in GCS to reduce coordinate precision and remove map ids. + + Payload keys: + - dry_run (bool) default True + - data_type (str) "gtfs" or "gbfs", default "gtfs" + - precision (int) default 5 + - bucket_name (str) optional, defaults to env var DATASETS_BUCKET_NAME or GBFS_SNAPSHOTS_BUCKET_NAME + - limit (int) + + """ + # Import GCS client at runtime to avoid dev environment import issues + try: + from google.cloud import storage + except Exception as e: + raise RuntimeError("google-cloud-storage is required at runtime: %s" % e) + + 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: 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)) + 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(): + logging.info("File does not exist: %s", file.name) + continue + logging.info("Processing file: %s", file.name) + text = file.download_as_text() + geojson = json.loads(text) + + geojson = process_geojson(geojson, precision) + if not geojson: + logging.warning("No valid GeoJSON features found in %s", file.name) + errors.append(feed.stable_id) + continue + + # Optionally upload processed geojson + if not dry_run: + _upload_file(bucket, file_path, geojson) + _update_feed_info(feed, timestamp) + logging.info("Updated feed %s", feed.stable_id) + + processed += 1 + 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 = { + "total_processed_files": processed, + "errors": errors, + "not_found_file": len(feeds) - processed - len(errors), + "params": { + "dry_run": dry_run, + "precision": precision, + "limit": limit, + }, + } + logging.info("update_geojson_files_handler result: %s", summary) + return summary diff --git a/functions-python/tasks_executor/tests/conftest.py b/functions-python/tasks_executor/tests/conftest.py index 89251bd58..70c1baede 100644 --- a/functions-python/tasks_executor/tests/conftest.py +++ b/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 @@ -46,6 +47,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 = [] @@ -56,6 +65,7 @@ def populate_database(db_session: Session | None = None): feed=feed, stable_id=f"dataset_stable_{i:04d}", downloaded_at=now - timedelta(days=i), + bounding_box="POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))", ) db_session.add(dataset) datasets.append(dataset) diff --git a/functions-python/tasks_executor/tests/tasks/geojson/test_update_geojson_files_precision.py b/functions-python/tasks_executor/tests/tasks/geojson/test_update_geojson_files_precision.py new file mode 100644 index 000000000..2478e9eba --- /dev/null +++ b/functions-python/tasks_executor/tests/tasks/geojson/test_update_geojson_files_precision.py @@ -0,0 +1,373 @@ +import os +import sys +import json +import types +import unittest +from unittest.mock import patch + +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, + update_geojson_files_precision_handler, + GEOLOCATION_FILENAME, +) +from test_shared.test_utils.database_utils import default_db_url + + +class _FakeBlobContext: + def __init__(self, bucket, name, blob_exists=True): + self.bucket = bucket + self.name = name + self.blob_exists = blob_exists + + def __enter__(self): + return self + + def exists(self): + return self.blob_exists + + def download_as_text(self): + return self.bucket.initial_blobs[self.name] + + +class _FakeUploadBlob: + def __init__(self, bucket, name): + self.bucket = bucket + self.name = name + + def upload_from_string(self, content, content_type=None): + # store as text for assertions + self.bucket.uploaded[self.name] = content + + def make_public(self): + return + + +class FakeBucket: + def __init__(self, initial_blobs=None): + # mapping name -> text + self.initial_blobs = initial_blobs or {} + self.uploaded = {} + + def blob(self, name): + # return an upload-capable blob + return _FakeUploadBlob(self, name) + + +class FakeClient: + def __init__(self, bucket): + self._bucket = bucket + + def bucket(self, name): + return self._bucket + + +class FakeStorageModule: + def __init__(self, bucket, blob_exists=True): + self._bucket = bucket + self._blob_exists = blob_exists + + def Client(self): + return FakeClient(self._bucket) + + # storage.Blob(...) used as a context manager in the handler + def Blob(self, *, bucket, name): + return _FakeBlobContext(bucket, name, self._blob_exists) + + +class TestUpdateGeojsonFilesPrecision(unittest.TestCase): + def setUp(self): + os.environ["DATASETS_BUCKET_NAME"] = "mock_bucket" + + def test_process_geojson_round_and_remove_osm_keys(self): + fc = { + "type": "FeatureCollection", + "features": [ + { + "type": "Feature", + "geometry": { + "type": "Point", + "coordinates": [12.3456789, -98.7654321], + }, + "properties": { + "name": "A", + "osm_id": 123, + "@id": "node/1", + "id": "way/2", + "keep": "yes", + }, + } + ], + } + + out = process_geojson(fc, precision=5) + self.assertIsNotNone(out) + coords = out["features"][0]["geometry"]["coordinates"] + self.assertEqual(coords, [round(12.3456789, 5), round(-98.7654321, 5)]) + + props = out["features"][0]["properties"] + self.assertNotIn("osm_id", props) + self.assertNotIn("@id", props) + self.assertNotIn("id", props) + self.assertEqual(props.get("name"), "A") + self.assertEqual(props.get("keep"), "yes") + + def test_process_geojson_single_feature_and_list_variants(self): + feat = { + "type": "Feature", + "geometry": {"type": "Point", "coordinates": [1.23456789, 2.3456789]}, + "properties": {"id": "123", "osm": "should_remove"}, + } + out1 = process_geojson(feat, precision=4) + self.assertIsInstance(out1, dict) + self.assertEqual( + out1["geometry"]["coordinates"], [round(1.23456789, 4), round(2.3456789, 4)] + ) + self.assertNotIn("osm", out1.get("properties", {})) + + lst = [feat] + out2 = process_geojson(lst, precision=3) + self.assertIsInstance(out2, list) + self.assertEqual( + out2[0]["geometry"]["coordinates"], + [round(1.23456789, 3), round(2.3456789, 3)], + ) + + @with_db_session(db_url=default_db_url) + def test_handler_uploads_and_updates_gtfs_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_feed = db_session.query(Gtfsfeed).limit(1).first() + feed_stable_id = testing_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, + "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(Gtfsfeed) + .filter(Gtfsfeed.id.__eq__(testing_feed.id)) + .limit(1) + .first() + ) + 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={}) + fake_storage = FakeStorageModule(fake_bucket, blob_exists=False) + + # 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, + "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 + ) + self.assertEqual( + { + "total_processed_files": 0, + "errors": [], + "not_found_file": 1, + "params": { + "dry_run": False, + "precision": 5, + "limit": 1, + }, + }, + result, + ) + + @with_db_session(db_url=default_db_url) + def test_handler_file_not_valid_file(self, db_session: Session): + geo = "{}" + testing_feed = db_session.query(Gtfsfeed).limit(1).first() + feed_stable_id = testing_feed.stable_id + blob_name = f"{feed_stable_id}/{GEOLOCATION_FILENAME}" + + fake_bucket = FakeBucket(initial_blobs={blob_name: 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, + "precision": 5, + "limit": 1, + } + testing_feed = db_session.query(Gtfsfeed).limit(1).first() + # 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 + ) + self.assertEqual( + { + "total_processed_files": 0, + "errors": [testing_feed.stable_id], + "not_found_file": 0, + "params": { + "dry_run": False, + "precision": 5, + "limit": 1, + }, + }, + result, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/infra/functions-python/main.tf b/infra/functions-python/main.tf index c31ce2485..f60b3fbfa 100644 --- a/infra/functions-python/main.tf +++ b/infra/functions-python/main.tf @@ -1193,6 +1193,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 PMTILES_BUILDER_QUEUE = google_cloud_tasks_queue.pmtiles_builder_task_queue.name SERVICE_ACCOUNT_EMAIL = google_service_account.functions_service_account.email GCP_REGION = var.gcp_region diff --git a/liquibase/changelog.xml b/liquibase/changelog.xml index 38eae1faa..d0f44884f 100644 --- a/liquibase/changelog.xml +++ b/liquibase/changelog.xml @@ -66,4 +66,5 @@ + diff --git a/liquibase/changes/feat_pt_152.sql b/liquibase/changes/feat_pt_152.sql new file mode 100644 index 000000000..94dad36be --- /dev/null +++ b/liquibase/changes/feat_pt_152.sql @@ -0,0 +1,9 @@ +-- Add geolocation file feed level columns +ALTER TABLE feed ADD COLUMN IF NOT EXISTS geolocation_file_created_date TIMESTAMP; +ALTER TABLE feed ADD COLUMN IF NOT EXISTS geolocation_file_dataset_id VARCHAR(255); + +ALTER TABLE feed + ADD CONSTRAINT fk_feed_geolocation_file_dataset + FOREIGN KEY (geolocation_file_dataset_id) + REFERENCES gtfsdataset(id) + ON DELETE SET NULL;