From 5e54e080a057327e48fb976065b1fe14ed211bb3 Mon Sep 17 00:00:00 2001
From: t0momi219 <t.kodama@datumstudio.jp>
Date: Tue, 12 Nov 2024 16:20:14 +0900
Subject: [PATCH 1/2] Add task sla and timeout support.

---
 cosmos/airflow/graph.py | 17 +++++++++++++++++
 cosmos/config.py        |  2 ++
 2 files changed, 19 insertions(+)

diff --git a/cosmos/airflow/graph.py b/cosmos/airflow/graph.py
index f507b03ac..9242026ca 100644
--- a/cosmos/airflow/graph.py
+++ b/cosmos/airflow/graph.py
@@ -2,6 +2,8 @@
 
 from typing import Any, Callable, Union
 
+from datetime import timedelta
+
 from airflow.models import BaseOperator
 from airflow.models.dag import DAG
 from airflow.utils.task_group import TaskGroup
@@ -135,6 +137,8 @@ def create_task_metadata(
     dbt_dag_task_group_identifier: str,
     use_task_group: bool = False,
     source_rendering_behavior: SourceRenderingBehavior = SourceRenderingBehavior.NONE,
+    model_timeout: bool = False,
+    model_sla: bool = False,
 ) -> TaskMetadata | None:
     """
     Create the metadata that will be used to instantiate the Airflow Task used to run the Dbt node.
@@ -166,6 +170,11 @@ def create_task_metadata(
             task_id = f"{node.name}_run"
             if use_task_group is True:
                 task_id = "run"
+            if model_timeout and "model_timeout" in node.config.keys():
+                logger.error(f'model_timeout: {node.config["model_timeout"]} in values')
+                args["execution_timeout"] = timedelta(seconds=int(node.config["model_timeout"]))
+            if model_sla and "model_sla" in node.config.keys():
+                args["sla"] = timedelta(seconds=int(node.config["model_sla"]))
         elif node.resource_type == DbtResourceType.SOURCE:
             if (source_rendering_behavior == SourceRenderingBehavior.NONE) or (
                 source_rendering_behavior == SourceRenderingBehavior.WITH_TESTS_OR_FRESHNESS
@@ -217,6 +226,8 @@ def generate_task_or_group(
     source_rendering_behavior: SourceRenderingBehavior,
     test_indirect_selection: TestIndirectSelection,
     on_warning_callback: Callable[..., Any] | None,
+    model_timeout: bool,
+    model_sla: bool,
     **kwargs: Any,
 ) -> BaseOperator | TaskGroup | None:
     task_or_group: BaseOperator | TaskGroup | None = None
@@ -234,6 +245,8 @@ def generate_task_or_group(
         dbt_dag_task_group_identifier=_get_dbt_dag_task_group_identifier(dag, task_group),
         use_task_group=use_task_group,
         source_rendering_behavior=source_rendering_behavior,
+        model_timeout=model_timeout,
+        model_sla=model_sla
     )
 
     # In most cases, we'll  map one DBT node to one Airflow task
@@ -335,6 +348,8 @@ def build_airflow_graph(
     node_converters = render_config.node_converters or {}
     test_behavior = render_config.test_behavior
     source_rendering_behavior = render_config.source_rendering_behavior
+    model_timeout = render_config.model_timeout
+    model_sla = render_config.model_sla
     tasks_map = {}
     task_or_group: TaskGroup | BaseOperator
 
@@ -356,6 +371,8 @@ def build_airflow_graph(
             source_rendering_behavior=source_rendering_behavior,
             test_indirect_selection=test_indirect_selection,
             on_warning_callback=on_warning_callback,
+            model_timeout=model_timeout,
+            model_sla=model_sla,
             node=node,
         )
         if task_or_group is not None:
diff --git a/cosmos/config.py b/cosmos/config.py
index 516a6787b..0275189c4 100644
--- a/cosmos/config.py
+++ b/cosmos/config.py
@@ -80,6 +80,8 @@ class RenderConfig:
     enable_mock_profile: bool = True
     source_rendering_behavior: SourceRenderingBehavior = SourceRenderingBehavior.NONE
     airflow_vars_to_purge_dbt_ls_cache: list[str] = field(default_factory=list)
+    model_timeout: bool = False
+    model_sla: bool = False
 
     def __post_init__(self, dbt_project_path: str | Path | None) -> None:
         if self.env_vars:

From 054e29118f1075f32d6d42b2040034db6c43b045 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 12 Nov 2024 07:36:16 +0000
Subject: [PATCH 2/2] =?UTF-8?q?=F0=9F=8E=A8=20[pre-commit.ci]=20Auto=20for?=
 =?UTF-8?q?mat=20from=20pre-commit.com=20hooks?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 cosmos/airflow/graph.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/cosmos/airflow/graph.py b/cosmos/airflow/graph.py
index 9242026ca..0f135bfa5 100644
--- a/cosmos/airflow/graph.py
+++ b/cosmos/airflow/graph.py
@@ -1,8 +1,7 @@
 from __future__ import annotations
 
-from typing import Any, Callable, Union
-
 from datetime import timedelta
+from typing import Any, Callable, Union
 
 from airflow.models import BaseOperator
 from airflow.models.dag import DAG
@@ -246,7 +245,7 @@ def generate_task_or_group(
         use_task_group=use_task_group,
         source_rendering_behavior=source_rendering_behavior,
         model_timeout=model_timeout,
-        model_sla=model_sla
+        model_sla=model_sla,
     )
 
     # In most cases, we'll  map one DBT node to one Airflow task