From 7de8f0fa8dba4667208591cd64cdf43ab38cee9b Mon Sep 17 00:00:00 2001
From: hellcassius <caiorogerio.santos@gmail.com>
Date: Wed, 13 Dec 2023 14:14:00 -0300
Subject: [PATCH] organize code, add default get_secret func

---
 pipelines/br_rj_riodejaneiro_brt_gps/tasks.py |  2 +-
 pipelines/schedules.py                        |  2 +-
 pipelines/tasks.py                            | 13 +++++-----
 pipelines/utils/__init__.py                   |  0
 pipelines/{ => utils}/implicit_ftp.py         |  0
 pipelines/utils/secret.py                     | 25 +++++++++++++++++++
 pipelines/{ => utils}/utils.py                |  5 +++-
 7 files changed, 38 insertions(+), 9 deletions(-)
 create mode 100644 pipelines/utils/__init__.py
 rename pipelines/{ => utils}/implicit_ftp.py (100%)
 create mode 100644 pipelines/utils/secret.py
 rename pipelines/{ => utils}/utils.py (99%)

diff --git a/pipelines/br_rj_riodejaneiro_brt_gps/tasks.py b/pipelines/br_rj_riodejaneiro_brt_gps/tasks.py
index ff4973a0..7cfe41f5 100644
--- a/pipelines/br_rj_riodejaneiro_brt_gps/tasks.py
+++ b/pipelines/br_rj_riodejaneiro_brt_gps/tasks.py
@@ -15,7 +15,7 @@
 # SMTR Imports #
 
 from pipelines.constants import constants
-from pipelines.utils import log_critical, map_dict_keys
+from pipelines.utils.utils import log_critical, map_dict_keys
 
 # Tasks #
 
diff --git a/pipelines/schedules.py b/pipelines/schedules.py
index 117043cb..4312feb9 100644
--- a/pipelines/schedules.py
+++ b/pipelines/schedules.py
@@ -8,7 +8,7 @@
 from prefect.schedules import Schedule
 from prefect.schedules.clocks import IntervalClock, CronClock
 from pipelines.constants import constants as emd_constants
-from pipelines.utils import generate_ftp_schedules
+from pipelines.utils.utils import generate_ftp_schedules
 from pipelines.constants import constants
 
 every_minute = Schedule(
diff --git a/pipelines/tasks.py b/pipelines/tasks.py
index bed66006..6884947f 100644
--- a/pipelines/tasks.py
+++ b/pipelines/tasks.py
@@ -22,7 +22,7 @@
 import requests
 
 from pipelines.constants import constants
-from pipelines.utils import (
+from pipelines.utils.utils import (
     create_or_append_table,
     bq_project,
     get_table_min_max_value,
@@ -39,9 +39,11 @@
     save_treated_local_func,
     save_raw_local_func,
     log_critical,
+    normalize_keys
 )
+from pipelines.utils.secret import get_secret
 from prefeitura_rio.pipelines_utils.dbt import run_dbt_model
-from prefeitura_rio.pipelines_utils.infisical import get_secret, inject_bd_credentials
+from prefeitura_rio.pipelines_utils.infisical import inject_bd_credentials
 from prefeitura_rio.pipelines_utils.logging import log
 
 ###############
@@ -588,8 +590,7 @@ def get_raw(  # pylint: disable=R0912
 
     try:
         if headers is not None:
-            headers = get_secret(headers)["data"]
-
+            headers = get_secret(secret_path=headers)
             # remove from headers, if present
             remove_headers = ["host", "databases"]
             for remove_header in remove_headers:
@@ -675,8 +676,8 @@ def create_request_params(
 
     elif dataset_id == constants.SUBSIDIO_SPPO_RECURSOS_DATASET_ID.value:
         extract_params["token"] = get_secret(
-            constants.SUBSIDIO_SPPO_RECURSO_API_SECRET_PATH.value
-        )["data"]["token"]
+            secret_path=constants.SUBSIDIO_SPPO_RECURSO_API_SECRET_PATH.value
+        )["token"]
         start = datetime.strftime(
             timestamp - timedelta(minutes=interval_minutes), "%Y-%m-%dT%H:%M:%S.%MZ"
         )
diff --git a/pipelines/utils/__init__.py b/pipelines/utils/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/pipelines/implicit_ftp.py b/pipelines/utils/implicit_ftp.py
similarity index 100%
rename from pipelines/implicit_ftp.py
rename to pipelines/utils/implicit_ftp.py
diff --git a/pipelines/utils/secret.py b/pipelines/utils/secret.py
new file mode 100644
index 00000000..e4120538
--- /dev/null
+++ b/pipelines/utils/secret.py
@@ -0,0 +1,25 @@
+from prefeitura_rio.pipelines_utils.infisical import get_infisical_client
+from pipelines.utils.utils import normalize_keys
+
+
+def get_secret(secret_path:str='/', secret_name:str=None, environment:str='dev'):
+    """
+    Fetches secrets from Infisical. If passing only `secret_path` and
+    no `secret_name`, returns all secrets inside a folder.
+
+    Args:
+        secret_name (str, optional): _description_. Defaults to None.
+        secret_path (str, optional): _description_. Defaults to '/'.
+        environment (str, optional): _description_. Defaults to 'dev'.
+
+    Returns:
+        _type_: _description_
+    """
+    client = get_infisical_client()
+    if not secret_path.startswith('/'):
+        secret_path = f"/{secret_path}"
+    if secret_path and not secret_name:
+        secrets = client.get_all_secrets(path=secret_path)
+        return normalize_keys({s.secret_name:s.secret_value for s in secrets})
+    secret = client.get_secret(secret_name=secret_name, path=secret_path, environment=environment)
+    return {secret_name:secret.secret_value}
\ No newline at end of file
diff --git a/pipelines/utils.py b/pipelines/utils/utils.py
similarity index 99%
rename from pipelines/utils.py
rename to pipelines/utils/utils.py
index 1319a284..70e40ad0 100644
--- a/pipelines/utils.py
+++ b/pipelines/utils/utils.py
@@ -34,7 +34,7 @@
 from pipelines.constants import constants
 
 
-from pipelines.implicit_ftp import ImplicitFtpTls
+from pipelines.utils.implicit_ftp import ImplicitFtpTls
 from pipelines.constants import constants
 
 from prefeitura_rio.pipelines_utils.logging import log #TODO: add or relocate imports
@@ -214,6 +214,9 @@ def map_dict_keys(data: dict, mapping: dict) -> None:
         data[new_key] = data.pop(old_key)
     return data
 
+def normalize_keys(data:dict):
+    _data = {key.lower():value for key, value in data.items()}
+    return _data
 
 def connect_ftp(secret_path: str = None, secure: bool = True):
     """Connect to FTP