Adiciona pipeline de captura e tratamento do gtfs

pipelines/rj_smtr/__init__.py

@@ -17,3 +17,4 @@
 from pipelines.rj_smtr.veiculo.flows import *
 from pipelines.rj_smtr.example.flows import *
 from pipelines.rj_smtr.br_rj_riodejaneiro_bilhetagem.flows import *
+from pipelines.rj_smtr.br_rj_riodejaneiro_gtfs.flows import *

pipelines/rj_smtr/br_rj_riodejaneiro_gtfs/flows.py

@@ -0,0 +1,135 @@
+# -*- coding: utf-8 -*-
+"""
+Flows for gtfs
+"""
+from copy import deepcopy
+
+# Imports #
+
+from prefect.run_configs import KubernetesRun
+from prefect.storage import GCS
+from prefect.tasks.prefect import create_flow_run, wait_for_flow_run
+from prefect.utilities.edges import unmapped
+from prefect import Parameter, case, task
+from prefect.tasks.control_flow import merge
+
+
+# EMD Imports #
+
+from pipelines.constants import constants as emd_constants
+from pipelines.utils.utils import set_default_parameters
+from pipelines.utils.decorators import Flow
+from pipelines.utils.tasks import (
+    rename_current_flow_run_now_time,
+    get_current_flow_labels,
+)
+
+# SMTR Imports #
+from pipelines.rj_smtr.constants import constants
+from pipelines.rj_smtr.tasks import (
+    get_current_timestamp,
+)
+
+from pipelines.rj_smtr.flows import default_capture_flow, default_materialization_flow
+
+# SETUP dos Flows
+
+gtfs_captura = deepcopy(default_capture_flow)
+gtfs_captura.name = "SMTR - Captura dos dados do GTFS"
+gtfs_captura.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
+gtfs_captura.run_config = KubernetesRun(
+    image=emd_constants.DOCKER_IMAGE.value,
+    labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
+)
+gtfs_captura.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
+gtfs_captura = set_default_parameters(
+    flow=gtfs_captura,
+    default_parameters=constants.GTFS_GENERAL_CAPTURE_PARAMS.value,
+)
+
+gtfs_materializacao = deepcopy(default_materialization_flow)
+gtfs_materializacao.name = "SMTR - Materialização dos dados do GTFS"
+gtfs_materializacao.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
+gtfs_materializacao.run_config = KubernetesRun(
+    image=emd_constants.DOCKER_IMAGE.value,
+    labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
+)
+gtfs_materializacao = set_default_parameters(
+    flow=gtfs_materializacao,
+    default_parameters=constants.GTFS_MATERIALIZACAO_PARAMS.value,
+)
+
+with Flow(
+    "SMTR: GTFS - Captura/Tratamento",
+    code_owners=["rodrigo", "carolinagomes"],
+) as gtfs_captura_tratamento:
+    # SETUP
+
+    data_versao_gtfs = Parameter("data_versao_gtfs", default=None)
+    capture = Parameter("capture", default=True)
+    materialize = Parameter("materialize", default=True)
+
+    timestamp = get_current_timestamp()
+
+    rename_flow_run = rename_current_flow_run_now_time(
+        prefix=gtfs_captura_tratamento.name + " ",
+        now_time=timestamp,
+    )
+
+    LABELS = get_current_flow_labels()
+
+    with case(capture, True):
+        gtfs_capture_parameters = [
+            {"timestamp": data_versao_gtfs, **d}
+            for d in constants.GTFS_TABLE_CAPTURE_PARAMS.value
+        ]
+
+        run_captura = create_flow_run.map(
+            flow_name=unmapped(gtfs_captura.name),
+            project_name=unmapped(emd_constants.PREFECT_DEFAULT_PROJECT.value),
+            parameters=gtfs_capture_parameters,
+            labels=unmapped(LABELS),
+        )
+
+        wait_captura_true = wait_for_flow_run.map(
+            run_captura,
+            stream_states=unmapped(True),
+            stream_logs=unmapped(True),
+            raise_final_state=unmapped(True),
+        )
+
+    with case(capture, False):
+        wait_captura_false = task(
+            lambda: [None], checkpoint=False, name="assign_none_to_previous_runs"
+        )()
+
+    wait_captura = merge(wait_captura_true, wait_captura_false)
+
+    with case(materialize, True):
+        gtfs_materializacao_parameters = {
+            "dbt_vars": {
+                "data_versao_gtfs": data_versao_gtfs,
+                "version": {},
+            },
+        }
+
+        run_materializacao = create_flow_run(
+            flow_name=gtfs_materializacao.name,
+            project_name=emd_constants.PREFECT_DEFAULT_PROJECT.value,
+            parameters=gtfs_materializacao_parameters,
+            labels=LABELS,
+            upstream_tasks=[wait_captura],
+        )
+
+        wait_materializacao = wait_for_flow_run(
+            run_materializacao,
+            stream_states=True,
+            stream_logs=True,
+            raise_final_state=True,
+        )
+
+gtfs_captura_tratamento.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
+gtfs_captura_tratamento.run_config = KubernetesRun(
+    image=emd_constants.DOCKER_IMAGE.value,
+    labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
+)

pipelines/rj_smtr/constants.py

@@ -323,3 +323,78 @@ class constants(Enum):  # pylint: disable=c0103
             "version": {},
         },
     }
+
+    # GTFS
+    GTFS_DATASET_ID = "br_rj_riodejaneiro_gtfs"
+
+    GTFS_GENERAL_CAPTURE_PARAMS = {
+        "partition_date_only": True,
+        "source_type": "gcs",
+        "dataset_id": "br_rj_riodejaneiro_gtfs",
+        "extract_params": {"filename": "gtfs"},
+        "partition_date_name": "data_versao",
+    }
+
+    GTFS_TABLE_CAPTURE_PARAMS = [
+        {
+            "table_id": "agency",
+            "primary_key": ["agency_id"],
+        },
+        {
+            "table_id": "calendar_dates",
+            "primary_key": ["service_id", "date"],
+        },
+        {
+            "table_id": "calendar",
+            "primary_key": ["service_id"],
+        },
+        {
+            "table_id": "feed_info",
+            "primary_key": ["feed_publisher_name"],
+        },
+        {
+            "table_id": "frequencies",
+            "primary_key": ["trip_id", "start_time"],
+        },
+        {
+            "table_id": "routes",
+            "primary_key": ["route_id"],
+        },
+        {
+            "table_id": "shapes",
+            "primary_key": ["shape_id", "shape_pt_sequence"],
+        },
+        {
+            "table_id": "stops",
+            "primary_key": ["stop_id"],
+        },
+        {
+            "table_id": "stop_times",
+            "primary_key": ["trip_id", "stop_sequence"],
+        },
+        {
+            "table_id": "trips",
+            "primary_key": ["trip_id"],
+        },
+        {
+            "table_id": "fare_attributes",
+            "primary_key": ["fare_id"],
+        },
+        {
+            "table_id": "fare_rules",
+            "primary_key": [],
+        },
+        {
+            "table_id": "ordem_servico",
+            "primary_key": ["servico"],
+            "extract_params": {"filename": "ordem_servico"},
+        },
+    ]
+
+    GTFS_MATERIALIZACAO_PARAMS = {
+        "dataset_id": GTFS_DATASET_ID,
+        "dbt_vars": {
+            "data_versao_gtfs": "",
+            "version": {},
+        },
+    }

pipelines/rj_smtr/flows.py

@@ -60,6 +60,8 @@
     interval_minutes = Parameter("interval_minutes", default=None)
     recapture = Parameter("recapture", default=False)
     recapture_window_days = Parameter("recapture_window_days", default=1)
+    timestamp = Parameter("timestamp", default=None)
+    partition_date_name = Parameter("partition_date_name", default=None)
 
     # Parâmetros Pré-tratamento #
     primary_key = Parameter("primary_key", default=None)
@@ -70,7 +72,9 @@
         checkpoint=False,
     )
 
-    current_timestamp = get_rounded_timestamp(interval_minutes=interval_minutes)
+    current_timestamp = get_rounded_timestamp(
+        timestamp=timestamp, interval_minutes=interval_minutes
+    )
 
     with case(recapture, True):
         _, recapture_timestamps, recapture_previous_errors = query_logs(
@@ -96,7 +100,9 @@
     )
 
     partitions = create_date_hour_partition.map(
-        timestamps, partition_date_only=unmapped(partition_date_only)
+        timestamps,
+        partition_date_name=unmapped(partition_date_name),
+        partition_date_only=unmapped(partition_date_only),
     )
 
     filenames = parse_timestamp_to_string.map(timestamps)

pipelines/rj_smtr/tasks.py

@@ -166,8 +166,8 @@ def create_dbt_run_vars(
 
     log(f"Creating DBT variables. Parameter received: {dbt_vars}")
 
-    if (not dbt_vars) or (not table_id):
-        log("dbt_vars or table_id are blank. Skiping task")
+    if not dbt_vars:
+        log("dbt_vars are blank. Skiping task...")
         return [None], None, False
 
     final_vars = []
@@ -187,6 +187,10 @@ def create_dbt_run_vars(
             }
         # Create date_range using Redis
         else:
+            if not table_id:
+                log("table_id are blank. Skiping task...")
+                return [None], None, False
+
             raw_table_id = raw_table_id or table_id
 
             date_var = get_materialization_date_range.run(
@@ -230,6 +234,16 @@ def create_dbt_run_vars(
 
         log(f"version created: {dataset_sha}")
 
+    if "data_versao_gtfs" in dbt_vars.keys():
+        log("Creating data_versao_gtfs variable")
+
+        temp_dict = {"data_versao_gtfs": dbt_vars["data_versao_gtfs"]}
+
+        if final_vars:
+            final_vars = get_join_dict.run(dict_list=final_vars, new_dict=temp_dict)
+        else:
+            final_vars.append(temp_dict)
+
     log(f"All variables was created, final value is: {final_vars}")
 
     return final_vars, date_var, flag_date_range
@@ -302,19 +316,22 @@ def get_current_timestamp(timestamp=None, truncate_minute: bool = True) -> datet
 
 @task
 def create_date_hour_partition(
-    timestamp: datetime, partition_date_only: bool = False
+    timestamp: datetime,
+    partition_date_name: str = "data",
+    partition_date_only: bool = False,
 ) -> str:
     """
     Create a date (and hour) Hive partition structure from timestamp.
 
     Args:
         timestamp (datetime): timestamp to be used as reference
+        partition_date_name (str, optional): partition name. Defaults to "data".
         partition_date_only (bool, optional): whether to add hour partition or not
 
     Returns:
         str: partition string
     """
-    partition = f"data={timestamp.strftime('%Y-%m-%d')}"
+    partition = f"{partition_date_name}={timestamp.strftime('%Y-%m-%d')}"
     if not partition_date_only:
         partition += f"/hora={timestamp.strftime('%H')}"
     return partition
@@ -388,11 +405,16 @@ def save_treated_local(file_path: str, status: dict, mode: str = "staging") -> s
     Returns:
         str: Path to the saved file
     """
+
+    log(f"Saving treated data to: {file_path}, {status}")
+
     _file_path = file_path.format(mode=mode, filetype="csv")
+
     Path(_file_path).parent.mkdir(parents=True, exist_ok=True)
     if status["error"] is None:
         status["data"].to_csv(_file_path, index=False)
         log(f"Treated data saved to: {_file_path}")
+
     return _file_path
 
 
@@ -545,7 +567,7 @@ def get_raw(  # pylint: disable=R0912
         params (dict, optional): Params to be sent on request
 
     Returns:
-        dict: Conatining keys
+        dict: Containing keys
           * `data` (json): data result
           * `error` (str): catched error, if any. Otherwise, returns None
     """
@@ -636,6 +658,9 @@ def create_request_params(
             "query": extract_params["query"].format(**datetime_range),
         }
 
+    elif dataset_id == constants.GTFS_DATASET_ID.value:
+        request_params = extract_params["filename"]
+
     return request_params, request_url
 
 

pipelines/rj_smtr/utils.py

@@ -581,11 +581,13 @@ def get_upload_storage_blob(
         Blob: blob object
     """
     bucket = bd.Storage(dataset_id="", table_id="")
+    log(f"Filename: {filename}, dataset_id: {dataset_id}")
     blob_list = list(
         bucket.client["storage_staging"]
         .bucket(bucket.bucket_name)
         .list_blobs(prefix=f"upload/{dataset_id}/{filename}.")
     )
+
     return blob_list[0]