Dump Serpro

pipelines/serpro/flows.py

@@ -1,19 +1,39 @@
 # -*- coding: utf-8 -*-
+from prefect import Parameter
 from prefect.run_configs import KubernetesRun
 from prefect.storage import GCS
+
+# from prefect.utilities.edges import unmapped
 from prefeitura_rio.pipelines_utils.custom import Flow
+from prefeitura_rio.pipelines_utils.state_handlers import handler_inject_bd_credentials
 
 from pipelines.constants import constants as smtr_constants
-from pipelines.serpro.tasks import wait_sleeping
+
+# from pipelines.migration.tasks import upload_raw_data_to_gcs
+from pipelines.serpro.tasks import dump_serpro, get_db_object, wait_sleeping
 from pipelines.serpro.utils import handler_setup_serpro
 
 with Flow("SMTR - Teste Conexão Serpro") as flow:
+    batch_size = Parameter("batch_size", default=100000)
     # setup_serpro()
+
+    jdbc = get_db_object()
+    csv_files = dump_serpro(jdbc, batch_size)
+
     wait_sleeping()
 
+    # upload_raw_data_to_gcs.map(
+    #     dataset_id=unmapped("radar_serpro"),
+    #     table_id=unmapped("tb_infracao_view"),
+    #     raw_filepath=csv_files,
+    #     partitions=unmapped(None),
+    #     error=unmapped(None),
+    #     bucket_name=unmapped("rj-smtr-dev"),
+    # )
+
 flow.storage = GCS(smtr_constants.GCS_FLOWS_BUCKET.value)
 flow.run_config = KubernetesRun(
     image=smtr_constants.DOCKER_IMAGE_FEDORA.value,
     labels=[smtr_constants.RJ_SMTR_AGENT_LABEL.value],
 )
-flow.state_handlers = [handler_setup_serpro]
+flow.state_handlers = [handler_setup_serpro, handler_inject_bd_credentials]

pipelines/serpro/tasks.py

@@ -1,16 +1,81 @@
 # -*- coding: utf-8 -*-
+import csv
+import os
+from datetime import timedelta
 from time import sleep
+from typing import List
 
 from prefect import task
 
 from pipelines.utils.jdbc import JDBC
+from pipelines.utils.secret import get_secret
+from pipelines.utils.utils import log
 
 
 @task
 def wait_sleeping(interval_seconds: int = 54000, wait=None):
     sleep(interval_seconds)
 
 
-@task
+@task(checkpoint=False, max_retries=3, retry_delay=timedelta(seconds=20))
 def get_db_object(secret_path="radar_serpro", environment: str = "dev"):
+    jar_path = get_secret(secret_path=secret_path, environment=environment)["jars"]
+
+    if not os.path.exists(jar_path):
+        raise Exception(f"Arquivo JAR '{jar_path}' não encontrado.")
+
     return JDBC(db_params_secret_path=secret_path, environment=environment)
+
+
+@task
+def dump_serpro(jdbc: JDBC, batch_size: int) -> List[str]:
+
+    index = 0
+    data_folder = os.getenv("DATA_FOLDER", "data")
+    file_path = f"{os.getcwd()}/{data_folder}/raw/radar_serpro/tb_infracao_view/"
+    csv_files = []
+
+    os.makedirs(os.path.dirname(file_path), exist_ok=True)
+
+    query = """            SELECT
+                *
+            FROM
+                dbpro_radar_view_SMTR_VBL.tb_infracao_view
+            WHERE
+                SUBSTRING(auinf_dt_infracao, 1, 10) = '2024-10-28'
+            """
+
+    jdbc.execute_query(query)
+
+    columns = jdbc.get_columns()
+
+    while True:
+        rows = jdbc.fetch_batch(batch_size)
+
+        if not rows:
+            break
+
+        output_file = file_path + f"dados_infracao_{index}.csv"
+
+        with open(output_file, "w", newline="") as csvfile:
+            writer = csv.writer(csvfile)
+            writer.writerow(columns)
+            writer.writerows(rows)
+
+        log(output_file)
+        csv_files.append(output_file)
+        index += 1
+
+    jdbc.close_connection()
+
+    return csv_files
+
+
+@task
+def list_files():
+
+    try:
+        files = os.listdir("/app")
+        log(f"Files: {files}")
+    except Exception as e:
+        log(f"Erro: {e}")

pipelines/utils/jdbc.py

@@ -5,6 +5,8 @@
 
 from pipelines.utils.secret import get_secret
 
+# from prefeitura_rio.pipelines_utils.logging import log
+
 
 class JDBC:
     def __init__(self, db_params_secret_path: str, environment: str = "staging") -> None:
@@ -68,3 +70,10 @@ def fetch_all(self) -> List[List]:
         Fetches all rows from the JDBC database.
         """
         return [list(item) for item in self._cursor.fetchall()]
+
+    def close_connection(self):
+        """
+        Closes the JDBC connection.
+        """
+        if self._connection:
+            self._connection.close()