diff --git a/functions-python/tasks_executor/README.md b/functions-python/tasks_executor/README.md index 821106aa4..1c1057558 100644 --- a/functions-python/tasks_executor/README.md +++ b/functions-python/tasks_executor/README.md @@ -13,26 +13,28 @@ The function receive the following payload: } ``` -Example: +Examples: ```json { "task": "rebuild_missing_validation_reports", "payload": { - "dry_run": true, - "filter_after_in_days": 14, - "filter_statuses": ["active", "inactive", "future"] - } + "dry_run": true, + "filter_after_in_days": 14, + "filter_statuses": ["active", "inactive", "future"] + } } ``` ```json { "task": "rebuild_missing_bounding_boxes", "payload": { - "dry_run": true, - "after_date": "2025-06-01" - } + "dry_run": true, + "after_date": "2025-06-01" + } } +``` +```json { "task": "refresh_materialized_view", "payload": { @@ -44,7 +46,7 @@ Example: To get the list of supported tasks use: ```json { -"name": "list_tasks", -"payload": {} + "name": "list_tasks", + "payload": {} } ``` diff --git a/functions-python/tasks_executor/function_config.json b/functions-python/tasks_executor/function_config.json index f8b8ee09b..b8099a2d9 100644 --- a/functions-python/tasks_executor/function_config.json +++ b/functions-python/tasks_executor/function_config.json @@ -2,12 +2,16 @@ "name": "tasks_executor", "description": "The Tasks Executor function runs maintenance tasks avoiding the creation of multiple functions for one-time execution", "entry_point": "tasks_executor", - "timeout": 540, - "memory": "4Gi", + "timeout": 1000, + "memory": "8Gi", "trigger_http": true, "include_folders": ["helpers"], "include_api_folders": ["database_gen", "database", "common"], - "environment_variables": [], + "environment_variables": [ + { + "key": "DATASETS_BUCKET_NAME" + } + ], "secret_environment_variables": [ { "key": "FEEDS_DATABASE_URL" @@ -21,5 +25,5 @@ "max_instance_request_concurrency": 1, "max_instance_count": 1, "min_instance_count": 0, - "available_cpu": 1 + "available_cpu": 2 } diff --git a/functions-python/tasks_executor/requirements.txt b/functions-python/tasks_executor/requirements.txt index b12efc831..31a091512 100644 --- a/functions-python/tasks_executor/requirements.txt +++ b/functions-python/tasks_executor/requirements.txt @@ -25,4 +25,6 @@ flask google-cloud-storage # Configuration -python-dotenv==1.0.0 \ No newline at end of file +python-dotenv==1.0.0 +tippecanoe + diff --git a/functions-python/tasks_executor/src/main.py b/functions-python/tasks_executor/src/main.py index c78c8b997..a244bb880 100644 --- a/functions-python/tasks_executor/src/main.py +++ b/functions-python/tasks_executor/src/main.py @@ -31,7 +31,7 @@ from tasks.missing_bounding_boxes.rebuild_missing_bounding_boxes import ( rebuild_missing_bounding_boxes_handler, ) - +from tasks.pmtiles_builder.build_pmtiles import build_pmtiles_handler init_logger() LIST_COMMAND: Final[str] = "list" @@ -63,6 +63,10 @@ "description": "Rebuilds missing dataset files for GTFS datasets.", "handler": rebuild_missing_dataset_files_handler, }, + "build_pmtiles": { + "description": "Build pmtiles for a given dataset.", + "handler": build_pmtiles_handler, + }, } diff --git a/functions-python/tasks_executor/src/tasks/pmtiles_builder/README.md b/functions-python/tasks_executor/src/tasks/pmtiles_builder/README.md new file mode 100644 index 000000000..c9399afc4 --- /dev/null +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/README.md @@ -0,0 +1,31 @@ +# Build pmtile for a specific GTFS dataset + +This task generates the pmtiles for a provided dataset. +pmtiles are used for displaying routes and stops in the UI + +## Task ID +Use task Id: `build_pmtiles` + +## Usage +The function receive the following payload: +``` + { + "feed_stable_id": str, + "dataset_stable_id*: str + } +``` + +Example: +```json + { + "feed_stable_id": "mdb-1004", + "dataset_stable_id": "mdb-1004-202507081807" + } +``` + +The task will verify that the dataset stable id starts with the feed stable id (mdb-1004 in our example) + +# GCP environment variables +The function uses the following environment variables: +- `ENV`: The environment to use. It can be `dev`, `staging` or `prod`. Default is `dev`. +- `DATASETS_BUCKET_NAME`: The bucket name where the datasets are stored. The task will fail if this is not defined. The variable has to include the suffix, like `-dev`, `-qa` or `-prod`. diff --git a/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py b/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py new file mode 100644 index 000000000..d899b628d --- /dev/null +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py @@ -0,0 +1,529 @@ +# +# +# 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. +# +# This module provides the PmtilesBuilder class and related functions to generate PMTiles files +# from GTFS datasets. It handles downloading required files from Google Cloud Storage, processing +# and indexing GTFS data, generating GeoJSON and JSON outputs, running Tippecanoe to create PMTiles, +# and uploading the results back to GCS. +import csv +import json +import logging +import os +import subprocess +import tempfile +from enum import Enum +from google.cloud import storage + +from shared.helpers.logger import get_logger + +STOP_TIMES_FILE = "stop_times.txt" +SHAPES_FILE = "shapes.txt" +TRIPS_FILE = "trips.txt" +ROUTES_FILE = "routes.txt" +STOPS_FILE = "stops.txt" + + +def build_pmtiles_handler(payload) -> dict: + """ + Entrypoint for building PMTiles files from a GTFS dataset. + """ + try: + # Create a temporary folder to work in. It will be deleted when exiting the block. + with tempfile.TemporaryDirectory(prefix="build_pmtiles_") as temp_dir: + # If DEBUG_WORKDIR is set, use it as the work directory so it survives at the end and can be examined. + # In that case temp_dir will not be used but still deleted at the end of the block. + + debug_workdir = os.getenv("DEBUG_WORKDIR") + if debug_workdir: + os.makedirs(debug_workdir, exist_ok=True) + workdir = debug_workdir + else: + workdir = temp_dir + feed_stable_id, dataset_stable_id = PmtilesBuilder._get_parameters(payload) + result = { + "params": { + "feed_stable_id": feed_stable_id, + "dataset_stable_id": dataset_stable_id, + }, + } + builder = PmtilesBuilder( + feed_stable_id=feed_stable_id, + dataset_stable_id=dataset_stable_id, + workdir=workdir, + ) + status, message = builder.build_pmtiles() + + # A failure at this point means the pmtiles could not be created because the data + # is not available. So it's not an error of the pmtiles creation. I n that case + # we log an warning instead of an error. + if status == PmtilesBuilder.OperationStatus.FAILURE: + result["warning"] = message + else: + result["message"] = "Successfully built pmtiles." + return result + + except Exception as e: + # We expect the creation of pmtiles to be run periodically (like every day). + # If it fails, we don't want GCP to retry automatically, which would be the case if we let the exception through + # So we log the error and return a message, and that will be taken as a success with no retries. + logging.exception("Failed to build PMTiles for dataset %s", dataset_stable_id) + return { + "error": f"Failed to build PMTiles for dataset {dataset_stable_id}: {e}" + } + + +class PmtilesBuilder: + """ + Orchestrates the end-to-end process of generating PMTiles files from GTFS datasets. + + This class manages downloading required files from Google Cloud Storage, processing and indexing GTFS data, + generating GeoJSON and JSON outputs, running Tippecanoe to create PMTiles, and uploading results back to GCS. + Temporary files are stored in the global `workdir` directory for local processing. + """ + + class OperationStatus(Enum): + SUCCESS = 1 + FAILURE = 2 + + def __init__( + self, + feed_stable_id: str | None = None, + dataset_stable_id: str | None = None, + workdir: str = "./workdir", + ): + self.bucket = None + self.feed_stable_id = feed_stable_id + self.dataset_stable_id = dataset_stable_id + self.bucket_name = os.getenv("DATASETS_BUCKET_NAME") + self.logger = get_logger(PmtilesBuilder.__name__, dataset_stable_id) + + self.stop_times_index = {} + self.stop_times_by_trip = None + + self.workdir = workdir + + self.logger.info("Using work directory: %s", self.workdir) + + def get_path(self, filename: str) -> str: + return os.path.join(self.workdir, filename) + + @staticmethod + def _get_parameters(payload): + """ + Get parameters from the payload and environment variables. + """ + feed_stable_id = payload.get("feed_stable_id", None) + dataset_stable_id = payload.get("dataset_stable_id", None) + return feed_stable_id, dataset_stable_id + + def build_pmtiles(self): + if not self.bucket_name: + raise Exception("DATASETS_BUCKET_NAME environment variable is not defined.") + + if not self.feed_stable_id or not self.dataset_stable_id: + raise Exception( + "Both feed_stable_id and dataset_stable_id must be defined." + ) + + if self.feed_stable_id not in self.dataset_stable_id: + raise Exception( + "feed_stable_id %s is not a prefix of dataset_stable_id %s." + % (self.feed_stable_id, self.dataset_stable_id) + ) + + self.logger.info("Starting PMTiles build") + unzipped_files_path = ( + f"{self.feed_stable_id}/{self.dataset_stable_id}/extracted" + ) + + status, message = self._download_files_from_gcs(unzipped_files_path) + if status == self.OperationStatus.FAILURE: + return status, message + + self._create_routes_geojson() + + self._run_tippecanoe("routes-output.geojson", "routes.pmtiles") + + self._create_stops_geojson() + + self._run_tippecanoe("stops-output.geojson", "stops.pmtiles") + + self._create_routes_json() + + files_to_upload = ["routes.pmtiles", "stops.pmtiles", "routes.json"] + self._upload_files_to_gcs(files_to_upload) + + return self.OperationStatus.SUCCESS, "success" + + def _download_files_from_gcs(self, unzipped_files_path): + self.logger.info( + "Downloading dataset from GCS bucket %s, directory %s", + self.bucket_name, + unzipped_files_path, + ) + try: + self.logger.debug("Initializing storage client") + try: + self.bucket = storage.Client().get_bucket(self.bucket_name) + except Exception as e: + msg = f"Bucket '{self.bucket_name}' does not exist or is inaccessible: {e}" + self.logger.warning(msg) + return self.OperationStatus.FAILURE, msg + + self.logger.debug("Getting blobs with prefix: %s", unzipped_files_path) + blobs = list(self.bucket.list_blobs(prefix=unzipped_files_path)) + self.logger.debug("Found %d blobs", len(blobs)) + if not blobs: + msg = f"Directory '{unzipped_files_path}' does not exist or is empty in bucket '{self.bucket_name}'." + self.logger.warning(msg) + return self.OperationStatus.FAILURE, msg + + files = [ + {"name": ROUTES_FILE, "required": True}, + {"name": STOP_TIMES_FILE, "required": True}, + {"name": TRIPS_FILE, "required": True}, + {"name": STOPS_FILE, "required": True}, + {"name": SHAPES_FILE, "required": False}, + ] + for file_info in files: + file_name = file_info["name"] + required = file_info["required"] + blob_path = f"{unzipped_files_path}/{file_name}" + blob = self.bucket.blob(blob_path) + if not blob.exists(): + if required: + msg = f"Required file '{blob_path}' does not exist in bucket '{self.bucket_name}'." + self.logger.warning(msg) + return self.OperationStatus.FAILURE, msg + self.logger.debug( + "Optional file %s does not exist in bucket %s", + blob_path, + self.bucket_name, + ) + continue + try: + blob.download_to_filename(self.get_path(file_name)) + except Exception as e: + if required: + msg = f"Error downloading required file '{blob_path}' from bucket '{self.bucket_name}': {e}" + self.logger.error(msg) + raise Exception(msg) from e + else: + msg = f"Cannot download optional file '{blob_path}' from bucket '{self.bucket_name}': {e}" + self.logger.warning(msg) + + msg = "All required files downloaded successfully." + return self.OperationStatus.SUCCESS, msg + except Exception as e: + msg = f"Error downloading files from GCS: {e}" + raise Exception(msg) from e + + def _upload_files_to_gcs(self, file_to_upload): + dest_prefix = f"{self.feed_stable_id}/{self.dataset_stable_id}/pmtiles" + self.logger.info( + "Uploading files to GCS bucket %s, directory %s", + self.bucket_name, + dest_prefix, + ) + try: + blobs_to_delete = list(self.bucket.list_blobs(prefix=dest_prefix + "/")) + for blob in blobs_to_delete: + blob.delete() + self.logger.debug("Deleted existing blob: %s", blob.name) + for file_name in file_to_upload: + file_path = os.path.join(self.workdir, file_name) + if not os.path.exists(file_path): + self.logger.warning("File not found: %s", file_path) + continue + blob_path = f"{dest_prefix}/{file_name}" + blob = self.bucket.blob(blob_path) + blob.upload_from_filename(file_path) + self.logger.debug( + "Uploaded %s to gs://%s/%s", + file_path, + self.bucket_name, + blob_path, + ) + except Exception as e: + raise Exception(f"Failed to upload files to GCS: {e}") from e + + def _create_shapes_index(self) -> dict: + """ + Create an index for shapes.txt file to quickly access shape points by shape_id. + We create the index to save memory. With the index, we keep a list of positions in the file for each shape. + If instead we read the whole file into memory, we would need 2 floats for the longitude and latitude plus an + int for the sequence number for each point. + The largest number of shapes we have currently in a dataset is 37 millions. + This means about 900 MB if we have the index, and 1.6 GB if read the coordinates in memory. + Returns: + A dictionary with key shaped_id and values a list of positions in the shapes.txt file. + """ + # TODO: see if we can get rid of the index by reading the shapes coordinates with the memory efficient numpy. + self.logger.info("Creating shapes index") + shapes_index = {} + try: + with open( + self.get_path(SHAPES_FILE), "r", encoding="utf-8", newline="" + ) as f: + header = f.readline() + columns = next(csv.reader([header])) + shapes_index["columns"] = columns + count = 0 + while True: + pos = f.tell() + line = f.readline() + if not line: + break + row = dict(zip(columns, next(csv.reader([line])))) + sid = row["shape_id"] + shapes_index.setdefault(sid, []).append(pos) + count += 1 + if count % 1000000 == 0: + self.logger.debug("Indexed %d lines so far...", count) + self.logger.debug("Total indexed lines: %d", count) + self.logger.debug("Total unique shape_ids: %d", len(shapes_index)) + except Exception as e: + self.logger.warning("Cannot read shapes file: %s", e) + return shapes_index + + def _read_csv(self, filename): + try: + self.logger.debug("Loading %s", filename) + with open(filename, newline="", encoding="utf-8") as f: + return list(csv.DictReader(f)) + except Exception as e: + raise Exception(f"Failed to read CSV file {filename}: {e}") from e + + def _get_shape_points(self, shape_id, index): + self.logger.debug("Getting shape points for shape_id %s", shape_id) + try: + points = [] + with open( + self.get_path(SHAPES_FILE), "r", encoding="utf-8", newline="" + ) as f: + for pos in index.get(shape_id, []): + f.seek(pos) + line = f.readline() + row = dict(zip(index["columns"], next(csv.reader([line])))) + points.append( + ( + float(row["shape_pt_lon"]), + float(row["shape_pt_lat"]), + int(row["shape_pt_sequence"]), + ) + ) + points.sort(key=lambda x: x[2]) + self.logger.debug( + " Found %d points for shape_id %s", len(points), shape_id + ) + return [pt[:2] for pt in points] + except Exception as e: + raise Exception(f"Failed to get shape points for {shape_id}: {e}") from e + + def _create_routes_geojson(self): + try: + shapes_index = self._create_shapes_index() + self.logger.info("Creating routes geojson (optimized for memory)") + + # Load stops into memory (usually not huge) + # Used only if there is no shapes for a route + coordinates_indexed_by_stop = { + s["stop_id"]: (float(s["stop_lon"]), float(s["stop_lat"])) + for s in self._read_csv(self.get_path(STOPS_FILE)) + } + + # We want the shape of a route. To do that we find one trip for each route. We will assume that all + # trips for a route have the same shape_id. If not I am not sure how to represent this in the route map + # when we select a route. + shape_map_indexed_by_route = {} + trip_map_indexed_by_route = {} + with open(self.get_path(TRIPS_FILE), newline="", encoding="utf-8") as f: + for row in csv.DictReader(f): + trip_id = row["trip_id"] + route_id = row["route_id"] + shape_id = row.get("shape_id", "") + if shape_id and route_id not in shape_map_indexed_by_route: + shape_map_indexed_by_route[route_id] = shape_id + if trip_id and route_id not in trip_map_indexed_by_route: + trip_map_indexed_by_route[route_id] = trip_id + + features = [] + missing_coordinates_routes = set() + routes_geojson = self.get_path("routes-output.geojson") + with open(routes_geojson, "w", encoding="utf-8") as geojson_file: + geojson_file.write('{"type": "FeatureCollection", "features": [\n') + first = True + with open( + self.get_path(ROUTES_FILE), newline="", encoding="utf-8" + ) as routes_file: + for i, route in enumerate(csv.DictReader(routes_file), 1): + route_id = route["route_id"] + + shape_id = shape_map_indexed_by_route.get(route_id, "") + + coordinates = [] + if shape_id: + coordinates = self._get_shape_points(shape_id, shapes_index) + if not coordinates: + # We don't have the coordinates for the shape, fallback on stop_times and stops + trip_id = trip_map_indexed_by_route.get(route_id, "") + + if trip_id: + trip_stop_times = self._get_trip_stop_times(trip_id) + # We assume stop_times is already sorted by stop_sequence in the file. + # According to the SPECS: + # The values must increase along the trip but do not need to be consecutive. + coordinates = [ + coordinates_indexed_by_stop[stop_id] + for stop_id in trip_stop_times + if stop_id in coordinates_indexed_by_stop + ] + if not coordinates: + missing_coordinates_routes.add(route_id) + continue + feature = { + "type": "Feature", + "properties": {k: route[k] for k in route}, + "geometry": { + "type": "LineString", + "coordinates": coordinates, + }, + } + + if not first: + geojson_file.write(",\n") + geojson_file.write(json.dumps(feature)) + first = False + + if i % 100 == 0 or i == 1: + self.logger.debug( + "Processed route %d (route_id: %s)", i, route_id + ) + + geojson_file.write("\n]}") + + if missing_coordinates_routes: + self.logger.info( + "Routes without coordinates: %s", list(missing_coordinates_routes) + ) + self.logger.debug( + "Wrote %d features to routes-output.geojson", len(features) + ) + except Exception as e: + raise Exception(f"Failed to create routes GeoJSON: {e}") from e + + def _get_trip_stop_times(self, trip_id): + # Lazy instantiation of the dictionary, because we may not need it al all if there is a shape. + if self.stop_times_by_trip is None: + self.stop_times_by_trip = {} + with open( + self.get_path(STOP_TIMES_FILE), newline="", encoding="utf-8" + ) as f: + for row in csv.DictReader(f): + self.stop_times_by_trip.setdefault(row["trip_id"], []).append( + row["stop_id"] + ) + + return self.stop_times_by_trip.get(trip_id, []) + + def _run_tippecanoe(self, input_file, output_file): + self.logger.info("Running tippecanoe for input file %s", input_file) + try: + cmd = [ + "tippecanoe", + "-o", + self.get_path(output_file), + "--force", + "--no-tile-size-limit", + "-zg", + self.get_path(input_file), + ] + self.logger.debug("Running command: %s", " ".join(cmd)) + result = subprocess.run(cmd, capture_output=True, text=True) + if result.stdout: + self.logger.debug("Tippecanoe output:\n%s", result.stdout) + if result.returncode != 0: + self.logger.error("Tippecanoe error:\n%s", result.stderr) + raise Exception(f"Tippecanoe failed with exit code {result.returncode}") + self.logger.debug("Tippecanoe command executed successfully.") + except Exception as e: + raise Exception( + f"Failed to run tippecanoe for output file {output_file}: {e}" + ) from e + + def _create_stops_geojson(self): + self.logger.info("Creating stops geojson...") + try: + stops = self._read_csv(self.get_path(STOPS_FILE)) + if isinstance(stops, dict) and "error" in stops: + return stops + self.logger.debug("Loaded %d stops.", len(stops)) + + features = [] + for i, stop in enumerate(stops, 1): + try: + lon = float(stop["stop_lon"]) + lat = float(stop["stop_lat"]) + except (KeyError, ValueError): + self.logger.info( + "Skipping stop %s: invalid coordinates", + stop.get("stop_id", ""), + ) + continue + features.append( + { + "type": "Feature", + "properties": {k: stop[k] for k in stop}, + "geometry": {"type": "Point", "coordinates": [lon, lat]}, + } + ) + + geojson = {"type": "FeatureCollection", "features": features} + stops_geojson = self.get_path("stops-output.geojson") + + self.logger.debug( + "Writing %d features to stops-output.geojson for dataset %s", + len(features), + self.dataset_stable_id, + ) + with open(stops_geojson, "w", encoding="utf-8") as f: + json.dump(geojson, f) + except Exception as e: + raise Exception( + f"Failed to create stops GeoJSON for dataset {self.dataset_stable_id}: {e}" + ) from e + + def _create_routes_json(self): + self.logger.info("Creating routes json...") + try: + routes = [] + with open(self.get_path(ROUTES_FILE), newline="", encoding="utf-8") as f: + reader = csv.DictReader(f) + for row in reader: + route = { + "routeId": row.get("route_id", ""), + "routeName": row.get("route_long_name", ""), + "color": f"#{row.get('route_color', '000000')}", + "textColor": f"#{row.get('route_text_color', 'FFFFFF')}", + "routeType": f"{row.get('route_type', 'unknown')}", + } + routes.append(route) + + with open(f"{self.workdir}/routes.json", "w", encoding="utf-8") as f: + json.dump(routes, f, ensure_ascii=False, indent=4) + + self.logger.debug("Converted %d routes to routes.json.", len(routes)) + except Exception as e: + raise Exception(f"Failed to create routes JSON for dataset: {e}") from e diff --git a/functions-python/tasks_executor/tests/tasks/pmtiles_builder/test_build_pmtiles.py b/functions-python/tasks_executor/tests/tasks/pmtiles_builder/test_build_pmtiles.py new file mode 100644 index 000000000..2f02ad3f3 --- /dev/null +++ b/functions-python/tasks_executor/tests/tasks/pmtiles_builder/test_build_pmtiles.py @@ -0,0 +1,510 @@ +import csv +import json +import logging +import tempfile +import unittest +from contextlib import contextmanager +from unittest.mock import patch, MagicMock +import os + +from tasks.pmtiles_builder.build_pmtiles import ( + PmtilesBuilder, + build_pmtiles_handler, +) + + +@contextmanager +def suppress_logging(level=logging.CRITICAL): + previous_level = logging.root.manager.disable + logging.disable(level) + try: + yield + finally: + logging.disable(previous_level) + + +class TestDownloadFilesFromGCS(unittest.TestCase): + def setUp(self): + self.builder = PmtilesBuilder( + feed_stable_id="feed123", + dataset_stable_id="feed123_dataset456", + workdir="/tmp", + ) + self.builder.bucket = MagicMock() + self.builder.logger = MagicMock() + self.builder.bucket_name = "test-bucket" # Ensure bucket_name is set + + @patch("tasks.pmtiles_builder.build_pmtiles.storage.Client") + @patch("tasks.pmtiles_builder.build_pmtiles.ROUTES_FILE", "routes.txt") + def test_required_file_missing(self, mock_storage_client): + mock_storage_client.return_value.get_bucket.return_value = self.builder.bucket + blob = MagicMock() + blob.exists.return_value = False + self.builder.bucket.blob.return_value = blob + # Simulate directory exists by returning a non-empty list + self.builder.bucket.list_blobs.return_value = [MagicMock()] + status, msg = self.builder._download_files_from_gcs("some/path") + self.assertEqual(status, self.builder.OperationStatus.FAILURE) + self.assertIn("Required file", msg) + self.builder.logger.warning.assert_called() + + @patch("tasks.pmtiles_builder.build_pmtiles.SHAPES_FILE", "shapes.txt") + def test_optional_file_missing(self): + blob = MagicMock() + blob.exists.return_value = False + self.builder.bucket.blob.return_value = blob + status, msg = self.builder._download_files_from_gcs("some/path") + self.builder.logger.debug.assert_called() + + @patch("tasks.pmtiles_builder.build_pmtiles.storage.Client") + @patch("tasks.pmtiles_builder.build_pmtiles.ROUTES_FILE", "routes.txt") + def test_file_download_success(self, mock_storage_client): + mock_storage_client.return_value.get_bucket.return_value = self.builder.bucket + blob = MagicMock() + blob.exists.return_value = True + blob.download_to_filename.return_value = None + self.builder.bucket.blob.return_value = blob + # Simulate directory exists by returning a non-empty list + self.builder.bucket.list_blobs.return_value = [MagicMock()] + status, msg = self.builder._download_files_from_gcs("some/path") + self.assertEqual(status, self.builder.OperationStatus.SUCCESS) + self.assertIn("downloaded successfully", msg) + + @patch("tasks.pmtiles_builder.build_pmtiles.storage.Client") + def test_bucket_not_exist(self, mock_client): + mock_client.return_value.get_bucket.side_effect = Exception("Bucket not found") + status, message = self.builder._download_files_from_gcs("some/path") + self.assertEqual(status, PmtilesBuilder.OperationStatus.FAILURE) + self.assertIn("Bucket not found", message) + + @patch("tasks.pmtiles_builder.build_pmtiles.storage.Client") + def test_download_required_file_error(self, mock_storage_client): + mock_storage_client.return_value.get_bucket.return_value = self.builder.bucket + blob = MagicMock() + blob.exists.return_value = True + blob.download_to_filename.side_effect = Exception("Download failed") + self.builder.bucket.blob.return_value = blob + self.builder.bucket.list_blobs.return_value = [MagicMock()] + # Only required files + with self.assertRaises(Exception) as context: + self.builder._download_files_from_gcs("some/path") + self.assertIn("Error downloading required file", str(context.exception)) + self.builder.logger.error.assert_called() + + @patch("tasks.pmtiles_builder.build_pmtiles.storage.Client") + def test_download_optional_file_error(self, mock_storage_client): + mock_storage_client.return_value.get_bucket.return_value = self.builder.bucket + blob = MagicMock() + blob.exists.return_value = True + + def download_side_effect(path): + if path.endswith("shapes.txt"): + raise Exception("Download failed") + # Simulate success for other files + + blob.download_to_filename.side_effect = download_side_effect + self.builder.bucket.blob.return_value = blob + self.builder.bucket.list_blobs.return_value = [MagicMock()] + with patch("tasks.pmtiles_builder.build_pmtiles.SHAPES_FILE", "shapes.txt"): + status, msg = self.builder._download_files_from_gcs("some/path") + self.assertEqual(status, self.builder.OperationStatus.SUCCESS) + self.builder.logger.warning.assert_called_with( + "Cannot download optional file 'some/path/shapes.txt' from bucket 'test-bucket': Download failed" + ) + self.assertEqual(msg, "All required files downloaded successfully.") + + +class TestPmtilesBuilder(unittest.TestCase): + def setUp(self): + self.feed_stable_id = "feed123" + self.dataset_stable_id = "feed123_dataset456" + os.environ["DATASETS_BUCKET_NAME"] = "test-bucket" + self.builder = PmtilesBuilder(self.feed_stable_id, self.dataset_stable_id) + + @patch("tasks.pmtiles_builder.build_pmtiles.subprocess.run") + def test_run_tippecanoe_success(self, mock_run): + mock_run.return_value = MagicMock(returncode=0, stdout="ok", stderr="") + self.builder._run_tippecanoe("input.geojson", "output.pmtiles") + mock_run.assert_called_once() + + @patch("tasks.pmtiles_builder.build_pmtiles.subprocess.run") + def test_run_tippecanoe_failure(self, mock_run): + mock_run.return_value = MagicMock(returncode=1, stdout="", stderr="error") + with self.assertRaises(Exception), suppress_logging(): + self.builder._run_tippecanoe("input.geojson", "output.pmtiles") + + @patch( + "tasks.pmtiles_builder.build_pmtiles.PmtilesBuilder._download_files_from_gcs" + ) + @patch("tasks.pmtiles_builder.build_pmtiles.PmtilesBuilder._create_shapes_index") + @patch("tasks.pmtiles_builder.build_pmtiles.PmtilesBuilder._create_routes_geojson") + @patch("tasks.pmtiles_builder.build_pmtiles.PmtilesBuilder._run_tippecanoe") + @patch("tasks.pmtiles_builder.build_pmtiles.PmtilesBuilder._create_stops_geojson") + @patch("tasks.pmtiles_builder.build_pmtiles.PmtilesBuilder._create_routes_json") + @patch("tasks.pmtiles_builder.build_pmtiles.PmtilesBuilder._upload_files_to_gcs") + def test_build_pmtiles_success( + self, + mock_upload, + mock_routes_json, + mock_stops_geojson, + mock_run_tippecanoe, + mock_routes_geojson, + mock_shapes_index, + mock_download, + ): + self.builder.bucket = MagicMock() + self.builder.bucket.list_blobs.return_value = [] + mock_download.return_value = ( + PmtilesBuilder.OperationStatus.SUCCESS, + "All required files downloaded successfully.", + ) + status, message = self.builder.build_pmtiles() + self.assertEqual(status, PmtilesBuilder.OperationStatus.SUCCESS) + self.assertEqual(message, "success") + + def test_get_parameters(self): + payload = {"feed_stable_id": "f", "dataset_stable_id": "d"} + f, d = PmtilesBuilder._get_parameters(payload) + self.assertEqual(f, "f") + self.assertEqual(d, "d") + + def tearDown(self): + if "DATASETS_BUCKET_NAME" in os.environ: + del os.environ["DATASETS_BUCKET_NAME"] + + def test_build_pmtiles_creates_correct_shapes_index(self): + # Prepare shapes.txt + with tempfile.TemporaryDirectory() as temp_dir: + self.builder.workdir = temp_dir + shapes_path = os.path.join(temp_dir, "shapes.txt") + with open(shapes_path, "w", encoding="utf-8") as f: + f.write("shape_id,shape_pt_lat,shape_pt_lon,shape_pt_sequence\n") + f.write("s1,45.0,-73.0,1\n") + f.write("s1,45.1,-73.1,2\n") + f.write("s2,46.0,-74.0,1\n") + + index = self.builder._create_shapes_index() + self.assertIn("columns", index) + self.assertEqual( + index["columns"], + ["shape_id", "shape_pt_lat", "shape_pt_lon", "shape_pt_sequence"], + ) + self.assertIn("s1", index) + self.assertIn("s2", index) + self.assertEqual(len(index["s1"]), 2) + self.assertEqual(len(index["s2"]), 1) + + def test_get_shape_points(self): + # Prepare shapes.txt + with tempfile.TemporaryDirectory() as temp_dir: + self.builder.workdir = temp_dir + shapes_path = self.builder.get_path("shapes.txt") + + with open(shapes_path, "w", encoding="utf-8") as f: + f.write("shape_id,shape_pt_lat,shape_pt_lon,shape_pt_sequence\n") + f.write("s1,45.0,-73.0,1\n") + f.write("s1,45.1,-73.1,2\n") + + # Build index with file positions + index = {} + with open(shapes_path, "r", encoding="utf-8") as f: + header = f.readline() + columns = next(csv.reader([header])) + pos1 = f.tell() + f.readline() + pos2 = f.tell() + f.readline() + index["s1"] = [pos1, pos2] + index["columns"] = columns + + # Call _get_shape_points + points = self.builder._get_shape_points("s1", index) + self.assertEqual(points, [(-73.0, 45.0), (-73.1, 45.1)]) + + def test_create_routes_geojson(self): + # Prepare minimal GTFS files + with tempfile.TemporaryDirectory() as temp_dir: + self.builder.workdir = temp_dir + with open(self.builder.get_path("routes.txt"), "w", encoding="utf-8") as f: + f.write( + "route_id,route_long_name,route_color,route_text_color,route_type\nr1,Route 1,FF0000,FFFFFF,3\n" + ) + with open(self.builder.get_path("trips.txt"), "w", encoding="utf-8") as f: + f.write("route_id,service_id,trip_id,shape_id\nr1,svc1,t1,s1\n") + with open(self.builder.get_path("stops.txt"), "w", encoding="utf-8") as f: + f.write( + "stop_id,stop_lat,stop_lon\nstop1,45.0,-73.0\nstop2,45.1,-73.1\n" + ) + with open( + self.builder.get_path("stop_times.txt"), "w", encoding="utf-8" + ) as f: + f.write("trip_id,arrival_time,departure_time,stop_id,stop_sequence\n") + f.write("t1,08:00:00,08:00:00,stop1,1\n") + f.write("t1,08:10:00,08:10:00,stop2,2\n") + with open(self.builder.get_path("shapes.txt"), "w", encoding="utf-8") as f: + f.write("shape_id,shape_pt_lat,shape_pt_lon,shape_pt_sequence\n") + f.write("s1,45.0,-73.0,1\n") + f.write("s1,45.1,-73.1,2\n") + + # Call the method + self.builder._create_routes_geojson() + + # Assert output file exists and is valid GeoJSON + output_path = self.builder.get_path("routes-output.geojson") + self.assertTrue(os.path.exists(output_path)) + with open(output_path, "r", encoding="utf-8") as f: + data = f.read() + self.assertIn("FeatureCollection", data) + + def test_create_stops_geojson(self): + # Prepare minimal stops.txt + with tempfile.TemporaryDirectory() as temp_dir: + self.builder.workdir = temp_dir + with open(self.builder.get_path("stops.txt"), "w", encoding="utf-8") as f: + f.write("stop_id,stop_lat,stop_lon\n") + f.write("stop1,45.0,-73.0\n") + f.write("stop2,45.1,-73.1\n") + + # Call the method + self.builder._create_stops_geojson() + + # Assert output file exists and is valid GeoJSON + output_path = self.builder.get_path("stops-output.geojson") + self.assertTrue(os.path.exists(output_path)) + with open(output_path, "r", encoding="utf-8") as f: + data = f.read() + self.assertIn("FeatureCollection", data) + + def test_create_routes_json(self): + # Prepare minimal routes.txt + with tempfile.TemporaryDirectory() as temp_dir: + self.builder.workdir = temp_dir + with open(self.builder.get_path("routes.txt"), "w", encoding="utf-8") as f: + f.write( + "route_id,route_long_name,route_color,route_text_color,route_type\n" + ) + f.write("r1,Route 1,FF0000,FFFFFF,3\n") + + # Call the method + self.builder._create_routes_json() + + # Assert output file exists and is valid JSON + output_path = self.builder.get_path("routes.json") + self.assertTrue(os.path.exists(output_path)) + with open(output_path, "r", encoding="utf-8") as f: + data = f.read() + self.assertIn("routeId", data) + self.assertIn("routeName", data) + + def test_create_routes_json_exception(self): + # Ensure routes.txt does not exist to trigger the exception + routes_path = self.builder.get_path("routes.txt") + if os.path.exists(routes_path): + os.remove(routes_path) + with self.assertRaises(Exception) as cm: + self.builder._create_routes_json() + self.assertIn("Failed to create routes JSON for dataset", str(cm.exception)) + + def test_build_pmtiles_exception(self): + # Set up builder with missing bucket_name to trigger an exception in _download_files_from_gcs + self.builder.bucket_name = "invalid-bucket" + # Patch _download_files_from_gcs to raise an exception + with patch.object( + self.builder, + "_download_files_from_gcs", + side_effect=Exception("Download failed"), + ): + # This is a test that purposely generates an exception to verify that error handling in build_pmtiles + # works as expected. The suppress_logging context manager is used to silence log output during the test. + with suppress_logging(): + with self.assertRaises(Exception) as cm: + self.builder.build_pmtiles() + self.assertIn("Download failed", str(cm.exception)) + + def test_upload_files_to_gcs_missing_file(self): + self.builder.bucket = MagicMock() + self.builder.bucket.blob.return_value = MagicMock() + self.builder.bucket.list_blobs.return_value = [] + + missing_file = "notfound.pmtiles" + missing_path = self.builder.get_path(missing_file) + if os.path.exists(missing_path): + os.remove(missing_path) + + with patch( + "os.path.exists", + side_effect=lambda path: False if missing_file in path else True, + ), self.assertLogs(level="WARNING") as log_cm: + self.builder._upload_files_to_gcs([missing_file]) + self.assertTrue( + any(f"File not found: {missing_path}" in msg for msg in log_cm.output) + ) + + def test_create_routes_geojson_fallback_to_stop_coordinates(self): + # Prepare minimal GTFS files with no shape_id for the trip + with tempfile.TemporaryDirectory() as temp_dir: + self.builder.workdir = temp_dir + with open(self.builder.get_path("routes.txt"), "w", encoding="utf-8") as f: + f.write( + "route_id,route_long_name,route_color,route_text_color,route_type\nr1,Route 1,FF0000,FFFFFF,3\n" + ) + with open(self.builder.get_path("trips.txt"), "w", encoding="utf-8") as f: + f.write( + "route_id,service_id,trip_id,shape_id\nr1,svc1,t1,\n" + ) # shape_id is empty + with open(self.builder.get_path("stops.txt"), "w", encoding="utf-8") as f: + f.write( + "stop_id,stop_lat,stop_lon\nstop1,45.0,-73.0\nstop2,45.1,-73.1\n" + ) + with open( + self.builder.get_path("stop_times.txt"), "w", encoding="utf-8" + ) as f: + f.write("trip_id,arrival_time,departure_time,stop_id,stop_sequence\n") + f.write("t1,08:00:00,08:00:00,stop1,1\n") + f.write("t1,08:10:00,08:10:00,stop2,2\n") + with open(self.builder.get_path("shapes.txt"), "w", encoding="utf-8") as f: + f.write("shape_id,shape_pt_lat,shape_pt_lon,shape_pt_sequence\n") + + # Call the method + self.builder._create_routes_geojson() + + # Assert output file exists and contains the expected coordinates + output_path = self.builder.get_path("routes-output.geojson") + self.assertTrue(os.path.exists(output_path)) + with open(output_path, "r", encoding="utf-8") as f: + data = json.load(f) + self.assertEqual(data["type"], "FeatureCollection") + self.assertEqual(len(data["features"]), 1) + coords = data["features"][0]["geometry"]["coordinates"] + self.assertEqual(coords, [[-73.0, 45.0], [-73.1, 45.1]]) + + def test_create_stops_geojson_invalid_coordinates(self): + with tempfile.TemporaryDirectory() as temp_dir: + self.builder.workdir = temp_dir + stops_path = self.builder.get_path("stops.txt") + with open(stops_path, "w", encoding="utf-8") as f: + f.write("stop_id,stop_lat,stop_lon\n") + f.write("stop1,45.0,-73.0\n") # valid + f.write("stop2,not_a_lat,-73.1\n") # invalid lat + + with self.assertLogs(level="INFO") as log_cm: + self.builder._create_stops_geojson() + self.assertTrue( + any( + "Skipping stop stop2: invalid coordinates" in msg + for msg in log_cm.output + ) + ) + + output_path = self.builder.get_path("stops-output.geojson") + self.assertTrue(os.path.exists(output_path)) + with open(output_path, "r", encoding="utf-8") as f: + data = json.load(f) + self.assertEqual(len(data["features"]), 1) + self.assertEqual(data["features"][0]["properties"]["stop_id"], "stop1") + + +class TestBuildPmtilesHandlerIntegration(unittest.TestCase): + def setUp(self): + self.test_dir = tempfile.TemporaryDirectory() + os.environ["DATASETS_BUCKET_NAME"] = "test-bucket" + self.feed_stable_id = "feed123" + self.dataset_stable_id = "feed123_dataset456" + self.builder = PmtilesBuilder( + self.feed_stable_id, self.dataset_stable_id, workdir=self.test_dir.name + ) + + # Create minimal GTFS files using builder.get_path + files = { + "routes.txt": ( + "route_id,route_long_name,route_color,route_text_color,route_type\n" + "r1,Route 1,FF0000,FFFFFF,3\n" + ), + "shapes.txt": ( + "shape_id,shape_pt_lat,shape_pt_lon,shape_pt_sequence\n" + "s1,45.0,-73.0,1\n" + "s1,45.1,-73.1,2\n" + ), + "trips.txt": ("route_id,service_id,trip_id,shape_id\n" "r1,svc1,t1,s1\n"), + "stops.txt": ( + "stop_id,stop_lat,stop_lon\n" "stop1,45.0,-73.0\n" "stop2,45.1,-73.1\n" + ), + "stop_times.txt": ( + "trip_id,arrival_time,departure_time,stop_id,stop_sequence\n" + "t1,08:00:00,08:00:00,stop1,1\n" + "t1,08:10:00,08:10:00,stop2,2\n" + ), + } + for fname, content in files.items(): + with open(self.builder.get_path(fname), "w", encoding="utf-8") as f: + f.write(content) + + def tearDown(self): + self.test_dir.cleanup() + if "DATASETS_BUCKET_NAME" in os.environ: + del os.environ["DATASETS_BUCKET_NAME"] + + def test_build_pmtiles_handler_missing_bucket_env(self): + os.environ.pop("DATASETS_BUCKET_NAME", None) + payload = { + "feed_stable_id": self.feed_stable_id, + "dataset_stable_id": self.dataset_stable_id, + } + with suppress_logging(): + result = build_pmtiles_handler(payload) + self.assertIn("error", result) + self.assertIn( + "DATASETS_BUCKET_NAME environment variable is not defined.", result["error"] + ) + + def test_build_pmtiles_handler_missing_ids(self): + payload = {} + with suppress_logging(): + result = build_pmtiles_handler(payload) + self.assertIn("error", result) + self.assertIn( + "Both feed_stable_id and dataset_stable_id must be defined.", + result["error"], + ) + + def test_build_pmtiles_handler_feed_not_prefix(self): + payload = { + "feed_stable_id": "notprefix", + "dataset_stable_id": self.dataset_stable_id, + } + with suppress_logging(): + result = build_pmtiles_handler(payload) + self.assertIn("error", result) + self.assertIn("is not a prefix of dataset_stable_id", result["error"]) + + +class TestPmtilesBuilderUpload(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + self.builder = PmtilesBuilder( + feed_stable_id="foo", + dataset_stable_id="foo_bar", + workdir=self.temp_dir.name, + ) + self.builder.bucket = MagicMock() + self.mock_blob = MagicMock() + self.builder.bucket.blob.return_value = self.mock_blob + self.builder.bucket.list_blobs.return_value = [] + + def tearDown(self): + self.temp_dir.cleanup() + + def test_upload_files_to_gcs(self): + test_file = self.builder.get_path("routes.pmtiles") + with open(test_file, "w") as f: + f.write("dummy data") + + self.builder._upload_files_to_gcs(["routes.pmtiles"]) + self.builder.bucket.blob.assert_called_with( + "foo/foo_bar/pmtiles/routes.pmtiles" + ) + self.mock_blob.upload_from_filename.assert_called_with(test_file) + + +if __name__ == "__main__": + unittest.main()