[smtr][brt] Migração da captura do BRT

pipelines/rj_smtr/br_rj_riodejaneiro_brt_gps/flows.py

@@ -19,27 +19,30 @@
 from pipelines.rj_smtr.constants import constants
 
 from pipelines.rj_smtr.schedules import (
-    # every_minute,
+    every_minute,
     every_hour,
 )
 from pipelines.rj_smtr.tasks import (
-    # create_current_date_hour_partition,
-    # create_local_partition_path,
+    create_date_hour_partition,
+    create_local_partition_path,
     fetch_dataset_sha,
+    get_current_timestamp,
     get_materialization_date_range,
     # get_local_dbt_client,
-    # get_raw,
-    # save_raw_local,
-    # save_treated_local,
+    get_raw,
+    parse_timestamp_to_string,
+    save_raw_local,
+    save_treated_local,
     set_last_run_timestamp,
-    # upload_logs_to_bq,
-    # bq_upload,
+    upload_logs_to_bq,
+    bq_upload,
 )
 from pipelines.utils.execute_dbt_model.tasks import run_dbt_model
 
-# from pipelines.rj_smtr.br_rj_riodejaneiro_brt_gps.tasks import (
-#     pre_treatment_br_rj_riodejaneiro_brt_gps,
-# )
+from pipelines.rj_smtr.br_rj_riodejaneiro_brt_gps.tasks import (
+    create_api_url_brt_gps,
+    pre_treatment_br_rj_riodejaneiro_brt_gps,
+)
 
 # Flows #
 
@@ -120,61 +123,61 @@
 materialize_brt.schedule = every_hour
 
 
-# with Flow(
-#     "SMTR: GPS BRT - Captura",
-#     code_owners=["caio", "fernanda"],
-# ) as captura_brt:
-
-#     # Rename flow run
-#     rename_flow_run = rename_current_flow_run_now_time(
-#         prefix="SMTR: GPS BRT - Captura - ", now_time=get_now_time()
-#     )
-
-#     # Get default parameters #
-#     dataset_id = Parameter("dataset_id", default=constants.GPS_BRT_RAW_DATASET_ID.value)
-#     table_id = Parameter("table_id", default=constants.GPS_BRT_RAW_TABLE_ID.value)
-#     url = Parameter("url", default=constants.GPS_BRT_API_BASE_URL.value)
-#     # secret_path = Parameter("secret_path", default=constants.GPS_BRT_API_SECRET_PATH.value)
-
-#     # Run tasks #
-#     file_dict = create_current_date_hour_partition()
-
-#     filepath = create_local_partition_path(
-#         dataset_id=dataset_id,
-#         table_id=table_id,
-#         filename=file_dict["filename"],
-#         partitions=file_dict["partitions"],
-#     )
-
-#     status_dict = get_raw(url=url)
-
-#     raw_filepath = save_raw_local(data=status_dict["data"], file_path=filepath)
-
-#     treated_status = pre_treatment_br_rj_riodejaneiro_brt_gps(status_dict=status_dict)
-
-#     upload_logs_to_bq(
-#         dataset_id=dataset_id,
-#         parent_table_id=table_id,
-#         timestamp=status_dict["timestamp"],
-#         error=status_dict["error"],
-#     )
-
-#     treated_filepath = save_treated_local(
-#         dataframe=treated_status["df"], file_path=filepath
-#     )
-
-#     bq_upload(
-#         dataset_id=dataset_id,
-#         table_id=table_id,
-#         filepath=treated_filepath,
-#         raw_filepath=raw_filepath,
-#         partitions=file_dict["partitions"],
-#     )
-#     captura_brt.set_dependencies(task=file_dict, upstream_tasks=[rename_flow_run])
-
-# captura_brt.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
-# captura_brt.run_config = KubernetesRun(
-#     image=emd_constants.DOCKER_IMAGE.value,
-#     labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
-# )
-# captura_brt.schedule = every_minute
+with Flow(
+    "SMTR: GPS BRT - Captura",
+    code_owners=["caio", "fernanda"],
+) as captura_brt:
+
+    timestamp = get_current_timestamp()
+
+    # Rename flow run
+    rename_flow_run = rename_current_flow_run_now_time(
+        prefix="SMTR: GPS BRT - Captura - ", now_time=timestamp
+    )
+
+    # SETUP LOCAL #
+    partitions = create_date_hour_partition(timestamp)
+
+    filename = parse_timestamp_to_string(timestamp)
+
+    filepath = create_local_partition_path(
+        dataset_id=constants.GPS_BRT_RAW_DATASET_ID.value,
+        table_id=constants.GPS_BRT_RAW_TABLE_ID.value,
+        filename=filename,
+        partitions=partitions,
+    )
+    # EXTRACT
+    url = create_api_url_brt_gps()
+
+    raw_status = get_raw(url=url)
+
+    raw_filepath = save_raw_local(status=raw_status, file_path=filepath)
+    # TREAT
+    treated_status = pre_treatment_br_rj_riodejaneiro_brt_gps(
+        status=raw_status, timestamp=timestamp
+    )
+
+    treated_filepath = save_treated_local(status=treated_status, file_path=filepath)
+    # LOAD
+    error = bq_upload(
+        dataset_id=constants.GPS_BRT_RAW_DATASET_ID.value,
+        table_id=constants.GPS_BRT_RAW_TABLE_ID.value,
+        filepath=treated_filepath,
+        raw_filepath=raw_filepath,
+        partitions=partitions,
+        status=treated_status,
+    )
+    upload_logs_to_bq(
+        dataset_id=constants.GPS_BRT_RAW_DATASET_ID.value,
+        parent_table_id=constants.GPS_BRT_RAW_TABLE_ID.value,
+        timestamp=timestamp,
+        error=error,
+    )
+    captura_brt.set_dependencies(task=partitions, upstream_tasks=[rename_flow_run])
+
+captura_brt.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
+captura_brt.run_config = KubernetesRun(
+    image=emd_constants.DOCKER_IMAGE.value,
+    labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
+)
+captura_brt.schedule = every_minute

pipelines/rj_smtr/br_rj_riodejaneiro_brt_gps/tasks.py

@@ -6,23 +6,37 @@
 from datetime import timedelta
 import traceback
 import pandas as pd
-import pendulum
 from prefect import task
 
 # EMD Imports #
 
-from pipelines.utils.utils import log
+from pipelines.utils.utils import get_vault_secret, log
 
 # SMTR Imports #
 
 from pipelines.rj_smtr.constants import constants
-from pipelines.rj_smtr.utils import log_critical, map_dict_keys, safe_cast
+from pipelines.rj_smtr.utils import log_critical, map_dict_keys
 
 # Tasks #
 
 
 @task
-def pre_treatment_br_rj_riodejaneiro_brt_gps(status_dict: dict):
+def create_api_url_brt_gps(secret_path: str = constants.GPS_BRT_SECRET_PATH.value):
+    """Create the url to request data from
+
+    Args:
+        secret_path (str, optional): secret path to the url.
+
+    Returns:
+        _str: url to request
+    """
+    url = get_vault_secret(secret_path=secret_path)["data"]["url"]
+    log(f"Request data from {url}")
+    return url
+
+
+@task
+def pre_treatment_br_rj_riodejaneiro_brt_gps(status: dict, timestamp):
     """Basic data treatment for brt gps data. Converts unix time to datetime,
     and apply filtering to stale data that may populate the API response.
 
@@ -35,82 +49,53 @@ def pre_treatment_br_rj_riodejaneiro_brt_gps(status_dict: dict):
     """
 
     # Check previous error
-    if status_dict["error"] is not None:
-        return {"df": pd.DataFrame(), "error": status_dict["error"]}
+    if status["error"] is not None:
+        return {"data": pd.DataFrame(), "error": status["error"]}
 
     error = None
-    data = status_dict["data"]
+    data = status["data"]
     timezone = constants.TIMEZONE.value
-    timestamp = status_dict["timestamp"]
-
+    log(
+        f"""
+    Received inputs:
+    - timestamp:\n{timestamp}
+    - data:\n{data[:10]}"""
+    )
     # Create dataframe sctructure
     key_column = "id_veiculo"
     columns = [key_column, "timestamp_gps", "timestamp_captura", "content"]
-    df = pd.DataFrame(data.json(), columns=columns)  # pylint: disable=c0103
-
-    # Populate dataframe
-    df["timestamp_captura"] = pd.to_datetime(timestamp)
+    df = pd.DataFrame(columns=columns)  # pylint: disable=c0103
 
     # map_dict_keys change data keys to match project data structure
-    map_keys = {
-        "vei_nro_gestor": "id_veiculo",
-        "linha": "servico",
-        "latitude": "latitude",
-        "longitude": "longitude",
-        "comunicacao": "timestamp_gps",
-        "velocidade": "velocidade",
-        "nomeItinerario": "sentido",
-    }
-    df["content"] = [map_dict_keys(piece, map_keys) for piece in data]
+    df["content"] = [
+        map_dict_keys(piece, constants.GPS_BRT_MAPPING_KEYS.value) for piece in data
+    ]
     df[key_column] = [piece[key_column] for piece in data]
+    df["timestamp_gps"] = [piece["timestamp_gps"] for piece in data]
+    df["timestamp_captura"] = timestamp
+    log(f"timestamp captura is:\n{df['timestamp_captura']}")
 
     # Remove timezone and force it to be config timezone
     log(f"Before converting, timestamp_gps is: \n{df['timestamp_gps']}")
-
-    df["timestamp_gps"] = df["content"].apply(
-        lambda x: pd.to_datetime(
-            pendulum.from_timestamp(safe_cast(x["timestamp_gps"], int, 0))
-            .replace(tzinfo=None)
-            .set(tz="UTC")
-            .isoformat()
-        )  # .tz_convert(timezone)
+    df["timestamp_gps"] = (
+        pd.to_datetime(df["timestamp_gps"], unit="s")
+        .dt.tz_localize("UTC")
+        .dt.tz_convert(timezone)
     )
-
     log(f"After converting the timezone, timestamp_gps is: \n{df['timestamp_gps']}")
 
     # Filter data for 0 <= time diff <= 1min
     try:
-        datahora_col = "timestamp_gps"
-        df_treated = df
-        try:
-            df_treated[datahora_col] = df_treated[datahora_col].apply(
-                lambda x: x.tz_convert(timezone)
-            )
-        except TypeError:
-            df_treated[datahora_col] = df_treated[datahora_col].apply(
-                lambda x: x.tz_localize(timezone)
-            )
-        try:
-            df_treated["timestamp_captura"] = df_treated["timestamp_captura"].apply(
-                lambda x: x.tz_convert(timezone)
-            )
-        except TypeError:
-            df_treated["timestamp_captura"] = df_treated["timestamp_captura"].apply(
-                lambda x: x.tz_localize(timezone)
-            )
-        mask = (df_treated["timestamp_captura"] - df_treated[datahora_col]).apply(
-            lambda x: timedelta(seconds=0) <= x <= timedelta(minutes=1)
-        )
-        df_treated = df_treated[mask]
         log(f"Shape antes da filtragem: {df.shape}")
-        log(f"Shape após a filtragem: {df_treated.shape}")
-        if df_treated.shape[0] == 0:
+
+        mask = (df["timestamp_captura"] - df["timestamp_gps"]) <= timedelta(minutes=1)
+        df = df[mask]  # pylint: disable=C0103
+        log(f"Shape após a filtragem: {df.shape}")
+        if df.shape[0] == 0:
             error = ValueError("After filtering, the dataframe is empty!")
             log_critical(f"@here\nFailed to filter BRT data: \n{error}")
-        df = df_treated  # pylint: disable=C0103
     except Exception:  # pylint: disable = W0703
         err = traceback.format_exc()
         log_critical(f"@here\nFailed to filter BRT data: \n{err}")
-    # log_critical(f"@here\n Got SPPO data at {timestamp} sucessfully")
 
-    return {"df": df, "error": error}
+    return {"data": df, "error": error}

pipelines/rj_smtr/constants.py

@@ -48,15 +48,22 @@ class constants(Enum):  # pylint: disable=c0103
     GPS_SPPO_CAPTURE_DELAY_V2 = 60
     GPS_SPPO_RECAPTURE_DELAY_V2 = 6
     # GPS BRT #
-    GPS_BRT_API_BASE_URL = (
-        "http://citgisbrj.tacom.srv.br:9977/gtfs-realtime-exporter/findAll/json"
-    )
-    # GPS_BRT_API_SECRET_PATH = "sppo_api"
-
+    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_MAPPING_KEYS = {
+        "vei_nro_gestor": "id_veiculo",
+        "linha": "servico",
+        "latitude": "latitude",
+        "longitude": "longitude",
+        "comunicacao": "timestamp_gps",
+        "velocidade": "velocidade",
+        "nomeItinerario": "sentido",
+        "nomeLinha": "vista",
+        "inicio_viagem": "timestamp_inicio_viagem",
+    }
 
     # SIGMOB (GTFS) #
     SIGMOB_GET_REQUESTS_TIMEOUT = 60

pipelines/rj_smtr/utils.py

@@ -154,11 +154,8 @@ def map_dict_keys(data: dict, mapping: dict) -> None:
     """
     Map old keys to new keys in a dict.
     """
-    pop_keys = [old_key for old_key in data.keys() if old_key not in mapping.keys()]
     for old_key, new_key in mapping.items():
         data[new_key] = data.pop(old_key)
-    for key in pop_keys:
-        data.pop(key)
     return data