Cria novo pull para a pipeline de recursos do SPPO

.gitignore

@@ -13,7 +13,7 @@ setup.py
 .vscode/*
 *.hdf
 *.DS_Store
-
+.idea/*
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

pipelines/rj_smtr/__init__.py

@@ -14,6 +14,7 @@
 from pipelines.rj_smtr.materialize_to_datario.flows import *
 from pipelines.rj_smtr.registros_ocr_rir.flows import *
 from pipelines.rj_smtr.projeto_subsidio_sppo.flows import *
+from pipelines.rj_smtr.br_rj_riodejaneiro_recurso.flows import *
 from pipelines.rj_smtr.veiculo.flows import *
 from pipelines.rj_smtr.example.flows import *
 from pipelines.rj_smtr.br_rj_riodejaneiro_bilhetagem.flows import *

pipelines/rj_smtr/br_rj_riodejaneiro_recurso/flows.py

@@ -0,0 +1,179 @@
+# -*- coding: utf-8 -*-
+"""
+Flows for br_rj_riodejaneiro_recurso
+"""
+from copy import deepcopy
+
+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 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
+from pipelines.rj_smtr.schedules import every_day
+
+
+# CAPTURA #
+
+sppo_recurso_captura = deepcopy(default_capture_flow)
+sppo_recurso_captura.name = "SMTR: Subsídio SPPO Recursos - Captura (subflow)"
+sppo_recurso_captura.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
+sppo_recurso_captura.run_config = KubernetesRun(
+    image=emd_constants.DOCKER_IMAGE.value,
+    labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
+)
+sppo_recurso_captura = set_default_parameters(
+    flow=sppo_recurso_captura,
+    default_parameters=constants.SUBSIDIO_SPPO_RECURSO_CAPTURE_PARAMS.value,
+)
+# RECAPTURA #
+sppo_recurso_recaptura = deepcopy(default_capture_flow)
+sppo_recurso_recaptura.name = "SMTR: Subsídio SPPO Recursos - Recaptura (subflow)"
+sppo_recurso_recaptura.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
+sppo_recurso_recaptura.run_config = KubernetesRun(
+    image=emd_constants.DOCKER_IMAGE.value,
+    labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
+)
+sppo_recurso_recaptura = set_default_parameters(
+    flow=sppo_recurso_recaptura,
+    default_parameters=constants.SUBSIDIO_SPPO_RECURSO_CAPTURE_PARAMS.value
+    | {"recapture": True},
+)
+
+# MATERIALIZAÇÃO #
+
+sppo_recurso_materializacao = deepcopy(default_materialization_flow)
+sppo_recurso_materializacao.name = (
+    "SMTR: Subsídio SPPO Recursos - Materialização (subflow)"
+)
+sppo_recurso_materializacao.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
+sppo_recurso_materializacao.run_config = KubernetesRun(
+    image=emd_constants.DOCKER_IMAGE.value,
+    labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
+)
+
+sppo_recurso_materializacao = set_default_parameters(
+    flow=sppo_recurso_materializacao,
+    default_parameters=constants.SUBSIDIO_SPPO_RECURSOS_MATERIALIZACAO_PARAMS.value,
+)
+
+with Flow(
+    "SMTR: Subsídio Recursos Viagens Individuais - Captura/Tratamento",
+    code_owners=["carolinagomes", "rafaelpinheiro"],
+) as subsidio_sppo_recurso:
+    capture = Parameter("capture", default=True)
+    materialize = Parameter("materialize", default=True)
+    recapture = Parameter("recapture", default=True)
+    data_recurso = Parameter("data_recurso", default=None)
+    interval_minutes = Parameter("interval_minutes", default=1440)
+    timestamp = get_current_timestamp(data_recurso, return_str=True)
+
+    rename_flow_run = rename_current_flow_run_now_time(
+        prefix=subsidio_sppo_recurso.name + " ",
+        now_time=timestamp,
+    )
+
+    LABELS = get_current_flow_labels()
+
+    # Captura
+
+    with case(capture, True):
+        run_captura = create_flow_run(
+            flow_name=sppo_recurso_captura.name,
+            project_name=emd_constants.PREFECT_DEFAULT_PROJECT.value,
+            parameters={"timestamp": timestamp},
+            labels=LABELS,
+        )
+
+        wait_captura_true = wait_for_flow_run(
+            run_captura,
+            stream_states=True,
+            stream_logs=True,
+            raise_final_state=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)
+
+    # Recaptura
+
+    with case(recapture, True):
+        run_recaptura = create_flow_run(
+            flow_name=sppo_recurso_recaptura.name,
+            project_name=emd_constants.PREFECT_DEFAULT_PROJECT.value,
+            labels=LABELS,
+        )
+        run_recaptura.set_upstream(wait_captura)
+
+        wait_recaptura_true = wait_for_flow_run(
+            run_recaptura,
+            stream_states=True,
+            stream_logs=True,
+            raise_final_state=True,
+        )
+
+    with case(recapture, False):
+        wait_recaptura_false = task(
+            lambda: [None], checkpoint=False, name="assign_none_to_previous_runs"
+        )()
+
+    wait_recaptura = merge(wait_recaptura_true, wait_recaptura_false)
+
+    # Materialização
+
+    with case(materialize, True):
+        run_materializacao = create_flow_run(
+            flow_name=sppo_recurso_materializacao.name,
+            project_name=emd_constants.PREFECT_DEFAULT_PROJECT.value,
+            labels=LABELS,
+            upstream_tasks=[wait_captura],
+        )
+
+        run_materializacao.set_upstream(wait_recaptura)
+
+        wait_materializacao_true = wait_for_flow_run(
+            run_materializacao,
+            stream_states=True,
+            stream_logs=True,
+            raise_final_state=True,
+        )
+
+    with case(materialize, False):
+        wait_materializacao_false = task(
+            lambda: [None], checkpoint=False, name="assign_none_to_previous_runs"
+        )()
+    wait_materializacao = merge(wait_materializacao_true, wait_materializacao_false)
+
+    subsidio_sppo_recurso.set_reference_tasks(
+        [wait_materializacao, wait_recaptura, wait_captura]
+    )
+
+subsidio_sppo_recurso.storage = GCS(emd_constants.GCS_FLOWS_BUCKET.value)
+subsidio_sppo_recurso.run_config = KubernetesRun(
+    image=emd_constants.DOCKER_IMAGE.value,
+    labels=[emd_constants.RJ_SMTR_AGENT_LABEL.value],
+)
+
+# Schedule
+subsidio_sppo_recurso.schedule = every_day

pipelines/rj_smtr/constants.py

@@ -527,3 +527,37 @@ class constants(Enum):  # pylint: disable=c0103
             "version": {},
         },
     }
+
+    # SUBSÍDIO RECURSOS VIAGENS INDIVIDUAIS
+    SUBSIDIO_SPPO_RECURSOS_DATASET_ID = "br_rj_riodejaneiro_recurso"
+    SUBSIDIO_SPPO_RECURSO_API_BASE_URL = "https://api.movidesk.com/public/v1/tickets?"
+    SUBSIDIO_SPPO_RECURSO_API_SECRET_PATH = "sppo_subsidio_recursos_api"
+    SUBSIDIO_SPPO_RECURSO_SERVICE = "serviceFull eq 'SPPO'"
+    SUBSIDIO_SPPO_RECURSO_CAPTURE_PARAMS = {
+        "partition_date_only": True,
+        "table_id": "recurso_sppo",
+        "dataset_id": SUBSIDIO_SPPO_RECURSOS_DATASET_ID,
+        "extract_params": {
+            "token": "",
+            "$select": "id,protocol,createdDate",
+            "$filter": "{dates} and serviceFull/any(serviceFull: {service})",
+            "$expand": "customFieldValues,customFieldValues($expand=items)",
+            "$orderby": "createdDate asc",
+        },
+        "interval_minutes": 1440,
+        "source_type": "movidesk",
+        "primary_key": ["protocol"],
+    }
+
+    SUBSIDIO_SPPO_RECURSOS_MATERIALIZACAO_PARAMS = {
+        "dataset_id": SUBSIDIO_SPPO_RECURSOS_DATASET_ID,
+        "table_id": SUBSIDIO_SPPO_RECURSO_CAPTURE_PARAMS["table_id"],
+        "upstream": True,
+        "dbt_vars": {
+            "date_range": {
+                "table_run_datetime_column_name": "data_recurso",
+                "delay_hours": 0,
+            },
+            "version": {},
+        },
+    }

pipelines/rj_smtr/tasks.py

@@ -31,6 +31,7 @@
     get_raw_data_api,
     get_raw_data_gcs,
     get_raw_data_db,
+    get_raw_recursos,
     upload_run_logs_to_bq,
     get_datetime_range,
     read_raw_data,
@@ -252,7 +253,7 @@ def create_dbt_run_vars(
 
 ###############
 #
-# Local file managment
+# Local file management
 #
 ###############
 
@@ -669,6 +670,24 @@ def create_request_params(
     elif dataset_id == constants.GTFS_DATASET_ID.value:
         request_params = extract_params["filename"]
 
+    elif dataset_id == constants.SUBSIDIO_SPPO_RECURSOS_DATASET_ID.value:
+        extract_params["token"] = get_vault_secret(
+            constants.SUBSIDIO_SPPO_RECURSO_API_SECRET_PATH.value
+        )["data"]["token"]
+        start = datetime.strftime(
+            timestamp - timedelta(minutes=interval_minutes), "%Y-%m-%dT%H:%M:%S.%MZ"
+        )
+        end = datetime.strftime(timestamp, "%Y-%m-%dT%H:%M:%S.%MZ")
+        log(f" Start date {start}, end date {end}")
+        recurso_params = {
+            "dates": f"createdDate ge {start} and createdDate le {end}",
+            "service": constants.SUBSIDIO_SPPO_RECURSO_SERVICE.value,
+        }
+        extract_params["$filter"] = extract_params["$filter"].format(**recurso_params)
+        request_params = extract_params
+
+        request_url = constants.SUBSIDIO_SPPO_RECURSO_API_BASE_URL.value
+
     return request_params, request_url
 
 
@@ -726,6 +745,10 @@ def get_raw_from_sources(
             error, data, filetype = get_raw_data_db(
                 host=source_path, secret_path=secret_path, **request_params
             )
+        elif source_type == "movidesk":
+            error, data, filetype = get_raw_recursos(
+                request_url=source_path, request_params=request_params
+            )
         else:
             raise NotImplementedError(f"{source_type} not supported")
 
@@ -1275,15 +1298,17 @@ def transform_raw_to_nested_structure(
             # Check empty dataframe
             if data.empty:
                 log("Empty dataframe, skipping transformation...")
+
             else:
                 log(f"Raw data:\n{data_info_str(data)}", level="info")
 
                 log("Adding captured timestamp column...", level="info")
                 data["timestamp_captura"] = timestamp
 
-                log("Striping string columns...", level="info")
-                for col in data.columns[data.dtypes == "object"].to_list():
-                    data[col] = data[col].str.strip()
+                if "customFieldValues" not in data:
+                    log("Striping string columns...", level="info")
+                    for col in data.columns[data.dtypes == "object"].to_list():
+                        data[col] = data[col].str.strip()
 
                 log(f"Finished cleaning! Data:\n{data_info_str(data)}", level="info")
 

pipelines/rj_smtr/utils.py

@@ -17,17 +17,21 @@
 import requests
 import basedosdados as bd
 from basedosdados import Table
+from basedosdados import Storage
+import math
 import pandas as pd
 from google.cloud.storage.blob import Blob
 import pymysql
 import psycopg2
 import psycopg2.extras
+import time
 
 
 from prefect.schedules.clocks import IntervalClock
 
 from pipelines.constants import constants as emd_constants
 
+
 from pipelines.rj_smtr.implicit_ftp import ImplicitFtpTls
 from pipelines.rj_smtr.constants import constants
 
@@ -828,3 +832,57 @@ def read_raw_data(filepath: str, csv_args: dict = None) -> tuple[str, pd.DataFra
         log(f"[CATCHED] Task failed with error: \n{error}", level="error")
 
     return error, data
+
+
+def get_raw_recursos(request_url: str, request_params: dict) -> tuple[str, str, str]:
+    """
+    Returns a dataframe with recursos data from movidesk api.
+    """
+    all_records = False
+    top = 1000
+    skip = 0
+    error = None
+    filetype = "json"
+    data = []
+
+    while not all_records:
+        try:
+            request_params["$top"] = top
+            request_params["$skip"] = skip
+
+            log(f"Request url {request_url}")
+
+            response = requests.get(
+                request_url,
+                params=request_params,
+                timeout=constants.MAX_TIMEOUT_SECONDS.value,
+            )
+            response.raise_for_status()
+
+            paginated_data = response.json()
+
+            if isinstance(paginated_data, dict):
+                paginated_data = [paginated_data]
+
+            if len(paginated_data) == top:
+                skip += top
+                time.sleep(36)
+
+            if len(paginated_data) == 0:
+                log("Nenhum dado para tratar.")
+
+            else:
+                all_records = True
+            data += paginated_data
+
+            log(f"Dados (paginados): {len(data)}")
+
+        except Exception as error:
+            error = traceback.format_exc()
+            log(f"[CATCHED] Task failed with error: \n{error}", level="error")
+            data = []
+            break
+
+    log(f"Request concluído, tamanho dos dados: {len(data)}.")
+
+    return error, data, filetype