From 920f412022c5e424eebb8500a0f24d2b5f79dc17 Mon Sep 17 00:00:00 2001 From: jcpitre Date: Fri, 25 Jul 2025 12:20:46 -0400 Subject: [PATCH 01/14] First draft --- .../tasks_executor/function_config.json | 4 +- .../tasks_executor/requirements.txt | 9 +- functions-python/tasks_executor/src/main.py | 6 +- .../tasks/pmtiles_builder/build_pmtiles.py | 199 ++++++++++++++++++ .../pmtiles_builder/create_routes_geojson.py | 111 ++++++++++ .../pmtiles_builder/create_shapes_index.py | 39 ++++ .../tasks/pmtiles_builder/run_tippecanoe.py | 19 ++ 7 files changed, 382 insertions(+), 5 deletions(-) create mode 100644 functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py create mode 100644 functions-python/tasks_executor/src/tasks/pmtiles_builder/create_routes_geojson.py create mode 100644 functions-python/tasks_executor/src/tasks/pmtiles_builder/create_shapes_index.py create mode 100644 functions-python/tasks_executor/src/tasks/pmtiles_builder/run_tippecanoe.py diff --git a/functions-python/tasks_executor/function_config.json b/functions-python/tasks_executor/function_config.json index 48578e272..eda4462f0 100644 --- a/functions-python/tasks_executor/function_config.json +++ b/functions-python/tasks_executor/function_config.json @@ -4,7 +4,7 @@ "entry_point": "tasks_executor", "timeout": 540, "memory": "4Gi", - "trigger_http": false, + "trigger_http": true, "include_folders": ["helpers"], "include_api_folders": ["database_gen", "database", "common"], "environment_variables": [], @@ -13,7 +13,7 @@ "key": "FEEDS_DATABASE_URL" } ], - "ingress_settings": "ALLOW_ALL", + "ingress_settings": "all", "max_instance_request_concurrency": 1, "max_instance_count": 1, "min_instance_count": 0, diff --git a/functions-python/tasks_executor/requirements.txt b/functions-python/tasks_executor/requirements.txt index 9822923f2..8ab832d4c 100644 --- a/functions-python/tasks_executor/requirements.txt +++ b/functions-python/tasks_executor/requirements.txt @@ -15,9 +15,14 @@ SQLAlchemy==2.0.23 geoalchemy2==0.14.7 # Google specific packages for this function -google-cloud-workflows google-cloud-pubsub +google-cloud-datastore +google-cloud-workflows +cloudevents~=1.10.1 flask # Configuration -python-dotenv==1.0.0 \ No newline at end of file +python-dotenv==1.0.0 +google-cloud-storage +tippecanoe + diff --git a/functions-python/tasks_executor/src/main.py b/functions-python/tasks_executor/src/main.py index 694f2e27c..bfc7af31a 100644 --- a/functions-python/tasks_executor/src/main.py +++ b/functions-python/tasks_executor/src/main.py @@ -25,7 +25,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" @@ -49,6 +49,10 @@ "description": "Rebuilds missing bounding boxes for GTFS datasets that contain valid stops.txt files.", "handler": rebuild_missing_bounding_boxes_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/build_pmtiles.py b/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py new file mode 100644 index 000000000..dc1f809d6 --- /dev/null +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py @@ -0,0 +1,199 @@ +# +# +# 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 logging +import os +import shutil +import subprocess +import sys + +from google.cloud import storage +from sqlalchemy.orm import Session + +sys.path.append(os.path.dirname(os.path.abspath(__file__))) # noqa: E402 + +from create_shapes_index import create_shapes_index # noqa: E402 +from create_routes_geojson import create_routes_geojson # noqa: E402 +from run_tippecanoe import run_tippecanoe # noqa: E402 + + +def build_pmtiles_handler(payload) -> dict: + """ + Rebuilds missing validation reports for GTFS datasets. + This function processes datasets with missing validation reports using the GTFS validator workflow. + The payload structure is: + { + "dry_run": bool, # [optional] If True, do not execute the workflow + "feed_stable_id": int, # [optional] Filter datasets older than this number of days(default: 14 days ago) + "dataset_stable_id": list[str] # [optional] Filter datasets by status(in) + } + Args: + payload (dict): The payload containing the task details. + Returns: + str: A message indicating the result of the operation with the total_processed datasets. + """ + dry_run: bool + ( + dry_run, + feed_stable_id, + dataset_stable_id, + ) = get_parameters(payload) + + return build_pmtiles( + dry_run=dry_run, + feed_stable_id=feed_stable_id, + dataset_stable_id=dataset_stable_id, + ) + + +def build_pmtiles( + dry_run: bool = True, + feed_stable_id: str | None = None, + dataset_stable_id: str | None = None, + db_session: Session | None = None, +) -> dict: + """ + Rebuilds missing validation reports for GTFS datasets. + + Args: + validator_endpoint: Validator endpoint URL + dry_run (bool): dry run flag. If True, do not execute the workflow. Default: True + filter_after_in_days (int): Filter the datasets older than this number of days. Default: 14 days ago + filter_statuses: [optional] Filter datasets by status(in). Default: None + prod_env (bool): True if target environment is production, false otherwise. Default: False + db_session: DB session + + Returns: + flask.Response: A response with message and total_processed datasets. + """ + bucket_name = os.getenv("DATASETS_BUCKET_NAME") + if not bucket_name: + return {"error": "DATASETS_BUCKET_NAME environment variable is not defined."} + + if not feed_stable_id or not dataset_stable_id: + return {"error": "Both feed_stable_id and dataset_stable_id must be defined."} + + if feed_stable_id not in dataset_stable_id: + return {"error": "feed_stable_id must be a substring of dataset_stable_id."} + + logging.info( + "Starting PMTiles build for feed %s and dataset %s on bucket %s", + feed_stable_id, + dataset_stable_id, + bucket_name, + ) + unzipped_files_path = f"{feed_stable_id}/{dataset_stable_id}/extracted" + + logging.info("Initializing storage client") + bucket = storage.Client().get_bucket(bucket_name) + logging.info("Getting blobs with prefix: %s", unzipped_files_path) + blobs = list(bucket.list_blobs(prefix=unzipped_files_path)) + logging.info("Found %d blobs", len(blobs)) + if not blobs: + return { + "error": f"Directory '{unzipped_files_path}' does not exist in bucket '{bucket_name}'." + } + + local_dir = "./unzipped" + if os.path.exists(local_dir): + shutil.rmtree(local_dir) + os.makedirs(local_dir, exist_ok=True) + download_files_from_gcs(bucket_name, unzipped_files_path, local_dir) + + create_shapes_index(local_dir) + create_routes_geojson(local_dir) + logging.info(os.getcwd()) + + result = subprocess.run(["which", "tippecanoe"], capture_output=True, text=True) + logging.info("REsult of which command: %s", result.stdout.strip()) + + run_tippecanoe("routes.pmtiles", "routes-output.geojson", local_dir) + + upload_files_to_gcs( + bucket_name, local_dir, ["routes.pmtiles"], feed_stable_id, dataset_stable_id + ) + + result = subprocess.run( + ["ls", "-l", "-R", local_dir], capture_output=True, text=True + ) + logging.info("Files created:\n%s", result.stdout.strip()) + return { + "message": f"Directory '{unzipped_files_path}' exists in bucket '{bucket_name}'." + } + + +def get_parameters(payload): + """ + Get parameters from the payload and environment variables. + + Args: + payload (dict): dictionary containing the payload data. + Returns: + dict: dict with: dry_run, filter_after_in_days, filter_statuses, prod_env, validator_endpoint parameters + """ + dry_run = payload.get("dry_run", True) + dry_run = dry_run if isinstance(dry_run, bool) else str(dry_run).lower() == "true" + feed_stable_id = payload.get("feed_stable_id", None) + dataset_stable_id = payload.get("dataset_stable_id", None) + + return dry_run, feed_stable_id, dataset_stable_id + + +def download_files_from_gcs(bucket_name, unzipped_files_path, local_dir): + file_names = [ + "routes.txt", + "shapes.txt", + "stop_times.txt", + "trips.txt", + "stops.txt", + ] + client = storage.Client() + bucket = client.get_bucket(bucket_name) + + for file_name in file_names: + blob_path = f"{unzipped_files_path}/{file_name}" + blob = bucket.blob(blob_path) + local_path = os.path.join(local_dir, file_name) + blob.download_to_filename(local_path) + logging.info("Downloaded %s to %s", blob_path, local_path) + + +def upload_files_to_gcs( + bucket_name: str, + source_dir: str, + file_names: list[str], + feed_stable_id: str, + dataset_stable_id: str, +): + client = storage.Client() + bucket = client.get_bucket(bucket_name) + dest_prefix = f"{feed_stable_id}/{dataset_stable_id}/pmtiles" + + # Delete existing files in the destination folder + blobs_to_delete = list(bucket.list_blobs(prefix=dest_prefix + "/")) + for blob in blobs_to_delete: + blob.delete() + logging.info("Deleted existing blob: %s", blob.name) + + # Upload new files + for file_name in file_names: + file_path = os.path.join(source_dir, file_name) + if not os.path.exists(file_path): + logging.warning("File not found: %s", file_path) + continue + blob_path = f"{dest_prefix}/{file_name}" + blob = bucket.blob(blob_path) + blob.upload_from_filename(file_path) + logging.info("Uploaded %s to gs://%s/%s", file_path, bucket_name, blob_path) diff --git a/functions-python/tasks_executor/src/tasks/pmtiles_builder/create_routes_geojson.py b/functions-python/tasks_executor/src/tasks/pmtiles_builder/create_routes_geojson.py new file mode 100644 index 000000000..c12e5079f --- /dev/null +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/create_routes_geojson.py @@ -0,0 +1,111 @@ +# Here’s how to integrate the indexed shape lookup into your full GeoJSON route creation script. +# This version loads shapes_index.pkl once, uses it for fast shape lookups, and prints progress. +import csv +import json +import pickle +import logging + + +def read_csv(filename): + print(f"Loading {filename}...") + with open(filename, newline="", encoding="utf-8") as f: + return list(csv.DictReader(f)) + + +def get_shape_points(shape_id, index, local_dir): + points = [] + shapes_file = f"{local_dir}/shapes.txt" + with open(shapes_file, "r", encoding="utf-8") 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]) + print(f" Found {len(points)} points for shape_id {shape_id}") + return [pt[:2] for pt in points] + + +def create_routes_geojson(local_dir): + logging.info("Loading shapes_index.pkl...") + shapes_index_file = f"{local_dir}/shapes_index.pkl" + shapes_file = f"{local_dir}/shapes.txt" + trips_file = f"{local_dir}/trips.txt" + routes_file = f"{local_dir}/routes.txt" + stops_file = f"{local_dir}/stops.txt" + stop_times_file = f"{local_dir}/stop_times.txt" + with open(shapes_index_file, "rb") as idxf: + shapes_index = pickle.load(idxf) + logging.info(f"Loaded index for {len(shapes_index)} shape_ids.") + + # Read header columns for shapes.txt (needed for manual parsing) + with open(shapes_file, "r", encoding="utf-8") as f: + header = f.readline() + shapes_columns = next(csv.reader([header])) + shapes_index["columns"] = shapes_columns + + routes = {r["route_id"]: r for r in read_csv(routes_file)} + logging.info(f"Loaded {len(routes)} routes.") + + trips = list(read_csv(trips_file)) + logging.info(f"Loaded {len(trips)} trips.") + + stops = { + s["stop_id"]: (float(s["stop_lon"]), float(s["stop_lat"])) + for s in read_csv(stops_file) + } + logging.info(f"Loaded {len(stops)} stops.") + + stop_times_by_trip = {} + print("Grouping stop_times by trip_id...") + with open(stop_times_file, newline="", encoding="utf-8") as f: + reader = csv.DictReader(f) + for row in reader: + stop_times_by_trip.setdefault(row["trip_id"], []).append(row) + logging.info(f"Grouped stop_times for {len(stop_times_by_trip)} trips.") + + features = [] + for i, (route_id, route) in enumerate(routes.items(), 1): + if i % 100 == 0 or i == 1: + logging.info( + f"Processing route {i}/{len(routes)} (route_id: {route_id})..." + ) + trip = next((t for t in trips if t["route_id"] == route_id), None) + if not trip: + logging.info(f" No trip found for route_id {route_id}, skipping.") + continue + coordinates = [] + if "shape_id" in trip and trip["shape_id"]: + logging.info(f" Using shape_id {trip['shape_id']} for route_id {route_id}") + coordinates = get_shape_points(trip["shape_id"], shapes_index, local_dir) + if not coordinates: + trip_stop_times = stop_times_by_trip.get(trip["trip_id"], []) + trip_stop_times.sort(key=lambda x: int(x["stop_sequence"])) + coordinates = [ + stops[st["stop_id"]] for st in trip_stop_times if st["stop_id"] in stops + ] + logging.info( + f" Used {len(coordinates)} stop coordinates for route_id {route_id}" + ) + if not coordinates: + logging.info(f" No coordinates found for route_id {route_id}, skipping.") + continue + features.append( + { + "type": "Feature", + "properties": {k: route[k] for k in route}, + "geometry": {"type": "LineString", "coordinates": coordinates}, + } + ) + + logging.info(f"Writing {len(features)} features to routes-output.geojson...") + routes_geojson = f"{local_dir}/routes-output.geojson" + with open(routes_geojson, "w", encoding="utf-8") as f: + json.dump({"type": "FeatureCollection", "features": features}, f) + logging.info("Done.") diff --git a/functions-python/tasks_executor/src/tasks/pmtiles_builder/create_shapes_index.py b/functions-python/tasks_executor/src/tasks/pmtiles_builder/create_shapes_index.py new file mode 100644 index 000000000..112d19e7c --- /dev/null +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/create_shapes_index.py @@ -0,0 +1,39 @@ +# Yes, indexing shapes.txt can greatly speed up lookups. You can preprocess shapes.txt once to build an on-disk index +# mapping each shape_id to its file offsets. Then, for each needed shape_id, seek directly to its entries. +# +# Explanation: +# +# +# First, scan shapes.txt and record the byte offsets for each shape_id in an index (e.g. a pickle or JSON file). +# When processing, use the index to seek and read only the relevant lines for each shape_id. +# Here’s a two-step approach: +import csv +import pickle +import logging + + +def create_shapes_index(local_dir): + index = {} + shapes = f"{local_dir}/shapes.txt" + outfile = f"{local_dir}/shapes_index.pkl" + with open(shapes, "r", encoding="utf-8") as f: + header = f.readline() + columns = next(csv.reader([header])) + 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"] + index.setdefault(sid, []).append(pos) + count += 1 + if count % 1000000 == 0: + logging.debug(f"Indexed {count} lines so far...") + + logging.info(f"Total indexed lines: {count}") + logging.info(f"Total unique shape_ids: {len(index)}") + with open(outfile, "wb") as idxf: + pickle.dump(index, idxf) + logging.info("Indexing complete. Saved to shapes_index.pkl.") diff --git a/functions-python/tasks_executor/src/tasks/pmtiles_builder/run_tippecanoe.py b/functions-python/tasks_executor/src/tasks/pmtiles_builder/run_tippecanoe.py new file mode 100644 index 000000000..cb33fe119 --- /dev/null +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/run_tippecanoe.py @@ -0,0 +1,19 @@ +import logging +import subprocess + + +def run_tippecanoe(input_file, output_file, local_dir="./unzipped"): + cmd = [ + "tippecanoe", + "-o", + f"{local_dir}/{input_file}", + "--force", + "--no-tile-size-limit", + "-zg", + f"{local_dir}/{output_file}", + ] + try: + subprocess.run(cmd, check=True) + logging.info("Tippecanoe command executed successfully.") + except subprocess.CalledProcessError as e: + logging.info(f"Error running tippecanoe: {e}") From d10552a95f6f91be69b89c2adc44b30e7984d383 Mon Sep 17 00:00:00 2001 From: jcpitre Date: Tue, 29 Jul 2025 06:49:35 -0400 Subject: [PATCH 02/14] Reformatted. --- .../tasks/pmtiles_builder/build_pmtiles.py | 626 +++++++++++++----- .../pmtiles_builder/create_routes_geojson.py | 111 ---- .../pmtiles_builder/create_shapes_index.py | 39 -- .../tasks/pmtiles_builder/run_tippecanoe.py | 19 - 4 files changed, 456 insertions(+), 339 deletions(-) delete mode 100644 functions-python/tasks_executor/src/tasks/pmtiles_builder/create_routes_geojson.py delete mode 100644 functions-python/tasks_executor/src/tasks/pmtiles_builder/create_shapes_index.py delete mode 100644 functions-python/tasks_executor/src/tasks/pmtiles_builder/run_tippecanoe.py 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 index dc1f809d6..aa6dd2e9c 100644 --- a/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py @@ -13,187 +13,473 @@ # 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 pickle import shutil import subprocess -import sys - +from logging import DEBUG, WARNING, INFO, ERROR from google.cloud import storage -from sqlalchemy.orm import Session - -sys.path.append(os.path.dirname(os.path.abspath(__file__))) # noqa: E402 -from create_shapes_index import create_shapes_index # noqa: E402 -from create_routes_geojson import create_routes_geojson # noqa: E402 -from run_tippecanoe import run_tippecanoe # noqa: E402 +# Files are stored locally to be able to run tippecanoe on them. This is the directory +local_dir = "./workdir" def build_pmtiles_handler(payload) -> dict: """ - Rebuilds missing validation reports for GTFS datasets. - This function processes datasets with missing validation reports using the GTFS validator workflow. - The payload structure is: - { - "dry_run": bool, # [optional] If True, do not execute the workflow - "feed_stable_id": int, # [optional] Filter datasets older than this number of days(default: 14 days ago) - "dataset_stable_id": list[str] # [optional] Filter datasets by status(in) - } - Args: - payload (dict): The payload containing the task details. - Returns: - str: A message indicating the result of the operation with the total_processed datasets. - """ - dry_run: bool - ( - dry_run, - feed_stable_id, - dataset_stable_id, - ) = get_parameters(payload) - - return build_pmtiles( - dry_run=dry_run, - feed_stable_id=feed_stable_id, - dataset_stable_id=dataset_stable_id, - ) - - -def build_pmtiles( - dry_run: bool = True, - feed_stable_id: str | None = None, - dataset_stable_id: str | None = None, - db_session: Session | None = None, -) -> dict: + Entrypoint for building PMTiles files from a GTFS dataset. """ - Rebuilds missing validation reports for GTFS datasets. - - Args: - validator_endpoint: Validator endpoint URL - dry_run (bool): dry run flag. If True, do not execute the workflow. Default: True - filter_after_in_days (int): Filter the datasets older than this number of days. Default: 14 days ago - filter_statuses: [optional] Filter datasets by status(in). Default: None - prod_env (bool): True if target environment is production, false otherwise. Default: False - db_session: DB session - - Returns: - flask.Response: A response with message and total_processed datasets. - """ - bucket_name = os.getenv("DATASETS_BUCKET_NAME") - if not bucket_name: - return {"error": "DATASETS_BUCKET_NAME environment variable is not defined."} - - if not feed_stable_id or not dataset_stable_id: - return {"error": "Both feed_stable_id and dataset_stable_id must be defined."} - - if feed_stable_id not in dataset_stable_id: - return {"error": "feed_stable_id must be a substring of dataset_stable_id."} - - logging.info( - "Starting PMTiles build for feed %s and dataset %s on bucket %s", - feed_stable_id, - dataset_stable_id, - bucket_name, - ) - unzipped_files_path = f"{feed_stable_id}/{dataset_stable_id}/extracted" - - logging.info("Initializing storage client") - bucket = storage.Client().get_bucket(bucket_name) - logging.info("Getting blobs with prefix: %s", unzipped_files_path) - blobs = list(bucket.list_blobs(prefix=unzipped_files_path)) - logging.info("Found %d blobs", len(blobs)) - if not blobs: - return { - "error": f"Directory '{unzipped_files_path}' does not exist in bucket '{bucket_name}'." - } - - local_dir = "./unzipped" - if os.path.exists(local_dir): - shutil.rmtree(local_dir) - os.makedirs(local_dir, exist_ok=True) - download_files_from_gcs(bucket_name, unzipped_files_path, local_dir) - - create_shapes_index(local_dir) - create_routes_geojson(local_dir) - logging.info(os.getcwd()) - - result = subprocess.run(["which", "tippecanoe"], capture_output=True, text=True) - logging.info("REsult of which command: %s", result.stdout.strip()) - - run_tippecanoe("routes.pmtiles", "routes-output.geojson", local_dir) - - upload_files_to_gcs( - bucket_name, local_dir, ["routes.pmtiles"], feed_stable_id, dataset_stable_id - ) - - result = subprocess.run( - ["ls", "-l", "-R", local_dir], capture_output=True, text=True - ) - logging.info("Files created:\n%s", result.stdout.strip()) - return { - "message": f"Directory '{unzipped_files_path}' exists in bucket '{bucket_name}'." - } - - -def get_parameters(payload): + try: + feed_stable_id, dataset_stable_id = PmtilesBuilder._get_parameters(payload) + builder = PmtilesBuilder( + feed_stable_id=feed_stable_id, dataset_stable_id=dataset_stable_id + ) + return builder.build_pmtiles() + except Exception as e: + return {"error": f"Failed to start PMTiles build: {e}"} + + +class PmtilesBuilder: """ - Get parameters from the payload and environment variables. + Orchestrates the end-to-end process of generating PMTiles files from GTFS datasets. - Args: - payload (dict): dictionary containing the payload data. - Returns: - dict: dict with: dry_run, filter_after_in_days, filter_statuses, prod_env, validator_endpoint parameters + 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. """ - dry_run = payload.get("dry_run", True) - dry_run = dry_run if isinstance(dry_run, bool) else str(dry_run).lower() == "true" - feed_stable_id = payload.get("feed_stable_id", None) - dataset_stable_id = payload.get("dataset_stable_id", None) - - return dry_run, feed_stable_id, dataset_stable_id - - -def download_files_from_gcs(bucket_name, unzipped_files_path, local_dir): - file_names = [ - "routes.txt", - "shapes.txt", - "stop_times.txt", - "trips.txt", - "stops.txt", - ] - client = storage.Client() - bucket = client.get_bucket(bucket_name) - - for file_name in file_names: - blob_path = f"{unzipped_files_path}/{file_name}" - blob = bucket.blob(blob_path) - local_path = os.path.join(local_dir, file_name) - blob.download_to_filename(local_path) - logging.info("Downloaded %s to %s", blob_path, local_path) - - -def upload_files_to_gcs( - bucket_name: str, - source_dir: str, - file_names: list[str], - feed_stable_id: str, - dataset_stable_id: str, -): - client = storage.Client() - bucket = client.get_bucket(bucket_name) - dest_prefix = f"{feed_stable_id}/{dataset_stable_id}/pmtiles" - - # Delete existing files in the destination folder - blobs_to_delete = list(bucket.list_blobs(prefix=dest_prefix + "/")) - for blob in blobs_to_delete: - blob.delete() - logging.info("Deleted existing blob: %s", blob.name) - - # Upload new files - for file_name in file_names: - file_path = os.path.join(source_dir, file_name) - if not os.path.exists(file_path): - logging.warning("File not found: %s", file_path) - continue - blob_path = f"{dest_prefix}/{file_name}" - blob = bucket.blob(blob_path) - blob.upload_from_filename(file_path) - logging.info("Uploaded %s to gs://%s/%s", file_path, bucket_name, blob_path) + + def __init__( + self, + feed_stable_id: str | None = None, + dataset_stable_id: str | None = None, + ): + 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") + + def _log(self, level, msg, *args): + logger = logging.getLogger() + if not logger.isEnabledFor(level): + return + formatted_msg = msg % args if args else msg + logger.log(level, "[%s] %s", self.dataset_stable_id, formatted_msg) + + @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) -> dict: + try: + if not self.bucket_name: + return { + "error": "DATASETS_BUCKET_NAME environment variable is not defined." + } + if not self.feed_stable_id or not self.dataset_stable_id: + return { + "error": "Both feed_stable_id and dataset_stable_id must be defined." + } + if self.feed_stable_id not in self.dataset_stable_id: + return { + "error": ( + "feed_stable_id %s is not a prefix of dataset_stable_id %s." + % (self.feed_stable_id, self.dataset_stable_id) + ) + } + + self._log( + INFO, "Starting PMTiles build for dataset %s", self.dataset_stable_id + ) + unzipped_files_path = ( + f"{self.feed_stable_id}/{self.dataset_stable_id}/extracted" + ) + + self._download_files_from_gcs(unzipped_files_path) + + self._create_shapes_index() + + 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) + + # List files in the relevant bucket folder instead of local_dir + + if logging.getLogger().isEnabledFor(DEBUG): + gcs_prefix = f"{self.feed_stable_id}/{self.dataset_stable_id}/pmtiles/" + try: # We don`t want an error here to abort the whole pmtiles operation. + blobs = list(self.bucket.list_blobs(prefix=gcs_prefix)) + file_list = "\n".join( + f"{blob.name} ({blob.size} bytes)" for blob in blobs + ) + self._log(DEBUG, "GCS files in %s:\n%s", gcs_prefix, file_list) + except Exception as e: + self._log( + ERROR, + "Could not list files in bucket %s for path %s: %s", + self.bucket_name, + gcs_prefix, + e, + ) + + return { + "message": f"Pmtiles successfully created for dataset {self.dataset_stable_id}." + } + except Exception as e: + logging.exception( + "Failed to build PMTiles for dataset %s", self.dataset_stable_id + ) + return { + "error": f"Failed to build PMTiles for dataset {self.dataset_stable_id}: {e}" + } + + def _download_files_from_gcs(self, unzipped_files_path): + self._log( + INFO, + "Downloading dataset from GCS bucket %s, directory %s", + self.bucket_name, + unzipped_files_path, + ) + try: + self._log(DEBUG, "Initializing storage client") + self.bucket = storage.Client().get_bucket(self.bucket_name) + self._log(DEBUG, "Getting blobs with prefix: %s", unzipped_files_path) + blobs = list(self.bucket.list_blobs(prefix=unzipped_files_path)) + self._log(DEBUG, "Found %d blobs", len(blobs)) + if not blobs: + raise { + f"Directory '{unzipped_files_path}' does not exist or is empty in bucket '{self.bucket_name}'." + } + + if os.path.exists(local_dir): + shutil.rmtree(local_dir) + os.makedirs(local_dir, exist_ok=True) + file_names = [ + "routes.txt", + "shapes.txt", + "stop_times.txt", + "trips.txt", + "stops.txt", + ] + for file_name in file_names: + blob_path = f"{unzipped_files_path}/{file_name}" + blob = self.bucket.blob(blob_path) + local_path = os.path.join(local_dir, file_name) + blob.download_to_filename(local_path) + self._log(DEBUG, "Downloaded %s to %s", blob_path, local_path) + return + except Exception as e: + raise Exception(f"Failed to download files from GCS: {e}") from e + + def _upload_files_to_gcs(self, file_to_upload): + dest_prefix = f"{self.feed_stable_id}/{self.dataset_stable_id}/pmtiles" + self._log( + 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._log(DEBUG, "Deleted existing blob: %s", blob.name) + for file_name in file_to_upload: + file_path = os.path.join(local_dir, file_name) + if not os.path.exists(file_path): + self._log(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._log( + 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): + self._log(INFO, "Creating shapes index") + try: + index = {} + shapes = f"{local_dir}/shapes.txt" + outfile = f"{local_dir}/shapes_index.pkl" + with open(shapes, "r", encoding="utf-8") as f: + header = f.readline() + columns = next(csv.reader([header])) + 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"] + index.setdefault(sid, []).append(pos) + count += 1 + if count % 1000000 == 0: + self._log(DEBUG, "Indexed %d lines so far...", count) + self._log(DEBUG, "Total indexed lines: %d, count") + self._log(DEBUG, "Total unique shape_ids: %d", len(index)) + with open(outfile, "wb") as idxf: + pickle.dump(index, idxf) + except Exception as e: + raise Exception(f"Failed to create shapes index: {e}") from e + + def _read_csv(self, filename): + try: + self._log(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._log(DEBUG, "Getting shape points for shape_id %s", shape_id) + try: + points = [] + shapes_file = f"{local_dir}/shapes.txt" + with open(shapes_file, "r", encoding="utf-8") 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._log(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): + self._log(INFO, "Creating routes geojson") + try: + self._log(DEBUG, "Loading shapes_index.pkl...") + shapes_index_file = f"{local_dir}/shapes_index.pkl" + shapes_file = f"{local_dir}/shapes.txt" + trips_file = f"{local_dir}/trips.txt" + routes_file = f"{local_dir}/routes.txt" + stops_file = f"{local_dir}/stops.txt" + stop_times_file = f"{local_dir}/stop_times.txt" + with open(shapes_index_file, "rb") as idxf: + shapes_index = pickle.load(idxf) + self._log(DEBUG, "Loaded index for %d shape_ids.", len(shapes_index)) + + with open(shapes_file, "r", encoding="utf-8") as f: + header = f.readline() + shapes_columns = next(csv.reader([header])) + shapes_index["columns"] = shapes_columns + + routes = {r["route_id"]: r for r in self._read_csv(routes_file)} + self._log(DEBUG, "Loaded %d routes.", len(routes)) + + trips = list(self._read_csv(trips_file)) + self._log(DEBUG, "Loaded %d trips.", len(trips)) + + stops = { + s["stop_id"]: (float(s["stop_lon"]), float(s["stop_lat"])) + for s in self._read_csv(stops_file) + } + self._log(DEBUG, "Loaded %d stops.", len(stops)) + + stop_times_by_trip = {} + self._log( + DEBUG, + "Grouping stop_times by trip_id for dataset %s", + self.dataset_stable_id, + ) + with open(stop_times_file, newline="", encoding="utf-8") as f: + reader = csv.DictReader(f) + for row in reader: + stop_times_by_trip.setdefault(row["trip_id"], []).append(row) + self._log( + DEBUG, "Grouped stop_times for %d trips.", len(stop_times_by_trip) + ) + + features = [] + for i, (route_id, route) in enumerate(routes.items(), 1): + if i % 100 == 0 or i == 1: + self._log( + DEBUG, + "Processing route %d/%d} (route_id: %s...", + i, + len(routes), + route_id, + ) + trip = next((t for t in trips if t["route_id"] == route_id), None) + if not trip: + self._log( + INFO, " No trip found for route_id %s, skipping.", route_id + ) + continue + coordinates = [] + if "shape_id" in trip and trip["shape_id"]: + self._log( + DEBUG, + " Using shape_id %s for route_id %s", + trip["shape_id"], + route_id, + ) + coordinates = self._get_shape_points(trip["shape_id"], shapes_index) + if isinstance(coordinates, dict) and "error" in coordinates: + raise Exception( + f"Error getting shape points for shape_id {trip['shape_id']}: {coordinates['error']}" + ) + if not coordinates: + trip_stop_times = stop_times_by_trip.get(trip["trip_id"], []) + trip_stop_times.sort(key=lambda x: int(x["stop_sequence"])) + coordinates = [ + stops[st["stop_id"]] + for st in trip_stop_times + if st["stop_id"] in stops + ] + self._log( + DEBUG, + " Used %d stop coordinates for route_id %s", + len(coordinates), + route_id, + ) + if not coordinates: + self._log( + INFO, + " No coordinates found for route_id %s, skipping.", + route_id, + ) + continue + features.append( + { + "type": "Feature", + "properties": {k: route[k] for k in route}, + "geometry": {"type": "LineString", "coordinates": coordinates}, + } + ) + + self._log( + DEBUG, "Writing %d features to routes-output.geojson...", len(features) + ) + routes_geojson = f"{local_dir}/routes-output.geojson" + with open(routes_geojson, "w", encoding="utf-8") as f: + json.dump({"type": "FeatureCollection", "features": features}, f) + except Exception as e: + raise Exception(f"Failed to create routes GeoJSON: {e}") from e + + def _run_tippecanoe(self, input_file, output_file): + self._log(INFO, "Running tippecanoe for input file %s", input_file) + try: + cmd = [ + "tippecanoe", + "-o", + f"{local_dir}/{output_file}", + "--force", + "--no-tile-size-limit", + "-zg", + f"{local_dir}/{input_file}", + ] + self._log(DEBUG, "Running command: %s", " ".join(cmd)) + result = subprocess.run(cmd, capture_output=True, text=True) + if result.stdout: + self._log(DEBUG, "Tippecanoe output:\n%s", result.stdout) + if result.returncode != 0: + self._log(ERROR, "Tippecanoe error:\n%s", result.stderr) + raise Exception(f"Tippecanoe failed with exit code {result.returncode}") + self._log(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._log(INFO, "Creating stops geojson...") + try: + stops = self._read_csv(f"{local_dir}/stops.txt") + if isinstance(stops, dict) and "error" in stops: + return stops + self._log(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._log( + 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 = f"{local_dir}/stops-output.geojson" + + self._log( + 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._log(INFO, "Creating routes json...") + try: + routes = [] + with open(f"{local_dir}/routes.txt", 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', '3')}", + } + routes.append(route) + + with open(f"{local_dir}/routes.json", "w", encoding="utf-8") as f: + json.dump(routes, f, ensure_ascii=False, indent=4) + + self._log(DEBUG, "Converted %d routes to routes.json.", len(routes)) + except Exception as e: + raise Exception("Failed to create routes JSON for dataset: {e}") from e diff --git a/functions-python/tasks_executor/src/tasks/pmtiles_builder/create_routes_geojson.py b/functions-python/tasks_executor/src/tasks/pmtiles_builder/create_routes_geojson.py deleted file mode 100644 index c12e5079f..000000000 --- a/functions-python/tasks_executor/src/tasks/pmtiles_builder/create_routes_geojson.py +++ /dev/null @@ -1,111 +0,0 @@ -# Here’s how to integrate the indexed shape lookup into your full GeoJSON route creation script. -# This version loads shapes_index.pkl once, uses it for fast shape lookups, and prints progress. -import csv -import json -import pickle -import logging - - -def read_csv(filename): - print(f"Loading {filename}...") - with open(filename, newline="", encoding="utf-8") as f: - return list(csv.DictReader(f)) - - -def get_shape_points(shape_id, index, local_dir): - points = [] - shapes_file = f"{local_dir}/shapes.txt" - with open(shapes_file, "r", encoding="utf-8") 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]) - print(f" Found {len(points)} points for shape_id {shape_id}") - return [pt[:2] for pt in points] - - -def create_routes_geojson(local_dir): - logging.info("Loading shapes_index.pkl...") - shapes_index_file = f"{local_dir}/shapes_index.pkl" - shapes_file = f"{local_dir}/shapes.txt" - trips_file = f"{local_dir}/trips.txt" - routes_file = f"{local_dir}/routes.txt" - stops_file = f"{local_dir}/stops.txt" - stop_times_file = f"{local_dir}/stop_times.txt" - with open(shapes_index_file, "rb") as idxf: - shapes_index = pickle.load(idxf) - logging.info(f"Loaded index for {len(shapes_index)} shape_ids.") - - # Read header columns for shapes.txt (needed for manual parsing) - with open(shapes_file, "r", encoding="utf-8") as f: - header = f.readline() - shapes_columns = next(csv.reader([header])) - shapes_index["columns"] = shapes_columns - - routes = {r["route_id"]: r for r in read_csv(routes_file)} - logging.info(f"Loaded {len(routes)} routes.") - - trips = list(read_csv(trips_file)) - logging.info(f"Loaded {len(trips)} trips.") - - stops = { - s["stop_id"]: (float(s["stop_lon"]), float(s["stop_lat"])) - for s in read_csv(stops_file) - } - logging.info(f"Loaded {len(stops)} stops.") - - stop_times_by_trip = {} - print("Grouping stop_times by trip_id...") - with open(stop_times_file, newline="", encoding="utf-8") as f: - reader = csv.DictReader(f) - for row in reader: - stop_times_by_trip.setdefault(row["trip_id"], []).append(row) - logging.info(f"Grouped stop_times for {len(stop_times_by_trip)} trips.") - - features = [] - for i, (route_id, route) in enumerate(routes.items(), 1): - if i % 100 == 0 or i == 1: - logging.info( - f"Processing route {i}/{len(routes)} (route_id: {route_id})..." - ) - trip = next((t for t in trips if t["route_id"] == route_id), None) - if not trip: - logging.info(f" No trip found for route_id {route_id}, skipping.") - continue - coordinates = [] - if "shape_id" in trip and trip["shape_id"]: - logging.info(f" Using shape_id {trip['shape_id']} for route_id {route_id}") - coordinates = get_shape_points(trip["shape_id"], shapes_index, local_dir) - if not coordinates: - trip_stop_times = stop_times_by_trip.get(trip["trip_id"], []) - trip_stop_times.sort(key=lambda x: int(x["stop_sequence"])) - coordinates = [ - stops[st["stop_id"]] for st in trip_stop_times if st["stop_id"] in stops - ] - logging.info( - f" Used {len(coordinates)} stop coordinates for route_id {route_id}" - ) - if not coordinates: - logging.info(f" No coordinates found for route_id {route_id}, skipping.") - continue - features.append( - { - "type": "Feature", - "properties": {k: route[k] for k in route}, - "geometry": {"type": "LineString", "coordinates": coordinates}, - } - ) - - logging.info(f"Writing {len(features)} features to routes-output.geojson...") - routes_geojson = f"{local_dir}/routes-output.geojson" - with open(routes_geojson, "w", encoding="utf-8") as f: - json.dump({"type": "FeatureCollection", "features": features}, f) - logging.info("Done.") diff --git a/functions-python/tasks_executor/src/tasks/pmtiles_builder/create_shapes_index.py b/functions-python/tasks_executor/src/tasks/pmtiles_builder/create_shapes_index.py deleted file mode 100644 index 112d19e7c..000000000 --- a/functions-python/tasks_executor/src/tasks/pmtiles_builder/create_shapes_index.py +++ /dev/null @@ -1,39 +0,0 @@ -# Yes, indexing shapes.txt can greatly speed up lookups. You can preprocess shapes.txt once to build an on-disk index -# mapping each shape_id to its file offsets. Then, for each needed shape_id, seek directly to its entries. -# -# Explanation: -# -# -# First, scan shapes.txt and record the byte offsets for each shape_id in an index (e.g. a pickle or JSON file). -# When processing, use the index to seek and read only the relevant lines for each shape_id. -# Here’s a two-step approach: -import csv -import pickle -import logging - - -def create_shapes_index(local_dir): - index = {} - shapes = f"{local_dir}/shapes.txt" - outfile = f"{local_dir}/shapes_index.pkl" - with open(shapes, "r", encoding="utf-8") as f: - header = f.readline() - columns = next(csv.reader([header])) - 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"] - index.setdefault(sid, []).append(pos) - count += 1 - if count % 1000000 == 0: - logging.debug(f"Indexed {count} lines so far...") - - logging.info(f"Total indexed lines: {count}") - logging.info(f"Total unique shape_ids: {len(index)}") - with open(outfile, "wb") as idxf: - pickle.dump(index, idxf) - logging.info("Indexing complete. Saved to shapes_index.pkl.") diff --git a/functions-python/tasks_executor/src/tasks/pmtiles_builder/run_tippecanoe.py b/functions-python/tasks_executor/src/tasks/pmtiles_builder/run_tippecanoe.py deleted file mode 100644 index cb33fe119..000000000 --- a/functions-python/tasks_executor/src/tasks/pmtiles_builder/run_tippecanoe.py +++ /dev/null @@ -1,19 +0,0 @@ -import logging -import subprocess - - -def run_tippecanoe(input_file, output_file, local_dir="./unzipped"): - cmd = [ - "tippecanoe", - "-o", - f"{local_dir}/{input_file}", - "--force", - "--no-tile-size-limit", - "-zg", - f"{local_dir}/{output_file}", - ] - try: - subprocess.run(cmd, check=True) - logging.info("Tippecanoe command executed successfully.") - except subprocess.CalledProcessError as e: - logging.info(f"Error running tippecanoe: {e}") From db34ac3cf33d22e22b519ef5fe86615cc3e590c9 Mon Sep 17 00:00:00 2001 From: jcpitre Date: Tue, 29 Jul 2025 07:40:25 -0400 Subject: [PATCH 03/14] Correction --- .../tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 index aa6dd2e9c..5129fbbf5 100644 --- a/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py @@ -166,9 +166,9 @@ def _download_files_from_gcs(self, unzipped_files_path): blobs = list(self.bucket.list_blobs(prefix=unzipped_files_path)) self._log(DEBUG, "Found %d blobs", len(blobs)) if not blobs: - raise { + raise Exception( f"Directory '{unzipped_files_path}' does not exist or is empty in bucket '{self.bucket_name}'." - } + ) if os.path.exists(local_dir): shutil.rmtree(local_dir) From d88f11ea1aa3111e8ca237ec816e76c3b162f6c9 Mon Sep 17 00:00:00 2001 From: jcpitre Date: Tue, 29 Jul 2025 11:05:15 -0400 Subject: [PATCH 04/14] Added tests --- .../pmtiles_builder/test_build_pmtiles.py | 447 ++++++++++++++++++ 1 file changed, 447 insertions(+) create mode 100644 functions-python/tasks_executor/tests/tasks/pmtiles_builder/test_build_pmtiles.py 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..49c898cfa --- /dev/null +++ b/functions-python/tasks_executor/tests/tasks/pmtiles_builder/test_build_pmtiles.py @@ -0,0 +1,447 @@ +import csv +import json +import logging +import pickle +import tempfile +import unittest +from contextlib import contextmanager +from unittest.mock import patch, MagicMock +import os + +from tasks.pmtiles_builder.build_pmtiles import ( + PmtilesBuilder, + local_dir, + 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 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.storage.Client") + def test_download_files_from_gcs_success(self, mock_client): + mock_bucket = MagicMock() + mock_blob = MagicMock() + mock_blob.download_to_filename = MagicMock() + mock_bucket.list_blobs.return_value = [mock_blob] * 5 + mock_client.return_value.get_bucket.return_value = mock_bucket + + with patch("os.path.exists", return_value=True), patch("shutil.rmtree"), patch( + "os.makedirs" + ): + self.builder._download_files_from_gcs("some/path") + self.assertTrue(mock_bucket.list_blobs.called) + + @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): + 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 = [] + result = self.builder.build_pmtiles() + self.assertIn("message", result) + + 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"] + + @patch( + "tasks.pmtiles_builder.build_pmtiles.PmtilesBuilder._download_files_from_gcs" + ) + @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_calls_create_shapes_index( + self, + mock_upload, + mock_routes_json, + mock_stops_geojson, + mock_run_tippecanoe, + mock_routes_geojson, + mock_download, + ): + self.builder.bucket = MagicMock() + self.builder.bucket.list_blobs.return_value = [] + # Create minimal shapes.txt in local_dir + os.makedirs(local_dir, exist_ok=True) + shapes_path = os.path.join(local_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\ns1,45.0,-73.0,1\n" + ) + result = self.builder.build_pmtiles() + self.assertIn("message", result) + + def test_get_shape_points(self): + # Prepare shapes.txt + os.makedirs(local_dir, exist_ok=True) + shapes_path = os.path.join(local_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") + + # 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 + os.makedirs(local_dir, exist_ok=True) + with open(os.path.join(local_dir, "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(os.path.join(local_dir, "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(os.path.join(local_dir, "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( + os.path.join(local_dir, "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(os.path.join(local_dir, "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") + # Create shapes_index.pkl + shapes_index = { + "s1": [ + len("shape_id,shape_pt_lat,shape_pt_lon,shape_pt_sequence\n"), + len( + "shape_id,shape_pt_lat,shape_pt_lon,shape_pt_sequence\ns1,45.0,-73.0,1\n" + ), + ], + "columns": [ + "shape_id", + "shape_pt_lat", + "shape_pt_lon", + "shape_pt_sequence", + ], + } + with open(os.path.join(local_dir, "shapes_index.pkl"), "wb") as f: + pickle.dump(shapes_index, f) + + # Call the method + self.builder._create_routes_geojson() + + # Assert output file exists and is valid GeoJSON + output_path = os.path.join(local_dir, "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 + os.makedirs(local_dir, exist_ok=True) + with open(os.path.join(local_dir, "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 = os.path.join(local_dir, "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 + os.makedirs(local_dir, exist_ok=True) + with open(os.path.join(local_dir, "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 = os.path.join(local_dir, "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 + if os.path.exists(os.path.join(local_dir, "routes.txt")): + os.remove(os.path.join(local_dir, "routes.txt")) + 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(): + result = self.builder.build_pmtiles() + self.assertIn("error", result) + self.assertIn("Failed to build PMTiles for dataset", result["error"]) + + def test_upload_files_to_gcs_missing_file(self): + builder = PmtilesBuilder(feed_stable_id="foo", dataset_stable_id="foo_bar") + builder.bucket = MagicMock() + builder.bucket.blob.return_value = MagicMock() + builder.bucket.list_blobs.return_value = [] + + # Ensure the file does not exist + missing_file = "notfound.pmtiles" + if os.path.exists(os.path.join(local_dir, missing_file)): + os.remove(os.path.join(local_dir, missing_file)) + + with patch( + "os.path.exists", + side_effect=lambda path: False if missing_file in path else True, + ), patch.object(builder, "_log") as mock_log: + builder._upload_files_to_gcs([missing_file]) + mock_log.assert_any_call( + logging.WARNING, + "File not found: %s", + os.path.join(local_dir, missing_file), + ) + + def test_create_routes_geojson_fallback_to_stop_coordinates(self): + # Prepare minimal GTFS files with no shape_id for the trip + os.makedirs(local_dir, exist_ok=True) + with open(os.path.join(local_dir, "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(os.path.join(local_dir, "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(os.path.join(local_dir, "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( + os.path.join(local_dir, "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") + # shapes.txt and shapes_index.pkl are still needed but not used in this case + with open(os.path.join(local_dir, "shapes.txt"), "w", encoding="utf-8") as f: + f.write("shape_id,shape_pt_lat,shape_pt_lon,shape_pt_sequence\n") + with open(os.path.join(local_dir, "shapes_index.pkl"), "wb") as f: + pickle.dump( + { + "columns": [ + "shape_id", + "shape_pt_lat", + "shape_pt_lon", + "shape_pt_sequence", + ] + }, + f, + ) + + # Call the method + self.builder._create_routes_geojson() + + # Assert output file exists and contains the expected coordinates + output_path = os.path.join(local_dir, "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): + # Prepare stops.txt with one valid and one invalid stop + os.makedirs(local_dir, exist_ok=True) + with open(os.path.join(local_dir, "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") # valid + f.write("stop2,not_a_lat,-73.1\n") # invalid lat + + with patch.object(self.builder, "_log") as mock_log: + self.builder._create_stops_geojson() + # Check that the log was called for the invalid stop + mock_log.assert_any_call( + logging.INFO, "Skipping stop %s: invalid coordinates", "stop2" + ) + + # Assert output file exists and only the valid stop is included + output_path = os.path.join(local_dir, "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): + # Patch local_dir to a temp directory + self.test_dir = tempfile.TemporaryDirectory() + self.old_local_dir = local_dir + self._patch_local_dir(self.test_dir.name) + os.environ["DATASETS_BUCKET_NAME"] = "test-bucket" + + # Create minimal GTFS files + 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( + os.path.join(self.test_dir.name, fname), "w", encoding="utf-8" + ) as f: + f.write(content) + + def tearDown(self): + self.test_dir.cleanup() + self._patch_local_dir(self.old_local_dir) + if "DATASETS_BUCKET_NAME" in os.environ: + del os.environ["DATASETS_BUCKET_NAME"] + + def _patch_local_dir(self, new_dir): + import tasks.pmtiles_builder.build_pmtiles as mod + + mod.local_dir = new_dir + + def test_build_pmtiles_handler_missing_bucket_env(self): + if "DATASETS_BUCKET_NAME" in os.environ: + del os.environ["DATASETS_BUCKET_NAME"] + payload = { + "feed_stable_id": "feed123", + "dataset_stable_id": "feed123_dataset456", + } + result = build_pmtiles_handler(payload) + self.assertIn("error", result) + self.assertIn("DATASETS_BUCKET_NAME", result["error"]) + + def test_build_pmtiles_handler_missing_ids(self): + os.environ["DATASETS_BUCKET_NAME"] = "test-bucket" + payload = {"feed_stable_id": "", "dataset_stable_id": ""} + result = build_pmtiles_handler(payload) + self.assertIn("error", result) + self.assertIn("must be defined", result["error"]) + + def test_build_pmtiles_handler_feed_not_prefix(self): + os.environ["DATASETS_BUCKET_NAME"] = "test-bucket" + payload = {"feed_stable_id": "foo", "dataset_stable_id": "barbaz"} + result = build_pmtiles_handler(payload) + self.assertIn("error", result) + self.assertIn("is not a prefix", result["error"]) + + +class TestPmtilesBuilderUpload(unittest.TestCase): + def test_upload_files_to_gcs(self): + builder = PmtilesBuilder(feed_stable_id="foo", dataset_stable_id="foo_bar") + builder.bucket = MagicMock() + mock_blob = MagicMock() + builder.bucket.blob.return_value = mock_blob + builder.bucket.list_blobs.return_value = [] + + # Create dummy files in local_dir for the test + os.makedirs(local_dir, exist_ok=True) + test_file = os.path.join(local_dir, "routes.pmtiles") + with open(test_file, "w") as f: + f.write("dummy data") + + builder._upload_files_to_gcs(["routes.pmtiles"]) + + builder.bucket.blob.assert_called_with("foo/foo_bar/pmtiles/routes.pmtiles") + mock_blob.upload_from_filename.assert_called_with(test_file) + + +if __name__ == "__main__": + unittest.main() From f29eb4d3955f16cb4d0bc328c4ed977ed00640c6 Mon Sep 17 00:00:00 2001 From: jcpitre Date: Tue, 29 Jul 2025 11:34:03 -0400 Subject: [PATCH 05/14] Corrected the documentation --- functions-python/tasks_executor/README.md | 27 ++++++++++++----------- 1 file changed, 14 insertions(+), 13 deletions(-) diff --git a/functions-python/tasks_executor/README.md b/functions-python/tasks_executor/README.md index 6f9654bac..2326a71dc 100644 --- a/functions-python/tasks_executor/README.md +++ b/functions-python/tasks_executor/README.md @@ -13,33 +13,34 @@ 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" + } } ``` To get the list of supported tasks use: -`` +```json { -"name": "list_tasks", -"payload": {} + "name": "list_tasks", + "payload": {} } - ``` -``` + From c7020ffb92328aa33bd769ef215c8a7d00d8a072 Mon Sep 17 00:00:00 2001 From: jcpitre Date: Tue, 29 Jul 2025 11:40:37 -0400 Subject: [PATCH 06/14] Corrected errors found by copilot --- .../src/tasks/pmtiles_builder/build_pmtiles.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 index 5129fbbf5..1f5d2593a 100644 --- a/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py @@ -242,7 +242,7 @@ def _create_shapes_index(self): count += 1 if count % 1000000 == 0: self._log(DEBUG, "Indexed %d lines so far...", count) - self._log(DEBUG, "Total indexed lines: %d, count") + self._log(DEBUG, "Total indexed lines: %d", count) self._log(DEBUG, "Total unique shape_ids: %d", len(index)) with open(outfile, "wb") as idxf: pickle.dump(index, idxf) @@ -330,7 +330,7 @@ def _create_routes_geojson(self): if i % 100 == 0 or i == 1: self._log( DEBUG, - "Processing route %d/%d} (route_id: %s...", + "Processing route %d/%d (route_id: %s)", i, len(routes), route_id, @@ -482,4 +482,4 @@ def _create_routes_json(self): self._log(DEBUG, "Converted %d routes to routes.json.", len(routes)) except Exception as e: - raise Exception("Failed to create routes JSON for dataset: {e}") from e + raise Exception(f"Failed to create routes JSON for dataset: {e}") from e From 6389ec7aad7c524a5fb6bef872514206b372440e Mon Sep 17 00:00:00 2001 From: jcpitre Date: Tue, 29 Jul 2025 12:03:26 -0400 Subject: [PATCH 07/14] Added documentation --- .../src/tasks/pmtiles_builder/README.md | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) create mode 100644 functions-python/tasks_executor/src/tasks/pmtiles_builder/README.md 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`. From cf53a53f97130e4d52dded0026292c22fc04098e Mon Sep 17 00:00:00 2001 From: jcpitre Date: Wed, 30 Jul 2025 07:56:05 -0400 Subject: [PATCH 08/14] Added env variables to the cloud function --- functions-python/tasks_executor/function_config.json | 6 +++++- infra/functions-python/main.tf | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/functions-python/tasks_executor/function_config.json b/functions-python/tasks_executor/function_config.json index eda4462f0..a33849100 100644 --- a/functions-python/tasks_executor/function_config.json +++ b/functions-python/tasks_executor/function_config.json @@ -7,7 +7,11 @@ "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" diff --git a/infra/functions-python/main.tf b/infra/functions-python/main.tf index 770546879..dc1f1de25 100644 --- a/infra/functions-python/main.tf +++ b/infra/functions-python/main.tf @@ -1222,6 +1222,7 @@ resource "google_cloudfunctions2_function" "tasks_executor" { PROJECT_ID = var.project_id ENV = var.environment PUBSUB_TOPIC_NAME = "rebuild-bounding-boxes-topic" + DATASETS_BUCKET_NAME = data.google_storage_bucket.datasets_bucket.name } available_memory = local.function_tasks_executor_config.memory timeout_seconds = local.function_tasks_executor_config.timeout From 56aa6f531769f6c343a06edbcfe1b053bc4fe9d0 Mon Sep 17 00:00:00 2001 From: jcpitre Date: Fri, 1 Aug 2025 11:56:22 -0400 Subject: [PATCH 09/14] Corrected a problem in the pmtiles builder where reading some files would loop forever. --- .../tasks_executor/function_config.json | 6 +- .../tasks/pmtiles_builder/build_pmtiles.py | 139 ++++++++---------- 2 files changed, 68 insertions(+), 77 deletions(-) diff --git a/functions-python/tasks_executor/function_config.json b/functions-python/tasks_executor/function_config.json index a33849100..81334ff0f 100644 --- a/functions-python/tasks_executor/function_config.json +++ b/functions-python/tasks_executor/function_config.json @@ -2,8 +2,8 @@ "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"], @@ -21,5 +21,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/src/tasks/pmtiles_builder/build_pmtiles.py b/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py index 1f5d2593a..eee764191 100644 --- a/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py @@ -24,9 +24,11 @@ import pickle import shutil import subprocess -from logging import DEBUG, WARNING, INFO, ERROR +from logging import DEBUG from google.cloud import storage +from shared.helpers.logger import get_logger + # Files are stored locally to be able to run tippecanoe on them. This is the directory local_dir = "./workdir" @@ -63,13 +65,7 @@ def __init__( self.feed_stable_id = feed_stable_id self.dataset_stable_id = dataset_stable_id self.bucket_name = os.getenv("DATASETS_BUCKET_NAME") - - def _log(self, level, msg, *args): - logger = logging.getLogger() - if not logger.isEnabledFor(level): - return - formatted_msg = msg % args if args else msg - logger.log(level, "[%s] %s", self.dataset_stable_id, formatted_msg) + self.logger = get_logger(PmtilesBuilder.__name__, dataset_stable_id) @staticmethod def _get_parameters(payload): @@ -98,8 +94,8 @@ def build_pmtiles(self) -> dict: ) } - self._log( - INFO, "Starting PMTiles build for dataset %s", self.dataset_stable_id + self.logger.info( + "Starting PMTiles build for dataset %s", self.dataset_stable_id ) unzipped_files_path = ( f"{self.feed_stable_id}/{self.dataset_stable_id}/extracted" @@ -131,10 +127,9 @@ def build_pmtiles(self) -> dict: file_list = "\n".join( f"{blob.name} ({blob.size} bytes)" for blob in blobs ) - self._log(DEBUG, "GCS files in %s:\n%s", gcs_prefix, file_list) + self.logger.debug("GCS files in %s:\n%s", gcs_prefix, file_list) except Exception as e: - self._log( - ERROR, + self.logger.error( "Could not list files in bucket %s for path %s: %s", self.bucket_name, gcs_prefix, @@ -153,18 +148,19 @@ def build_pmtiles(self) -> dict: } def _download_files_from_gcs(self, unzipped_files_path): - self._log( - INFO, + self.logger.info( "Downloading dataset from GCS bucket %s, directory %s", self.bucket_name, unzipped_files_path, ) try: - self._log(DEBUG, "Initializing storage client") - self.bucket = storage.Client().get_bucket(self.bucket_name) - self._log(DEBUG, "Getting blobs with prefix: %s", unzipped_files_path) + self.logger.debug("Initializing storage client") + self.bucket = storage.Client( + # client_options={"api_endpoint": "http://localhost:4443"} + ).get_bucket(self.bucket_name) + self.logger.debug("Getting blobs with prefix: %s", unzipped_files_path) blobs = list(self.bucket.list_blobs(prefix=unzipped_files_path)) - self._log(DEBUG, "Found %d blobs", len(blobs)) + self.logger.debug("Found %d blobs", len(blobs)) if not blobs: raise Exception( f"Directory '{unzipped_files_path}' does not exist or is empty in bucket '{self.bucket_name}'." @@ -185,15 +181,14 @@ def _download_files_from_gcs(self, unzipped_files_path): blob = self.bucket.blob(blob_path) local_path = os.path.join(local_dir, file_name) blob.download_to_filename(local_path) - self._log(DEBUG, "Downloaded %s to %s", blob_path, local_path) + self.logger.debug("Downloaded %s to %s", blob_path, local_path) return except Exception as e: raise Exception(f"Failed to download files from GCS: {e}") from e def _upload_files_to_gcs(self, file_to_upload): dest_prefix = f"{self.feed_stable_id}/{self.dataset_stable_id}/pmtiles" - self._log( - INFO, + self.logger.info( "Uploading files to GCS bucket %s, directory %s", self.bucket_name, dest_prefix, @@ -202,17 +197,16 @@ def _upload_files_to_gcs(self, file_to_upload): blobs_to_delete = list(self.bucket.list_blobs(prefix=dest_prefix + "/")) for blob in blobs_to_delete: blob.delete() - self._log(DEBUG, "Deleted existing blob: %s", blob.name) + self.logger.debug("Deleted existing blob: %s", blob.name) for file_name in file_to_upload: file_path = os.path.join(local_dir, file_name) if not os.path.exists(file_path): - self._log(WARNING, "File not found: %s", 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._log( - DEBUG, + self.logger.debug( "Uploaded %s to gs://%s/%s", file_path, self.bucket_name, @@ -222,12 +216,12 @@ def _upload_files_to_gcs(self, file_to_upload): raise Exception(f"Failed to upload files to GCS: {e}") from e def _create_shapes_index(self): - self._log(INFO, "Creating shapes index") + self.logger.info("Creating shapes index") try: index = {} shapes = f"{local_dir}/shapes.txt" outfile = f"{local_dir}/shapes_index.pkl" - with open(shapes, "r", encoding="utf-8") as f: + with open(shapes, "r", encoding="utf-8", newline="") as f: header = f.readline() columns = next(csv.reader([header])) count = 0 @@ -241,9 +235,9 @@ def _create_shapes_index(self): index.setdefault(sid, []).append(pos) count += 1 if count % 1000000 == 0: - self._log(DEBUG, "Indexed %d lines so far...", count) - self._log(DEBUG, "Total indexed lines: %d", count) - self._log(DEBUG, "Total unique shape_ids: %d", len(index)) + 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(index)) with open(outfile, "wb") as idxf: pickle.dump(index, idxf) except Exception as e: @@ -251,18 +245,18 @@ def _create_shapes_index(self): def _read_csv(self, filename): try: - self._log(DEBUG, "Loading %s", filename) + 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._log(DEBUG, "Getting shape points for shape_id %s", shape_id) + self.logger.debug("Getting shape points for shape_id %s", shape_id) try: points = [] shapes_file = f"{local_dir}/shapes.txt" - with open(shapes_file, "r", encoding="utf-8") as f: + with open(shapes_file, "r", encoding="utf-8", newline="") as f: for pos in index.get(shape_id, []): f.seek(pos) line = f.readline() @@ -275,15 +269,17 @@ def _get_shape_points(self, shape_id, index): ) ) points.sort(key=lambda x: x[2]) - self._log(DEBUG, " Found %d points for shape_id %s", len(points), shape_id) + 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): - self._log(INFO, "Creating routes geojson") + self.logger.info("Creating routes geojson") try: - self._log(DEBUG, "Loading shapes_index.pkl...") + self.logger.debug("Loading shapes_index.pkl...") shapes_index_file = f"{local_dir}/shapes_index.pkl" shapes_file = f"{local_dir}/shapes.txt" trips_file = f"{local_dir}/trips.txt" @@ -292,28 +288,27 @@ def _create_routes_geojson(self): stop_times_file = f"{local_dir}/stop_times.txt" with open(shapes_index_file, "rb") as idxf: shapes_index = pickle.load(idxf) - self._log(DEBUG, "Loaded index for %d shape_ids.", len(shapes_index)) + self.logger.debug("Loaded index for %d shape_ids.", len(shapes_index)) - with open(shapes_file, "r", encoding="utf-8") as f: + with open(shapes_file, "r", encoding="utf-8", newline="") as f: header = f.readline() shapes_columns = next(csv.reader([header])) shapes_index["columns"] = shapes_columns routes = {r["route_id"]: r for r in self._read_csv(routes_file)} - self._log(DEBUG, "Loaded %d routes.", len(routes)) + self.logger.debug("Loaded %d routes.", len(routes)) trips = list(self._read_csv(trips_file)) - self._log(DEBUG, "Loaded %d trips.", len(trips)) + self.logger.debug("Loaded %d trips.", len(trips)) stops = { s["stop_id"]: (float(s["stop_lon"]), float(s["stop_lat"])) for s in self._read_csv(stops_file) } - self._log(DEBUG, "Loaded %d stops.", len(stops)) + self.logger.debug("Loaded %d stops.", len(stops)) stop_times_by_trip = {} - self._log( - DEBUG, + self.logger.debug( "Grouping stop_times by trip_id for dataset %s", self.dataset_stable_id, ) @@ -321,15 +316,15 @@ def _create_routes_geojson(self): reader = csv.DictReader(f) for row in reader: stop_times_by_trip.setdefault(row["trip_id"], []).append(row) - self._log( - DEBUG, "Grouped stop_times for %d trips.", len(stop_times_by_trip) + self.logger.debug( + "Grouped stop_times for %d trips.", len(stop_times_by_trip) ) features = [] + missing_coordinates_routes = set() for i, (route_id, route) in enumerate(routes.items(), 1): if i % 100 == 0 or i == 1: - self._log( - DEBUG, + self.logger.debug( "Processing route %d/%d (route_id: %s)", i, len(routes), @@ -337,14 +332,13 @@ def _create_routes_geojson(self): ) trip = next((t for t in trips if t["route_id"] == route_id), None) if not trip: - self._log( - INFO, " No trip found for route_id %s, skipping.", route_id + self.logger.iunfo( + " No trip found for route_id %s, skipping.", route_id ) continue coordinates = [] if "shape_id" in trip and trip["shape_id"]: - self._log( - DEBUG, + self.logger.debug( " Using shape_id %s for route_id %s", trip["shape_id"], route_id, @@ -362,18 +356,13 @@ def _create_routes_geojson(self): for st in trip_stop_times if st["stop_id"] in stops ] - self._log( - DEBUG, + self.logger.debug( " Used %d stop coordinates for route_id %s", len(coordinates), route_id, ) if not coordinates: - self._log( - INFO, - " No coordinates found for route_id %s, skipping.", - route_id, - ) + missing_coordinates_routes.add(route_id) continue features.append( { @@ -383,8 +372,12 @@ def _create_routes_geojson(self): } ) - self._log( - DEBUG, "Writing %d features to routes-output.geojson...", len(features) + if missing_coordinates_routes: + self.logger.info( + "Routes without coordinates: %s", list(missing_coordinates_routes) + ) + self.logger.debug( + "Writing %d features to routes-output.geojson...", len(features) ) routes_geojson = f"{local_dir}/routes-output.geojson" with open(routes_geojson, "w", encoding="utf-8") as f: @@ -393,7 +386,7 @@ def _create_routes_geojson(self): raise Exception(f"Failed to create routes GeoJSON: {e}") from e def _run_tippecanoe(self, input_file, output_file): - self._log(INFO, "Running tippecanoe for input file %s", input_file) + self.logger.info("Running tippecanoe for input file %s", input_file) try: cmd = [ "tippecanoe", @@ -404,26 +397,26 @@ def _run_tippecanoe(self, input_file, output_file): "-zg", f"{local_dir}/{input_file}", ] - self._log(DEBUG, "Running command: %s", " ".join(cmd)) + self.logger.debug("Running command: %s", " ".join(cmd)) result = subprocess.run(cmd, capture_output=True, text=True) if result.stdout: - self._log(DEBUG, "Tippecanoe output:\n%s", result.stdout) + self.logger.debug("Tippecanoe output:\n%s", result.stdout) if result.returncode != 0: - self._log(ERROR, "Tippecanoe error:\n%s", result.stderr) + self.logger.error("Tippecanoe error:\n%s", result.stderr) raise Exception(f"Tippecanoe failed with exit code {result.returncode}") - self._log(DEBUG, "Tippecanoe command executed successfully.") + 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._log(INFO, "Creating stops geojson...") + self.logger.info("Creating stops geojson...") try: stops = self._read_csv(f"{local_dir}/stops.txt") if isinstance(stops, dict) and "error" in stops: return stops - self._log(DEBUG, "Loaded %d stops.", len(stops)) + self.logger.debug("Loaded %d stops.", len(stops)) features = [] for i, stop in enumerate(stops, 1): @@ -431,8 +424,7 @@ def _create_stops_geojson(self): lon = float(stop["stop_lon"]) lat = float(stop["stop_lat"]) except (KeyError, ValueError): - self._log( - INFO, + self.logger.info( "Skipping stop %s: invalid coordinates", stop.get("stop_id", ""), ) @@ -448,8 +440,7 @@ def _create_stops_geojson(self): geojson = {"type": "FeatureCollection", "features": features} stops_geojson = f"{local_dir}/stops-output.geojson" - self._log( - DEBUG, + self.logger.debug( "Writing %d features to stops-output.geojson for dataset %s", len(features), self.dataset_stable_id, @@ -462,7 +453,7 @@ def _create_stops_geojson(self): ) from e def _create_routes_json(self): - self._log(INFO, "Creating routes json...") + self.logger.info("Creating routes json...") try: routes = [] with open(f"{local_dir}/routes.txt", newline="", encoding="utf-8") as f: @@ -480,6 +471,6 @@ def _create_routes_json(self): with open(f"{local_dir}/routes.json", "w", encoding="utf-8") as f: json.dump(routes, f, ensure_ascii=False, indent=4) - self._log(DEBUG, "Converted %d routes to routes.json.", len(routes)) + 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 From d742775997ce21421114e67d0804e2c6188c5a67 Mon Sep 17 00:00:00 2001 From: jcpitre Date: Fri, 1 Aug 2025 16:02:21 -0400 Subject: [PATCH 10/14] Repaired tests. --- .../pmtiles_builder/test_build_pmtiles.py | 22 ++++++++++--------- 1 file changed, 12 insertions(+), 10 deletions(-) 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 index 49c898cfa..cefc5ed62 100644 --- 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 @@ -263,7 +263,6 @@ def test_upload_files_to_gcs_missing_file(self): builder.bucket.blob.return_value = MagicMock() builder.bucket.list_blobs.return_value = [] - # Ensure the file does not exist missing_file = "notfound.pmtiles" if os.path.exists(os.path.join(local_dir, missing_file)): os.remove(os.path.join(local_dir, missing_file)) @@ -271,12 +270,13 @@ def test_upload_files_to_gcs_missing_file(self): with patch( "os.path.exists", side_effect=lambda path: False if missing_file in path else True, - ), patch.object(builder, "_log") as mock_log: + ), self.assertLogs(level="WARNING") as log_cm: builder._upload_files_to_gcs([missing_file]) - mock_log.assert_any_call( - logging.WARNING, - "File not found: %s", - os.path.join(local_dir, missing_file), + self.assertTrue( + any( + f"File not found: {os.path.join(local_dir, missing_file)}" in msg + for msg in log_cm.output + ) ) def test_create_routes_geojson_fallback_to_stop_coordinates(self): @@ -335,11 +335,13 @@ def test_create_stops_geojson_invalid_coordinates(self): f.write("stop1,45.0,-73.0\n") # valid f.write("stop2,not_a_lat,-73.1\n") # invalid lat - with patch.object(self.builder, "_log") as mock_log: + with self.assertLogs(level="INFO") as log_cm: self.builder._create_stops_geojson() - # Check that the log was called for the invalid stop - mock_log.assert_any_call( - logging.INFO, "Skipping stop %s: invalid coordinates", "stop2" + self.assertTrue( + any( + "Skipping stop stop2: invalid coordinates" in msg + for msg in log_cm.output + ) ) # Assert output file exists and only the valid stop is included From 6cf66b01f85ab54f9818b2b35ef7e0ef8976fd68 Mon Sep 17 00:00:00 2001 From: jcpitre Date: Wed, 6 Aug 2025 23:24:59 -0400 Subject: [PATCH 11/14] Optimized memory. --- .../tasks/pmtiles_builder/build_pmtiles.py | 216 +++++++++--------- 1 file changed, 112 insertions(+), 104 deletions(-) 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 index eee764191..f26bcfe84 100644 --- a/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py @@ -21,7 +21,6 @@ import json import logging import os -import pickle import shutil import subprocess from logging import DEBUG @@ -67,6 +66,15 @@ def __init__( 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.stop_times_file = f"{local_dir}/stop_times.txt" + self.shapes_file = f"{local_dir}/shapes.txt" + self.trips_file = f"{local_dir}/trips.txt" + self.routes_file = f"{local_dir}/routes.txt" + self.stops_file = f"{local_dir}/stops.txt" + @staticmethod def _get_parameters(payload): """ @@ -103,8 +111,6 @@ def build_pmtiles(self) -> dict: self._download_files_from_gcs(unzipped_files_path) - self._create_shapes_index() - self._create_routes_geojson() self._run_tippecanoe("routes-output.geojson", "routes.pmtiles") @@ -215,15 +221,24 @@ def _upload_files_to_gcs(self, file_to_upload): except Exception as e: raise Exception(f"Failed to upload files to GCS: {e}") from e - def _create_shapes_index(self): + 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. + """ self.logger.info("Creating shapes index") + shapes_index = {} try: - index = {} - shapes = f"{local_dir}/shapes.txt" - outfile = f"{local_dir}/shapes_index.pkl" - with open(shapes, "r", encoding="utf-8", newline="") as f: + with open(self.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() @@ -232,16 +247,15 @@ def _create_shapes_index(self): break row = dict(zip(columns, next(csv.reader([line])))) sid = row["shape_id"] - index.setdefault(sid, []).append(pos) + 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(index)) - with open(outfile, "wb") as idxf: - pickle.dump(index, idxf) + self.logger.debug("Total unique shape_ids: %d", len(shapes_index)) except Exception as e: raise Exception(f"Failed to create shapes index: {e}") from e + return shapes_index def _read_csv(self, filename): try: @@ -255,8 +269,7 @@ def _get_shape_points(self, shape_id, index): self.logger.debug("Getting shape points for shape_id %s", shape_id) try: points = [] - shapes_file = f"{local_dir}/shapes.txt" - with open(shapes_file, "r", encoding="utf-8", newline="") as f: + with open(self.shapes_file, "r", encoding="utf-8", newline="") as f: for pos in index.get(shape_id, []): f.seek(pos) line = f.readline() @@ -277,114 +290,109 @@ def _get_shape_points(self, shape_id, index): raise Exception(f"Failed to get shape points for {shape_id}: {e}") from e def _create_routes_geojson(self): - self.logger.info("Creating routes geojson") try: - self.logger.debug("Loading shapes_index.pkl...") - shapes_index_file = f"{local_dir}/shapes_index.pkl" - shapes_file = f"{local_dir}/shapes.txt" - trips_file = f"{local_dir}/trips.txt" - routes_file = f"{local_dir}/routes.txt" - stops_file = f"{local_dir}/stops.txt" - stop_times_file = f"{local_dir}/stop_times.txt" - with open(shapes_index_file, "rb") as idxf: - shapes_index = pickle.load(idxf) - self.logger.debug("Loaded index for %d shape_ids.", len(shapes_index)) - - with open(shapes_file, "r", encoding="utf-8", newline="") as f: - header = f.readline() - shapes_columns = next(csv.reader([header])) - shapes_index["columns"] = shapes_columns - - routes = {r["route_id"]: r for r in self._read_csv(routes_file)} - self.logger.debug("Loaded %d routes.", len(routes)) + shapes_index = self._create_shapes_index() + self.logger.info("Creating routes geojson (optimized for memory)") - trips = list(self._read_csv(trips_file)) - self.logger.debug("Loaded %d trips.", len(trips)) - - stops = { + # 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(stops_file) + for s in self._read_csv(self.stops_file) } - self.logger.debug("Loaded %d stops.", len(stops)) - stop_times_by_trip = {} - self.logger.debug( - "Grouping stop_times by trip_id for dataset %s", - self.dataset_stable_id, - ) - with open(stop_times_file, newline="", encoding="utf-8") as f: - reader = csv.DictReader(f) - for row in reader: - stop_times_by_trip.setdefault(row["trip_id"], []).append(row) - self.logger.debug( - "Grouped stop_times for %d trips.", len(stop_times_by_trip) - ) + # 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.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() - for i, (route_id, route) in enumerate(routes.items(), 1): - if i % 100 == 0 or i == 1: - self.logger.debug( - "Processing route %d/%d (route_id: %s)", - i, - len(routes), - route_id, - ) - trip = next((t for t in trips if t["route_id"] == route_id), None) - if not trip: - self.logger.iunfo( - " No trip found for route_id %s, skipping.", route_id - ) - continue - coordinates = [] - if "shape_id" in trip and trip["shape_id"]: - self.logger.debug( - " Using shape_id %s for route_id %s", - trip["shape_id"], - route_id, - ) - coordinates = self._get_shape_points(trip["shape_id"], shapes_index) - if isinstance(coordinates, dict) and "error" in coordinates: - raise Exception( - f"Error getting shape points for shape_id {trip['shape_id']}: {coordinates['error']}" - ) - if not coordinates: - trip_stop_times = stop_times_by_trip.get(trip["trip_id"], []) - trip_stop_times.sort(key=lambda x: int(x["stop_sequence"])) - coordinates = [ - stops[st["stop_id"]] - for st in trip_stop_times - if st["stop_id"] in stops - ] - self.logger.debug( - " Used %d stop coordinates for route_id %s", - len(coordinates), - route_id, - ) - if not coordinates: - missing_coordinates_routes.add(route_id) - continue - features.append( - { - "type": "Feature", - "properties": {k: route[k] for k in route}, - "geometry": {"type": "LineString", "coordinates": coordinates}, - } - ) + routes_geojson = f"{local_dir}/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.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( - "Writing %d features to routes-output.geojson...", len(features) + "Wrote %d features to routes-output.geojson", len(features) ) - routes_geojson = f"{local_dir}/routes-output.geojson" - with open(routes_geojson, "w", encoding="utf-8") as f: - json.dump({"type": "FeatureCollection", "features": features}, f) 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.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: From 6c9880f724177b3f104483a28a5f418e1310322c Mon Sep 17 00:00:00 2001 From: jcpitre Date: Wed, 6 Aug 2025 23:41:14 -0400 Subject: [PATCH 12/14] Minor corrections. --- .../src/tasks/pmtiles_builder/build_pmtiles.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 index f26bcfe84..38089690d 100644 --- a/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py @@ -232,6 +232,7 @@ def _create_shapes_index(self) -> dict: 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: @@ -421,7 +422,7 @@ def _run_tippecanoe(self, input_file, output_file): def _create_stops_geojson(self): self.logger.info("Creating stops geojson...") try: - stops = self._read_csv(f"{local_dir}/stops.txt") + stops = self._read_csv(self.stops_file) if isinstance(stops, dict) and "error" in stops: return stops self.logger.debug("Loaded %d stops.", len(stops)) @@ -464,7 +465,7 @@ def _create_routes_json(self): self.logger.info("Creating routes json...") try: routes = [] - with open(f"{local_dir}/routes.txt", newline="", encoding="utf-8") as f: + with open(self.routes_file, newline="", encoding="utf-8") as f: reader = csv.DictReader(f) for row in reader: route = { From 6ebd55e778b7b89c5fb1b38244e7dc005648be44 Mon Sep 17 00:00:00 2001 From: jcpitre Date: Wed, 6 Aug 2025 23:48:31 -0400 Subject: [PATCH 13/14] Modified tests --- .../pmtiles_builder/test_build_pmtiles.py | 42 ++++++++----------- 1 file changed, 17 insertions(+), 25 deletions(-) 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 index cefc5ed62..608193806 100644 --- 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 @@ -92,34 +92,26 @@ def tearDown(self): if "DATASETS_BUCKET_NAME" in os.environ: del os.environ["DATASETS_BUCKET_NAME"] - @patch( - "tasks.pmtiles_builder.build_pmtiles.PmtilesBuilder._download_files_from_gcs" - ) - @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_calls_create_shapes_index( - self, - mock_upload, - mock_routes_json, - mock_stops_geojson, - mock_run_tippecanoe, - mock_routes_geojson, - mock_download, - ): - self.builder.bucket = MagicMock() - self.builder.bucket.list_blobs.return_value = [] - # Create minimal shapes.txt in local_dir + def test_build_pmtiles_creates_correct_shapes_index(self): + # Prepare shapes.txt os.makedirs(local_dir, exist_ok=True) shapes_path = os.path.join(local_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\ns1,45.0,-73.0,1\n" - ) - result = self.builder.build_pmtiles() - self.assertIn("message", result) + 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 From ae947d8c9b8ba1961318aaea2dfc5793db58bd90 Mon Sep 17 00:00:00 2001 From: jcpitre Date: Mon, 11 Aug 2025 16:58:50 -0400 Subject: [PATCH 14/14] Modified according to PR comments. --- .../tasks/pmtiles_builder/build_pmtiles.py | 262 ++++---- .../pmtiles_builder/test_build_pmtiles.py | 563 ++++++++++-------- 2 files changed, 469 insertions(+), 356 deletions(-) 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 index 38089690d..d899b628d 100644 --- a/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py +++ b/functions-python/tasks_executor/src/tasks/pmtiles_builder/build_pmtiles.py @@ -21,15 +21,18 @@ import json import logging import os -import shutil import subprocess -from logging import DEBUG +import tempfile +from enum import Enum from google.cloud import storage from shared.helpers.logger import get_logger -# Files are stored locally to be able to run tippecanoe on them. This is the directory -local_dir = "./workdir" +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: @@ -37,13 +40,48 @@ def build_pmtiles_handler(payload) -> dict: Entrypoint for building PMTiles files from a GTFS dataset. """ try: - feed_stable_id, dataset_stable_id = PmtilesBuilder._get_parameters(payload) - builder = PmtilesBuilder( - feed_stable_id=feed_stable_id, dataset_stable_id=dataset_stable_id - ) - return builder.build_pmtiles() + # 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: - return {"error": f"Failed to start PMTiles build: {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: @@ -55,10 +93,15 @@ class PmtilesBuilder: 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 @@ -69,11 +112,12 @@ def __init__( self.stop_times_index = {} self.stop_times_by_trip = None - self.stop_times_file = f"{local_dir}/stop_times.txt" - self.shapes_file = f"{local_dir}/shapes.txt" - self.trips_file = f"{local_dir}/trips.txt" - self.routes_file = f"{local_dir}/routes.txt" - self.stops_file = f"{local_dir}/stops.txt" + 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): @@ -84,74 +128,44 @@ def _get_parameters(payload): dataset_stable_id = payload.get("dataset_stable_id", None) return feed_stable_id, dataset_stable_id - def build_pmtiles(self) -> dict: - try: - if not self.bucket_name: - return { - "error": "DATASETS_BUCKET_NAME environment variable is not defined." - } - if not self.feed_stable_id or not self.dataset_stable_id: - return { - "error": "Both feed_stable_id and dataset_stable_id must be defined." - } - if self.feed_stable_id not in self.dataset_stable_id: - return { - "error": ( - "feed_stable_id %s is not a prefix of dataset_stable_id %s." - % (self.feed_stable_id, self.dataset_stable_id) - ) - } + def build_pmtiles(self): + if not self.bucket_name: + raise Exception("DATASETS_BUCKET_NAME environment variable is not defined.") - self.logger.info( - "Starting PMTiles build for dataset %s", self.dataset_stable_id + 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." ) - unzipped_files_path = ( - f"{self.feed_stable_id}/{self.dataset_stable_id}/extracted" + + 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._download_files_from_gcs(unzipped_files_path) + self.logger.info("Starting PMTiles build") + unzipped_files_path = ( + f"{self.feed_stable_id}/{self.dataset_stable_id}/extracted" + ) - self._create_routes_geojson() + status, message = self._download_files_from_gcs(unzipped_files_path) + if status == self.OperationStatus.FAILURE: + return status, message - self._run_tippecanoe("routes-output.geojson", "routes.pmtiles") + self._create_routes_geojson() - self._create_stops_geojson() + self._run_tippecanoe("routes-output.geojson", "routes.pmtiles") - self._run_tippecanoe("stops-output.geojson", "stops.pmtiles") + self._create_stops_geojson() - self._create_routes_json() + self._run_tippecanoe("stops-output.geojson", "stops.pmtiles") - files_to_upload = ["routes.pmtiles", "stops.pmtiles", "routes.json"] - self._upload_files_to_gcs(files_to_upload) + self._create_routes_json() - # List files in the relevant bucket folder instead of local_dir + files_to_upload = ["routes.pmtiles", "stops.pmtiles", "routes.json"] + self._upload_files_to_gcs(files_to_upload) - if logging.getLogger().isEnabledFor(DEBUG): - gcs_prefix = f"{self.feed_stable_id}/{self.dataset_stable_id}/pmtiles/" - try: # We don`t want an error here to abort the whole pmtiles operation. - blobs = list(self.bucket.list_blobs(prefix=gcs_prefix)) - file_list = "\n".join( - f"{blob.name} ({blob.size} bytes)" for blob in blobs - ) - self.logger.debug("GCS files in %s:\n%s", gcs_prefix, file_list) - except Exception as e: - self.logger.error( - "Could not list files in bucket %s for path %s: %s", - self.bucket_name, - gcs_prefix, - e, - ) - - return { - "message": f"Pmtiles successfully created for dataset {self.dataset_stable_id}." - } - except Exception as e: - logging.exception( - "Failed to build PMTiles for dataset %s", self.dataset_stable_id - ) - return { - "error": f"Failed to build PMTiles for dataset {self.dataset_stable_id}: {e}" - } + return self.OperationStatus.SUCCESS, "success" def _download_files_from_gcs(self, unzipped_files_path): self.logger.info( @@ -161,36 +175,60 @@ def _download_files_from_gcs(self, unzipped_files_path): ) try: self.logger.debug("Initializing storage client") - self.bucket = storage.Client( - # client_options={"api_endpoint": "http://localhost:4443"} - ).get_bucket(self.bucket_name) + 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: - raise Exception( - f"Directory '{unzipped_files_path}' does not exist or is empty in bucket '{self.bucket_name}'." - ) - - if os.path.exists(local_dir): - shutil.rmtree(local_dir) - os.makedirs(local_dir, exist_ok=True) - file_names = [ - "routes.txt", - "shapes.txt", - "stop_times.txt", - "trips.txt", - "stops.txt", + 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_name in file_names: + 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) - local_path = os.path.join(local_dir, file_name) - blob.download_to_filename(local_path) - self.logger.debug("Downloaded %s to %s", blob_path, local_path) - return + 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: - raise Exception(f"Failed to download files from GCS: {e}") from 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" @@ -205,7 +243,7 @@ def _upload_files_to_gcs(self, file_to_upload): blob.delete() self.logger.debug("Deleted existing blob: %s", blob.name) for file_name in file_to_upload: - file_path = os.path.join(local_dir, file_name) + 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 @@ -236,7 +274,9 @@ def _create_shapes_index(self) -> dict: self.logger.info("Creating shapes index") shapes_index = {} try: - with open(self.shapes_file, "r", encoding="utf-8", newline="") as f: + 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 @@ -255,7 +295,7 @@ def _create_shapes_index(self) -> dict: self.logger.debug("Total indexed lines: %d", count) self.logger.debug("Total unique shape_ids: %d", len(shapes_index)) except Exception as e: - raise Exception(f"Failed to create shapes index: {e}") from e + self.logger.warning("Cannot read shapes file: %s", e) return shapes_index def _read_csv(self, filename): @@ -270,7 +310,9 @@ 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.shapes_file, "r", encoding="utf-8", newline="") as f: + 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() @@ -299,7 +341,7 @@ def _create_routes_geojson(self): # 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.stops_file) + 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 @@ -307,7 +349,7 @@ def _create_routes_geojson(self): # when we select a route. shape_map_indexed_by_route = {} trip_map_indexed_by_route = {} - with open(self.trips_file, newline="", encoding="utf-8") as f: + 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"] @@ -319,12 +361,12 @@ def _create_routes_geojson(self): features = [] missing_coordinates_routes = set() - routes_geojson = f"{local_dir}/routes-output.geojson" + 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.routes_file, newline="", encoding="utf-8" + 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"] @@ -386,7 +428,9 @@ 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.stop_times_file, newline="", encoding="utf-8") as f: + 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"] @@ -400,11 +444,11 @@ def _run_tippecanoe(self, input_file, output_file): cmd = [ "tippecanoe", "-o", - f"{local_dir}/{output_file}", + self.get_path(output_file), "--force", "--no-tile-size-limit", "-zg", - f"{local_dir}/{input_file}", + self.get_path(input_file), ] self.logger.debug("Running command: %s", " ".join(cmd)) result = subprocess.run(cmd, capture_output=True, text=True) @@ -422,7 +466,7 @@ def _run_tippecanoe(self, input_file, output_file): def _create_stops_geojson(self): self.logger.info("Creating stops geojson...") try: - stops = self._read_csv(self.stops_file) + 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)) @@ -447,7 +491,7 @@ def _create_stops_geojson(self): ) geojson = {"type": "FeatureCollection", "features": features} - stops_geojson = f"{local_dir}/stops-output.geojson" + stops_geojson = self.get_path("stops-output.geojson") self.logger.debug( "Writing %d features to stops-output.geojson for dataset %s", @@ -465,7 +509,7 @@ def _create_routes_json(self): self.logger.info("Creating routes json...") try: routes = [] - with open(self.routes_file, newline="", encoding="utf-8") as f: + with open(self.get_path(ROUTES_FILE), newline="", encoding="utf-8") as f: reader = csv.DictReader(f) for row in reader: route = { @@ -473,11 +517,11 @@ def _create_routes_json(self): "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', '3')}", + "routeType": f"{row.get('route_type', 'unknown')}", } routes.append(route) - with open(f"{local_dir}/routes.json", "w", encoding="utf-8") as f: + 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)) 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 index 608193806..2f02ad3f3 100644 --- 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 @@ -1,7 +1,6 @@ import csv import json import logging -import pickle import tempfile import unittest from contextlib import contextmanager @@ -10,7 +9,6 @@ from tasks.pmtiles_builder.build_pmtiles import ( PmtilesBuilder, - local_dir, build_pmtiles_handler, ) @@ -25,6 +23,97 @@ def suppress_logging(level=logging.CRITICAL): 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" @@ -32,20 +121,6 @@ def setUp(self): os.environ["DATASETS_BUCKET_NAME"] = "test-bucket" self.builder = PmtilesBuilder(self.feed_stable_id, self.dataset_stable_id) - @patch("tasks.pmtiles_builder.build_pmtiles.storage.Client") - def test_download_files_from_gcs_success(self, mock_client): - mock_bucket = MagicMock() - mock_blob = MagicMock() - mock_blob.download_to_filename = MagicMock() - mock_bucket.list_blobs.return_value = [mock_blob] * 5 - mock_client.return_value.get_bucket.return_value = mock_bucket - - with patch("os.path.exists", return_value=True), patch("shutil.rmtree"), patch( - "os.makedirs" - ): - self.builder._download_files_from_gcs("some/path") - self.assertTrue(mock_bucket.list_blobs.called) - @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="") @@ -55,7 +130,7 @@ def test_run_tippecanoe_success(self, mock_run): @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): + with self.assertRaises(Exception), suppress_logging(): self.builder._run_tippecanoe("input.geojson", "output.pmtiles") @patch( @@ -79,8 +154,13 @@ def test_build_pmtiles_success( ): self.builder.bucket = MagicMock() self.builder.bucket.list_blobs.return_value = [] - result = self.builder.build_pmtiles() - self.assertIn("message", result) + 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"} @@ -94,15 +174,16 @@ def tearDown(self): def test_build_pmtiles_creates_correct_shapes_index(self): # Prepare shapes.txt - os.makedirs(local_dir, exist_ok=True) - shapes_path = os.path.join(local_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() + 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"], @@ -115,120 +196,111 @@ def test_build_pmtiles_creates_correct_shapes_index(self): def test_get_shape_points(self): # Prepare shapes.txt - os.makedirs(local_dir, exist_ok=True) - shapes_path = os.path.join(local_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") - - # 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) + 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 - os.makedirs(local_dir, exist_ok=True) - with open(os.path.join(local_dir, "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(os.path.join(local_dir, "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(os.path.join(local_dir, "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( - os.path.join(local_dir, "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(os.path.join(local_dir, "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") - # Create shapes_index.pkl - shapes_index = { - "s1": [ - len("shape_id,shape_pt_lat,shape_pt_lon,shape_pt_sequence\n"), - len( - "shape_id,shape_pt_lat,shape_pt_lon,shape_pt_sequence\ns1,45.0,-73.0,1\n" - ), - ], - "columns": [ - "shape_id", - "shape_pt_lat", - "shape_pt_lon", - "shape_pt_sequence", - ], - } - with open(os.path.join(local_dir, "shapes_index.pkl"), "wb") as f: - pickle.dump(shapes_index, f) - - # Call the method - self.builder._create_routes_geojson() - - # Assert output file exists and is valid GeoJSON - output_path = os.path.join(local_dir, "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) + 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 - os.makedirs(local_dir, exist_ok=True) - with open(os.path.join(local_dir, "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 = os.path.join(local_dir, "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) + 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 - os.makedirs(local_dir, exist_ok=True) - with open(os.path.join(local_dir, "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") + 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() + # Call the method + self.builder._create_routes_json() - # Assert output file exists and is valid JSON - output_path = os.path.join(local_dir, "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) + # 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 - if os.path.exists(os.path.join(local_dir, "routes.txt")): - os.remove(os.path.join(local_dir, "routes.txt")) + 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)) @@ -245,115 +317,104 @@ def test_build_pmtiles_exception(self): # 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(): - result = self.builder.build_pmtiles() - self.assertIn("error", result) - self.assertIn("Failed to build PMTiles for dataset", result["error"]) + 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): - builder = PmtilesBuilder(feed_stable_id="foo", dataset_stable_id="foo_bar") - builder.bucket = MagicMock() - builder.bucket.blob.return_value = MagicMock() - builder.bucket.list_blobs.return_value = [] + self.builder.bucket = MagicMock() + self.builder.bucket.blob.return_value = MagicMock() + self.builder.bucket.list_blobs.return_value = [] missing_file = "notfound.pmtiles" - if os.path.exists(os.path.join(local_dir, missing_file)): - os.remove(os.path.join(local_dir, missing_file)) + 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: - builder._upload_files_to_gcs([missing_file]) + self.builder._upload_files_to_gcs([missing_file]) self.assertTrue( - any( - f"File not found: {os.path.join(local_dir, missing_file)}" in msg - for msg in log_cm.output - ) + 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 - os.makedirs(local_dir, exist_ok=True) - with open(os.path.join(local_dir, "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(os.path.join(local_dir, "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(os.path.join(local_dir, "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( - os.path.join(local_dir, "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") - # shapes.txt and shapes_index.pkl are still needed but not used in this case - with open(os.path.join(local_dir, "shapes.txt"), "w", encoding="utf-8") as f: - f.write("shape_id,shape_pt_lat,shape_pt_lon,shape_pt_sequence\n") - with open(os.path.join(local_dir, "shapes_index.pkl"), "wb") as f: - pickle.dump( - { - "columns": [ - "shape_id", - "shape_pt_lat", - "shape_pt_lon", - "shape_pt_sequence", - ] - }, - f, - ) - - # Call the method - self.builder._create_routes_geojson() - - # Assert output file exists and contains the expected coordinates - output_path = os.path.join(local_dir, "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]]) + 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): - # Prepare stops.txt with one valid and one invalid stop - os.makedirs(local_dir, exist_ok=True) - with open(os.path.join(local_dir, "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") # 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 + 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 + ) ) - ) - # Assert output file exists and only the valid stop is included - output_path = os.path.join(local_dir, "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") + 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): - # Patch local_dir to a temp directory self.test_dir = tempfile.TemporaryDirectory() - self.old_local_dir = local_dir - self._patch_local_dir(self.test_dir.name) 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 + # Create minimal GTFS files using builder.get_path files = { "routes.txt": ( "route_id,route_long_name,route_color,route_text_color,route_type\n" @@ -375,66 +436,74 @@ def setUp(self): ), } for fname, content in files.items(): - with open( - os.path.join(self.test_dir.name, fname), "w", encoding="utf-8" - ) as f: + with open(self.builder.get_path(fname), "w", encoding="utf-8") as f: f.write(content) def tearDown(self): self.test_dir.cleanup() - self._patch_local_dir(self.old_local_dir) if "DATASETS_BUCKET_NAME" in os.environ: del os.environ["DATASETS_BUCKET_NAME"] - def _patch_local_dir(self, new_dir): - import tasks.pmtiles_builder.build_pmtiles as mod - - mod.local_dir = new_dir - def test_build_pmtiles_handler_missing_bucket_env(self): - if "DATASETS_BUCKET_NAME" in os.environ: - del os.environ["DATASETS_BUCKET_NAME"] + os.environ.pop("DATASETS_BUCKET_NAME", None) payload = { - "feed_stable_id": "feed123", - "dataset_stable_id": "feed123_dataset456", + "feed_stable_id": self.feed_stable_id, + "dataset_stable_id": self.dataset_stable_id, } - result = build_pmtiles_handler(payload) + with suppress_logging(): + result = build_pmtiles_handler(payload) self.assertIn("error", result) - self.assertIn("DATASETS_BUCKET_NAME", result["error"]) + self.assertIn( + "DATASETS_BUCKET_NAME environment variable is not defined.", result["error"] + ) def test_build_pmtiles_handler_missing_ids(self): - os.environ["DATASETS_BUCKET_NAME"] = "test-bucket" - payload = {"feed_stable_id": "", "dataset_stable_id": ""} - result = build_pmtiles_handler(payload) + payload = {} + with suppress_logging(): + result = build_pmtiles_handler(payload) self.assertIn("error", result) - self.assertIn("must be defined", result["error"]) + self.assertIn( + "Both feed_stable_id and dataset_stable_id must be defined.", + result["error"], + ) def test_build_pmtiles_handler_feed_not_prefix(self): - os.environ["DATASETS_BUCKET_NAME"] = "test-bucket" - payload = {"feed_stable_id": "foo", "dataset_stable_id": "barbaz"} - result = build_pmtiles_handler(payload) + 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", result["error"]) + 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): - builder = PmtilesBuilder(feed_stable_id="foo", dataset_stable_id="foo_bar") - builder.bucket = MagicMock() - mock_blob = MagicMock() - builder.bucket.blob.return_value = mock_blob - builder.bucket.list_blobs.return_value = [] - - # Create dummy files in local_dir for the test - os.makedirs(local_dir, exist_ok=True) - test_file = os.path.join(local_dir, "routes.pmtiles") + test_file = self.builder.get_path("routes.pmtiles") with open(test_file, "w") as f: f.write("dummy data") - builder._upload_files_to_gcs(["routes.pmtiles"]) - - builder.bucket.blob.assert_called_with("foo/foo_bar/pmtiles/routes.pmtiles") - mock_blob.upload_from_filename.assert_called_with(test_file) + 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__":