From 79add9ac98203aafd6c901d197bcc099f420ecb2 Mon Sep 17 00:00:00 2001
From: Kiyan Mair <kiyan.mair@gmail.com>
Date: Sun, 14 Jul 2024 13:40:54 -0400
Subject: [PATCH] feat: Add task_kwargs param to dbt_flow

---
 prefect_dbt_flow/dbt/tasks.py | 3 +++
 prefect_dbt_flow/flow.py      | 3 +++
 2 files changed, 6 insertions(+)

diff --git a/prefect_dbt_flow/dbt/tasks.py b/prefect_dbt_flow/dbt/tasks.py
index a4aebc7..9e48705 100644
--- a/prefect_dbt_flow/dbt/tasks.py
+++ b/prefect_dbt_flow/dbt/tasks.py
@@ -214,6 +214,7 @@ def generate_tasks_dag(
     profile: Optional[DbtProfile],
     dag_options: Optional[DbtDagOptions],
     dbt_graph: List[DbtNode],
+    task_kwargs: Optional[Dict] = None,
     run_test_after_model: bool = False,
 ) -> None:
     """
@@ -224,6 +225,7 @@ def generate_tasks_dag(
         profile: A class that represents a dbt profile configuration.
         dag_options: A class to add dbt DAG configurations.
         dbt_graph: A list of dbt nodes (models) to include in the DAG.
+        task_kwargs: Additional task configuration for running each model.
         run_test_after_model: If True, run tests after running each model.
 
     Returns:
@@ -237,6 +239,7 @@ def generate_tasks_dag(
             profile=profile,
             dag_options=dag_options,
             dbt_node=dbt_node,
+            task_kwargs=task_kwargs,
         )
         for dbt_node in dbt_graph
     }
diff --git a/prefect_dbt_flow/flow.py b/prefect_dbt_flow/flow.py
index 81fe8f7..2a986ef 100644
--- a/prefect_dbt_flow/flow.py
+++ b/prefect_dbt_flow/flow.py
@@ -11,6 +11,7 @@ def dbt_flow(
     profile: Optional[DbtProfile] = None,
     dag_options: Optional[DbtDagOptions] = None,
     flow_kwargs: Optional[dict] = None,
+    task_kwargs: Optional[dict] = None,
 ) -> Flow:
     """
     Create a PrefectFlow for executing a dbt project.
@@ -20,6 +21,7 @@ def dbt_flow(
         profile: A Class that represents a dbt profile configuration.
         dag_options: A Class to add dbt DAG configurations.
         flow_kwargs: A dict of prefect @flow arguments
+        task_kwargs: A dict of Prefect @task arguments
 
     Returns:
         dbt_flow: A Prefec Flow.
@@ -44,6 +46,7 @@ def dbt_flow():
             profile,
             dag_options,
             dbt_graph,
+            task_kwargs,
             dag_options.run_test_after_model if dag_options else False,
         )