add dependencies to utils, fix imports, add init to templates

prefeitura-rio · Dec 14, 2023 · aa357c5 · aa357c5
1 parent f184556
commit aa357c5
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 4 deletions.
diff --git a/pipelines/br_rj_riodejaneiro_bilhetagem/flows.py b/pipelines/br_rj_riodejaneiro_bilhetagem/flows.py
@@ -29,7 +29,7 @@
 
 # SMTR Imports #
 
-from pipelines.flows import (
+from pipelines.templates.flows import (
     default_capture_flow,
     default_materialization_flow,
 )

diff --git a/pipelines/templates/__init__.py b/pipelines/templates/__init__.py
diff --git a/pipelines/utils/utils.py b/pipelines/utils/utils.py
@@ -19,13 +19,12 @@
 from basedosdados import Table
 from basedosdados import Storage
 from pytz import timezone
-import math
+import prefect
 import pandas as pd
 from google.cloud.storage.blob import Blob
 import pymysql
 import psycopg2
 import psycopg2.extras
-from redis_pal import RedisPal
 import time
 
 
@@ -956,4 +955,15 @@ def generate_ftp_schedules(
                     labels=[label],
                 )
             )
-    return clocks
+    return clocks
+
+def set_default_parameters(
+    flow: prefect.Flow, default_parameters: dict
+) -> prefect.Flow:
+    """
+    Sets default parameters for a flow.
+    """
+    for parameter in flow.parameters():
+        if parameter.name in default_parameters:
+            parameter.default = default_parameters[parameter.name]
+    return flow