[WIP] Refatora recapturas de GPS do SPPO

pipelines/rj_smtr/br_rj_riodejaneiro_brt_gps/flows.py

@@ -12,7 +12,12 @@
 from pipelines.constants import constants as emd_constants
 from pipelines.utils.decorators import Flow
 from pipelines.utils.execute_dbt_model.tasks import get_k8s_dbt_client
-from pipelines.utils.tasks import get_now_time, rename_current_flow_run_now_time
+from pipelines.utils.tasks import (
+    get_now_time,
+    rename_current_flow_run_now_time,
+    get_current_flow_mode,
+    get_current_flow_labels,
+)
 
 # SMTR Imports #
 
@@ -67,8 +72,11 @@
     table_id = Parameter("table_id", default=constants.GPS_BRT_TREATED_TABLE_ID.value)
     rebuild = Parameter("rebuild", False)
 
+    LABELS = get_current_flow_labels()
+    MODE = get_current_flow_mode(LABELS)
+
     # Set dbt client #
-    dbt_client = get_k8s_dbt_client(mode="prod", wait=rename_flow_run)
+    dbt_client = get_k8s_dbt_client(mode=MODE, wait=rename_flow_run)
     # Use the command below to get the dbt client in dev mode:
     # dbt_client = get_local_dbt_client(host="localhost", port=3001)
 
@@ -79,6 +87,8 @@
         raw_dataset_id=raw_dataset_id,
         raw_table_id=raw_table_id,
         table_date_column_name="data",
+        mode=MODE,
+        delay_hours=constants.GPS_BRT_MATERIALIZE_DELAY_HOURS.value,
     )
     dataset_sha = fetch_dataset_sha(
         dataset_id=dataset_id,
@@ -99,6 +109,7 @@
             table_id=table_id,
             timestamp=date_range["date_range_end"],
             wait=RUN,
+            mode=MODE,
         )
     with case(rebuild, False):
         RUN = run_dbt_model(
@@ -113,12 +124,13 @@
             table_id=table_id,
             timestamp=date_range["date_range_end"],
             wait=RUN,
+            mode=MODE,
         )
 
 materialize_brt.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
 materialize_brt.run_config = KubernetesRun(
     image=emd_constants.DOCKER_IMAGE.value,
-    labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
+    labels=[emd_constants.RJ_SMTR_DEV_AGENT_LABEL.value],
 )
 materialize_brt.schedule = every_hour
 

pipelines/rj_smtr/br_rj_riodejaneiro_onibus_gps/flows.py

@@ -7,6 +7,7 @@
 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.tasks.control_flow import merge
 from prefect.utilities.edges import unmapped
 
 
@@ -38,7 +39,7 @@
     get_materialization_date_range,
     # get_local_dbt_client,
     get_raw,
-    parse_timestamp_to_string,
+    get_bool,
     query_logs,
     save_raw_local,
     save_treated_local,
@@ -90,6 +91,7 @@
         raw_table_id=raw_table_id,
         table_date_column_name="data",
         mode=MODE,
+        delay_hours=constants.GPS_SPPO_MATERIALIZE_DELAY_HOURS.value,
     )
     dataset_sha = fetch_dataset_sha(
         dataset_id=dataset_id,
@@ -131,7 +133,7 @@
 materialize_sppo.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
 materialize_sppo.run_config = KubernetesRun(
     image=emd_constants.DOCKER_IMAGE.value,
-    labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
+    labels=[emd_constants.RJ_SMTR_DEV_AGENT_LABEL.value],
 )
 
 
@@ -142,25 +144,33 @@
 
     version = Parameter("version", default=2)
 
+    # RECAPTURE PARAMETERS
+    datetime_filter = Parameter("datetime_filter", default=None)
+    recapture = Parameter("recapture", default=False)
+    previous_error = Parameter("previous_error", default=None)
+
     # SETUP #
-    timestamp = get_current_timestamp()
+    with case(get_bool(datetime_filter), False):
+        timestamp_now = get_current_timestamp()
+    with case(get_bool(datetime_filter), True):
+        timestamp_default = datetime_filter
+
+    timestamp = merge(timestamp_default, timestamp_now)
 
     rename_flow_run = rename_current_flow_run_now_time(
         prefix="GPS SPPO: ", now_time=timestamp
     )
 
     partitions = create_date_hour_partition(timestamp)
 
-    filename = parse_timestamp_to_string(timestamp)
-
     filepath = create_local_partition_path(
         dataset_id=constants.GPS_SPPO_RAW_DATASET_ID.value,
         table_id=constants.GPS_SPPO_RAW_TABLE_ID.value,
-        filename=filename,
+        filename=timestamp,
         partitions=partitions,
     )
 
-    url = create_api_url_onibus_gps(version=version)
+    url = create_api_url_onibus_gps(version=version, timestamp=timestamp)
 
     # EXTRACT #
     raw_status = get_raw(url)
@@ -183,18 +193,24 @@
         partitions=partitions,
         status=treated_status,
     )
-
     upload_logs_to_bq(
         dataset_id=constants.GPS_SPPO_RAW_DATASET_ID.value,
         parent_table_id=constants.GPS_SPPO_RAW_TABLE_ID.value,
         error=error,
+        previous_error=previous_error,
         timestamp=timestamp,
+        recapture=recapture,
     )
 
+
 captura_sppo_v2.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
 captura_sppo_v2.run_config = KubernetesRun(
     image=emd_constants.DOCKER_IMAGE.value,
-    labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
+    labels=[emd_constants.RJ_SMTR_DEV_AGENT_LABEL.value],
+    cpu_limit="50m",
+    cpu_request="10m",
+    memory_limit="256Mi",
+    memory_request="64Mi",
 )
 captura_sppo_v2.schedule = every_minute
 
@@ -205,20 +221,22 @@
     datetime_filter = Parameter("datetime_filter", default=None)
     # SETUP #
     LABELS = get_current_flow_labels()
-    errors, timestamps = query_logs(
+    errors, recapture_parameters = query_logs(
         dataset_id=constants.GPS_SPPO_RAW_DATASET_ID.value,
         table_id=constants.GPS_SPPO_RAW_TABLE_ID.value,
         datetime_filter=datetime_filter,
     )
 
     rename_flow_run = rename_current_flow_run_now_time(
-        prefix="GPS SPPO Recapturas: ", now_time=get_now_time(), wait=timestamps
+        prefix="GPS SPPO Recapturas: ",
+        now_time=get_now_time(),
+        wait=recapture_parameters,
     )
 
     with case(errors, False):
         materialize_no_errors = create_flow_run(
             flow_name=materialize_sppo.name,
-            project_name=emd_constants.PREFECT_DEFAULT_PROJECT.value,
+            project_name="staging",
             labels=LABELS,
             run_name=materialize_sppo.name,
         )
@@ -229,69 +247,35 @@
             raise_final_state=True,
         )
     with case(errors, True):
-        # SETUP #
-        partitions = create_date_hour_partition.map(timestamps)
-        filename = parse_timestamp_to_string.map(timestamps)
-
-        filepath = create_local_partition_path.map(
-            dataset_id=unmapped(constants.GPS_SPPO_RAW_DATASET_ID.value),
-            table_id=unmapped(constants.GPS_SPPO_RAW_TABLE_ID.value),
-            filename=filename,
-            partitions=partitions,
-        )
-
-        url = create_api_url_onibus_gps.map(
-            version=unmapped(version), timestamp=timestamps
-        )
-
-        # EXTRACT #
-        raw_status = get_raw.map(url)
-
-        raw_filepath = save_raw_local.map(status=raw_status, file_path=filepath)
-
-        # # CLEAN #
-        trated_status = pre_treatment_br_rj_riodejaneiro_onibus_gps.map(
-            status=raw_status, timestamp=timestamps, version=unmapped(version)
-        )
-
-        treated_filepath = save_treated_local.map(
-            status=trated_status, file_path=filepath
+        recaptura_runs = create_flow_run.map(
+            flow_name=unmapped(captura_sppo_v2.name),
+            project_name=unmapped("staging"),
+            labels=unmapped(LABELS),
+            run_name=unmapped(captura_sppo_v2.name),
+            parameters=recapture_parameters,
         )
-
-        # # LOAD #
-        error = bq_upload.map(
-            dataset_id=unmapped(constants.GPS_SPPO_RAW_DATASET_ID.value),
-            table_id=unmapped(constants.GPS_SPPO_RAW_TABLE_ID.value),
-            filepath=treated_filepath,
-            raw_filepath=raw_filepath,
-            partitions=partitions,
-            status=trated_status,
-        )
-
-        UPLOAD_LOGS = upload_logs_to_bq.map(
-            dataset_id=unmapped(constants.GPS_SPPO_RAW_DATASET_ID.value),
-            parent_table_id=unmapped(constants.GPS_SPPO_RAW_TABLE_ID.value),
-            error=error,
-            timestamp=timestamps,
-            recapture=unmapped(True),
+        wait_for_recapturas = wait_for_flow_run.map(
+            recaptura_runs,
+            stream_states=unmapped(True),
+            stream_logs=unmapped(True),
+            raise_final_state=unmapped(True),
         )
         materialize = create_flow_run(
             flow_name=materialize_sppo.name,
-            project_name=emd_constants.PREFECT_DEFAULT_PROJECT.value,
+            project_name="staging",
             labels=LABELS,
             run_name=materialize_sppo.name,
         )
-        wait_materialize = wait_for_flow_run(
-            materialize,
-            stream_states=True,
-            stream_logs=True,
-            raise_final_state=True,
+        wait_for_materialize = wait_for_flow_run(
+            materialize, stream_states=True, stream_logs=True, raise_final_state=True
         )
-    recaptura.set_dependencies(task=materialize, upstream_tasks=[UPLOAD_LOGS])
-
+    materialize.set_upstream(wait_for_recapturas)
 recaptura.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
 recaptura.run_config = KubernetesRun(
     image=emd_constants.DOCKER_IMAGE.value,
-    labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
+    labels=[emd_constants.RJ_SMTR_DEV_AGENT_LABEL.value],
+    memory_limit="256Mi",
+    memory_request="64Mi",
+    cpu_limit="50m",
 )
 recaptura.schedule = every_hour_minute_six

pipelines/rj_smtr/br_rj_riodejaneiro_onibus_gps/tasks.py

@@ -22,7 +22,7 @@
 
 
 @task
-def create_api_url_onibus_gps(version: str, timestamp: datetime = None) -> str:
+def create_api_url_onibus_gps(version: str, timestamp: str = None) -> str:
     """
     Generates the complete URL to get data from API.
     """
@@ -38,6 +38,8 @@ def create_api_url_onibus_gps(version: str, timestamp: datetime = None) -> str:
         timestamp = pendulum.now(constants.TIMEZONE.value).replace(
             second=0, microsecond=0
         )
+    else:
+        timestamp = datetime.fromisoformat(timestamp)
 
     headers = get_vault_secret(source)["data"]
     key = list(headers)[0]
@@ -83,6 +85,7 @@ def pre_treatment_br_rj_riodejaneiro_onibus_gps(
     timezone = constants.TIMEZONE.value
 
     log(f"Data received to treat: \n{status['data'][:5]}")
+
     df = pd.DataFrame(status["data"])  # pylint: disable=c0103
     df["timestamp_captura"] = timestamp
     log(f"Before converting, datahora is: \n{df['datahora']}")

pipelines/rj_smtr/constants.py

@@ -47,12 +47,14 @@ class constants(Enum):  # pylint: disable=c0103
     GPS_SPPO_CAPTURE_DELAY_V1 = 1
     GPS_SPPO_CAPTURE_DELAY_V2 = 60
     GPS_SPPO_RECAPTURE_DELAY_V2 = 6
+    GPS_SPPO_MATERIALIZE_DELAY_HOURS = 1
     # GPS BRT #
     GPS_BRT_SECRET_PATH = "brt_api"
     GPS_BRT_DATASET_ID = "br_rj_riodejaneiro_veiculos"
     GPS_BRT_RAW_DATASET_ID = "br_rj_riodejaneiro_brt_gps"
     GPS_BRT_RAW_TABLE_ID = "registros"
     GPS_BRT_TREATED_TABLE_ID = "gps_brt"
+    GPS_BRT_MATERIALIZE_DELAY_HOURS = 0
     GPS_BRT_MAPPING_KEYS = {
         "vei_nro_gestor": "id_veiculo",
         "linha": "servico",

pipelines/rj_smtr/schedules.py

@@ -59,7 +59,7 @@
                 2021, 1, 1, 0, 6, 0, tzinfo=timezone(constants.TIMEZONE.value)
             ),
             labels=[
-                emd_constants.RJ_SMTR_AGENT_LABEL.value,
+                emd_constants.RJ_SMTR_DEV_AGENT_LABEL.value,
             ],
         ),
     ]