added pipeline tasks to batch processing

Author: cka-y

functions-python/batch_process_dataset/src/main.py

@@ -41,6 +41,7 @@
     get_hash_from_file,
     download_from_gcs,
 )
+from pipeline_tasks import create_pipeline_tasks
 
 init_logger()
 
@@ -300,7 +301,8 @@ def process_from_bucket_latest(self, public=True) -> DatasetFile or None:
                     else None
                 ),
             )
-            self.create_dataset_entities(dataset_file, skip_dataset_creation=True)
+            dataset = self.create_dataset_entities(dataset_file, skip_dataset_creation=True)
+            create_pipeline_tasks(dataset)
         finally:
             if temp_file_path and os.path.exists(temp_file_path):
                 os.remove(temp_file_path)
@@ -352,6 +354,7 @@ def create_dataset_entities(
             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:
                 dataset = Gtfsdataset(
                     id=str(uuid.uuid4()),
@@ -394,6 +397,7 @@ def create_dataset_entities(
             self.logger.info(f"[{self.feed_stable_id}] Dataset created successfully.")
 
             create_refresh_materialized_view_task()
+            return latest_dataset if skip_dataset_creation else dataset
         except Exception as e:
             raise Exception(f"Error creating dataset: {e}")
 
@@ -407,7 +411,8 @@ def process_from_producer_url(self) -> DatasetFile or None:
         if dataset_file is None:
             self.logger.info(f"[{self.feed_stable_id}] No database update required.")
             return None
-        self.create_dataset_entities(dataset_file)
+        dataset = self.create_dataset_entities(dataset_file)
+        create_pipeline_tasks(dataset)
         return dataset_file
 
 

functions-python/batch_process_dataset/src/pipeline_tasks.py

@@ -0,0 +1,79 @@
+import json
+import os
+
+from google.cloud import tasks_v2
+
+from shared.database_gen.sqlacodegen_models import Gtfsdataset
+from shared.helpers.utils import create_http_task
+
+
+def create_http_reverse_geolocation_processor_task(
+    stable_id: str,
+    dataset_stable_id: str,
+    stops_url: str,
+) -> None:
+    """
+    Create a task to process reverse geolocation for a dataset.
+    """
+    client = tasks_v2.CloudTasksClient()
+    body = json.dumps(
+        {"stable_id": stable_id, "stops_url": stops_url, "dataset_id": dataset_stable_id}
+    ).encode()
+    queue_name = os.getenv("REVERSE_GEOLOCATION_QUEUE_NAME")
+    project_id = os.getenv("PROJECT_ID")
+    gcp_region = os.getenv("GCP_REGION")
+
+    create_http_task(
+        client,
+        body,
+        f"https://{gcp_region}-{project_id}.cloudfunctions.net/reverse-geolocation-processor",
+        project_id,
+        gcp_region,
+        queue_name,
+    )
+
+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_NAME")
+    project_id = os.getenv("PROJECT_ID")
+    gcp_region = os.getenv("GCP_REGION")
+
+    create_http_task(
+        client,
+        body,
+        f"https://{gcp_region}-{project_id}.cloudfunctions.net/pmtiles_builder",
+        project_id,
+        gcp_region,
+        queue_name,
+    )
+
+
+def create_pipeline_tasks(
+    dataset: Gtfsdataset
+) -> None:
+    """
+    Create pipeline tasks for a dataset.
+    """
+    stable_id = dataset.feed.stable_id
+    dataset_stable_id = dataset.stable_id
+    gtfs_files = dataset.gtfsfiles
+    stops_file = next(
+        (file for file in gtfs_files if file.file_name == "stops.txt"), None
+    )
+    stops_url = stops_file.hosted_url if stops_file else None
+
+    # Create reverse geolocation task
+    if stops_url:
+        create_http_reverse_geolocation_processor_task(stable_id, dataset_stable_id, stops_url)
+
+    # Create PMTiles builder task
+    create_http_pmtiles_builder_task(stable_id, dataset_stable_id)

functions-python/batch_process_dataset/tests/test_pipeline_tasks.py

@@ -0,0 +1 @@
+# TODO

functions-python/tasks_executor/requirements.txt

@@ -26,5 +26,3 @@ google-cloud-storage
 
 # Configuration
 python-dotenv==1.0.0
-tippecanoe
-