feat: batch fill pmtiles files (#1336)

Author: cka-y

functions-python/batch_process_dataset/src/main.py

@@ -262,29 +262,22 @@ def upload_dataset(self, public=True) -> DatasetFile or None:
         return None
 
     @with_db_session
-    def process_from_bucket_latest(
-        self, db_session, public=True
-    ) -> DatasetFile or None:
+    def process_from_bucket(self, db_session, public=True) -> Optional[DatasetFile]:
         """
-        Uploads a dataset to a GCP bucket as <feed_stable_id>/latest.zip and
-        <feed_stable_id>/<feed_stable_id>-<upload_datetime>.zip
-        if the dataset hash is different from the latest dataset stored
-        :return: the file hash and the hosted url as a tuple or None if no upload is required
+        Process an existing dataset from the GCP bucket updates the related database entities
+        :return: The DatasetFile object created
         """
         temp_file_path = None
         try:
-            self.logger.info("Accessing URL %s", self.producer_url)
             temp_file_path = self.generate_temp_filename()
-            blob_file_path = f"{self.feed_stable_id}/latest.zip"
+            blob_file_path = f"{self.feed_stable_id}/{self.dataset_stable_id}/{self.dataset_stable_id}.zip"
+            self.logger.info(f"Processing dataset from bucket: {blob_file_path}")
             download_from_gcs(
                 os.getenv("DATASETS_BUCKET_NAME"), blob_file_path, temp_file_path
             )
 
             extracted_files_path = self.unzip_files(temp_file_path)
-            dataset_full_path = f"{self.feed_stable_id}/{self.dataset_stable_id}/{self.dataset_stable_id}.zip"
-            self.logger.info(
-                f"Creating file {dataset_full_path} in bucket {self.bucket_name}"
-            )
+
             _, extracted_files = self.upload_files_to_storage(
                 temp_file_path,
                 self.dataset_stable_id,
@@ -296,7 +289,7 @@ def process_from_bucket_latest(
             dataset_file = DatasetFile(
                 stable_id=self.dataset_stable_id,
                 file_sha256_hash=self.latest_hash,
-                hosted_url=f"{self.public_hosted_datasets_url}/{dataset_full_path}",
+                hosted_url=f"{self.public_hosted_datasets_url}/{blob_file_path}",
                 extracted_files=extracted_files,
                 zipped_size=(
                     os.path.getsize(temp_file_path)
@@ -307,11 +300,21 @@ def process_from_bucket_latest(
             dataset = self.create_dataset_entities(
                 dataset_file, skip_dataset_creation=True, db_session=db_session
             )
-            create_pipeline_tasks(dataset)
+            if dataset and dataset.latest:
+                self.logger.info(
+                    f"Creating pipeline tasks for latest dataset {dataset.stable_id}"
+                )
+                create_pipeline_tasks(dataset)
+            elif dataset:
+                self.logger.info(
+                    f"Dataset {dataset.stable_id} is not the latest, skipping pipeline tasks creation."
+                )
+            else:
+                raise ValueError("Dataset update failed, dataset is None.")
+            return dataset_file
         finally:
             if temp_file_path and os.path.exists(temp_file_path):
                 os.remove(temp_file_path)
-        return None
 
     def unzip_files(self, temp_file_path):
         extracted_files_path = os.path.join(temp_file_path.split(".")[0], "extracted")
@@ -356,11 +359,11 @@ def create_dataset_entities(
                     f"[{self.feed_stable_id}] No latest dataset found for feed."
                 )
 
-            self.logger.info(
-                f"[{self.feed_stable_id}] Creating new dataset for feed with stable id {dataset_file.stable_id}."
-            )
             dataset = None
             if not skip_dataset_creation:
+                self.logger.info(
+                    f"[{self.feed_stable_id}] Creating new dataset for feed with stable id {dataset_file.stable_id}."
+                )
                 dataset = Gtfsdataset(
                     id=str(uuid.uuid4()),
                     feed_id=self.feed_id,
@@ -377,22 +380,20 @@ def create_dataset_entities(
                         else []
                     ),
                     zipped_size_bytes=dataset_file.zipped_size,
-                    unzipped_size_bytes=(
-                        sum([ex.file_size_bytes for ex in dataset_file.extracted_files])
-                        if dataset_file.extracted_files
-                        else None
-                    ),
+                    unzipped_size_bytes=self._get_unzipped_size(dataset_file),
                 )
                 db_session.add(dataset)
             elif skip_dataset_creation and latest_dataset:
+                self.logger.info(
+                    f"[{self.feed_stable_id}] Updating latest dataset for feed with stable id "
+                    f"{latest_dataset.stable_id}."
+                )
                 latest_dataset.gtfsfiles = (
                     dataset_file.extracted_files if dataset_file.extracted_files else []
                 )
                 latest_dataset.zipped_size_bytes = dataset_file.zipped_size
-                latest_dataset.unzipped_size_bytes = (
-                    sum([ex.file_size_bytes for ex in dataset_file.extracted_files])
-                    if dataset_file.extracted_files
-                    else None
+                latest_dataset.unzipped_size_bytes = self._get_unzipped_size(
+                    dataset_file
                 )
 
             if latest_dataset and not skip_dataset_creation:
@@ -406,11 +407,19 @@ def create_dataset_entities(
         except Exception as e:
             raise Exception(f"Error creating dataset: {e}")
 
+    @staticmethod
+    def _get_unzipped_size(dataset_file):
+        return (
+            sum([ex.file_size_bytes for ex in dataset_file.extracted_files])
+            if dataset_file.extracted_files
+            else None
+        )
+
     @with_db_session
-    def process_from_producer_url(self, db_session) -> DatasetFile or None:
+    def process_from_producer_url(self, db_session) -> Optional[DatasetFile]:
         """
         Process the dataset and store new version in GCP bucket if any changes are detected
-        :return: the file hash and the hosted url as a tuple or None if no upload is required
+        :return: the DatasetFile object created
         """
         dataset_file = self.upload_dataset()
 
@@ -531,7 +540,7 @@ def process_dataset(cloud_event: CloudEvent):
             json_payload.get("dataset_stable_id"),
         )
         if json_payload.get("use_bucket_latest", False):
-            dataset_file = processor.process_from_bucket_latest()
+            dataset_file = processor.process_from_bucket()
         else:
             dataset_file = processor.process_from_producer_url()
     except Exception as e:

functions-python/batch_process_dataset/src/pipeline_tasks.py

@@ -8,7 +8,7 @@
 
 from shared.database.database import with_db_session
 from shared.database_gen.sqlacodegen_models import Gtfsdataset
-from shared.helpers.utils import create_http_task
+from shared.helpers.utils import create_http_task, create_http_pmtiles_builder_task
 
 
 def create_http_reverse_geolocation_processor_task(
@@ -41,32 +41,6 @@ def create_http_reverse_geolocation_processor_task(
     )
 
 
-def create_http_pmtiles_builder_task(
-    stable_id: str,
-    dataset_stable_id: str,
-) -> None:
-    """
-    Create a task to generate PMTiles for a dataset.
-    """
-    client = tasks_v2.CloudTasksClient()
-    body = json.dumps(
-        {"feed_stable_id": stable_id, "dataset_stable_id": dataset_stable_id}
-    ).encode()
-    queue_name = os.getenv("PMTILES_BUILDER_QUEUE")
-    project_id = os.getenv("PROJECT_ID")
-    gcp_region = os.getenv("GCP_REGION")
-    gcp_env = os.getenv("ENVIRONMENT")
-
-    create_http_task(
-        client,
-        body,
-        f"https://{gcp_region}-{project_id}.cloudfunctions.net/pmtiles-builder-{gcp_env}",
-        project_id,
-        gcp_region,
-        queue_name,
-    )
-
-
 @with_db_session
 def get_changed_files(
     dataset: Gtfsdataset,

functions-python/batch_process_dataset/tests/test_batch_process_dataset_main.py

@@ -488,16 +488,18 @@ def test_process_from_bucket_latest_happy_path(
         )
 
         # Act
-        result = processor.process_from_bucket_latest(public=True)
+        result = processor.process_from_bucket(public=True)
 
         # Assert: function returns None in current implementation
-        self.assertIsNone(result)
+        self.assertIsNone(result.zipped_size)
 
         # Assert: downloads from the bucket latest.zip for this feed
         mock_download_from_gcs.assert_called_once()
         args, kwargs = mock_download_from_gcs.call_args
         self.assertEqual(args[0], "test-bucket")  # bucket name
-        self.assertEqual(args[1], "feed_stable_id/latest.zip")  # blob path
+        self.assertEqual(
+            args[1], "feed_stable_id/dataset-stable-id-123/dataset-stable-id-123.zip"
+        )  # blob path
         self.assertIsNotNone(
             args[2]
         )  # temp file path (random), so just ensure it exists

functions-python/batch_process_dataset/tests/test_pipeline_tasks.py

@@ -5,7 +5,6 @@
 
 from pipeline_tasks import (
     create_http_reverse_geolocation_processor_task,
-    create_http_pmtiles_builder_task,
     get_changed_files,
     create_pipeline_tasks,
 )
@@ -86,48 +85,6 @@ def test_create_http_reverse_geolocation_processor_task(
         self.assertEqual(args[4], "northamerica-northeast1")
         self.assertEqual(args[5], "rev-geo-queue")
 
-    @patch.dict(
-        os.environ,
-        {
-            "PMTILES_BUILDER_QUEUE": "pmtiles-queue",
-            "PROJECT_ID": "my-project",
-            "GCP_REGION": "northamerica-northeast1",
-            "ENVIRONMENT": "dev",
-        },
-        clear=False,
-    )
-    @patch("pipeline_tasks.create_http_task")
-    @patch("pipeline_tasks.tasks_v2.CloudTasksClient")
-    def test_create_http_pmtiles_builder_task(
-        self, mock_client_cls, mock_create_http_task
-    ):
-        client_instance = MagicMock()
-        mock_client_cls.return_value = client_instance
-
-        stable_id = "feed-456"
-        dataset_stable_id = "dataset-def"
-
-        create_http_pmtiles_builder_task(
-            stable_id=stable_id, dataset_stable_id=dataset_stable_id
-        )
-
-        mock_client_cls.assert_called_once()
-        self.assertEqual(mock_create_http_task.call_count, 1)
-        args, _ = mock_create_http_task.call_args
-
-        payload = json.loads(args[1].decode("utf-8"))
-        self.assertEqual(
-            payload,
-            {"feed_stable_id": stable_id, "dataset_stable_id": dataset_stable_id},
-        )
-        self.assertEqual(
-            args[2],
-            "https://northamerica-northeast1-my-project.cloudfunctions.net/pmtiles-builder-dev",
-        )
-        self.assertEqual(args[3], "my-project")
-        self.assertEqual(args[4], "northamerica-northeast1")
-        self.assertEqual(args[5], "pmtiles-queue")
-
 
 class TestHasFileChanged(unittest.TestCase):
     def _make_mock_session_chain(self, previous_dataset):

functions-python/helpers/tests/test_helpers.py

@@ -212,3 +212,44 @@ def test_create_http_task(self):
         url = "test"
         create_http_task(client, body, url, "test", "test", "test")
         client.create_task.assert_called_once()
+
+    @patch.dict(
+        os.environ,
+        {
+            "PMTILES_BUILDER_QUEUE": "pmtiles-queue",
+            "PROJECT_ID": "my-project",
+            "GCP_REGION": "northamerica-northeast1",
+            "ENVIRONMENT": "dev",
+        },
+        clear=False,
+    )
+    @patch("utils.create_http_task")
+    @patch("google.cloud.tasks_v2.CloudTasksClient")
+    def test_create_http_pmtiles_builder_task(
+        self, mock_client_cls, mock_create_http_task
+    ):
+        from utils import create_http_pmtiles_builder_task
+        import json
+
+        client_instance = MagicMock()
+        mock_client_cls.return_value = client_instance
+        stable_id = "feed-456"
+        dataset_stable_id = "dataset-def"
+        create_http_pmtiles_builder_task(
+            stable_id=stable_id, dataset_stable_id=dataset_stable_id
+        )
+        mock_client_cls.assert_called_once()
+        self.assertEqual(mock_create_http_task.call_count, 1)
+        args, _ = mock_create_http_task.call_args
+        payload = json.loads(args[1].decode("utf-8"))
+        self.assertEqual(
+            payload,
+            {"feed_stable_id": stable_id, "dataset_stable_id": dataset_stable_id},
+        )
+        self.assertEqual(
+            args[2],
+            "https://northamerica-northeast1-my-project.cloudfunctions.net/pmtiles-builder-dev",
+        )
+        self.assertEqual(args[3], "my-project")
+        self.assertEqual(args[4], "northamerica-northeast1")
+        self.assertEqual(args[5], "pmtiles-queue")

functions-python/helpers/utils.py

@@ -221,6 +221,35 @@ def create_http_task(
     )
 
 
+def create_http_pmtiles_builder_task(
+    stable_id: str,
+    dataset_stable_id: str,
+) -> None:
+    """
+    Create a task to generate PMTiles for a dataset.
+    """
+    from google.cloud import tasks_v2
+    import json
+
+    client = tasks_v2.CloudTasksClient()
+    body = json.dumps(
+        {"feed_stable_id": stable_id, "dataset_stable_id": dataset_stable_id}
+    ).encode()
+    queue_name = os.getenv("PMTILES_BUILDER_QUEUE")
+    project_id = os.getenv("PROJECT_ID")
+    gcp_region = os.getenv("GCP_REGION")
+    gcp_env = os.getenv("ENVIRONMENT")
+
+    create_http_task(
+        client,
+        body,
+        f"https://{gcp_region}-{project_id}.cloudfunctions.net/pmtiles-builder-{gcp_env}",
+        project_id,
+        gcp_region,
+        queue_name,
+    )
+
+
 def get_execution_id(json_payload: dict, stable_id: Optional[str]) -> str:
     """
     Extracts the execution_id from the JSON payload.

functions-python/tasks_executor/src/main.py

@@ -18,18 +18,22 @@
 
 import flask
 import functions_framework
+
 from shared.helpers.logger import init_logger
-from tasks.refresh_feedsearch_view.refresh_materialized_view import (
-    refresh_materialized_view_handler,
-)
 from tasks.dataset_files.rebuild_missing_dataset_files import (
     rebuild_missing_dataset_files_handler,
 )
+from tasks.missing_bounding_boxes.rebuild_missing_bounding_boxes import (
+    rebuild_missing_bounding_boxes_handler,
+)
+from tasks.refresh_feedsearch_view.refresh_materialized_view import (
+    refresh_materialized_view_handler,
+)
 from tasks.validation_reports.rebuild_missing_validation_reports import (
     rebuild_missing_validation_reports_handler,
 )
-from tasks.missing_bounding_boxes.rebuild_missing_bounding_boxes import (
-    rebuild_missing_bounding_boxes_handler,
+from tasks.visualization_files.rebuild_missing_visualization_files import (
+    rebuild_missing_visualization_files_handler,
 )
 
 init_logger()
@@ -62,6 +66,10 @@
         "description": "Rebuilds missing dataset files for GTFS datasets.",
         "handler": rebuild_missing_dataset_files_handler,
     },
+    "rebuild_missing_visualization_files": {
+        "description": "Rebuilds missing visualization files for GTFS datasets.",
+        "handler": rebuild_missing_visualization_files_handler,
+    },
 }