Cast callbacks to functions when set with default_args on task groups

dagfactory/__init__.py

@@ -1,2 +1,3 @@
 """Modules and methods to export for easier access"""
+
 from .dagfactory import DagFactory, load_yaml_dags

dagfactory/__version__.py

@@ -1,2 +1,3 @@
 """Module contains the version of dag-factory"""
+
 __version__ = "0.19.0"

dagfactory/dagbuilder.py

@@ -1,4 +1,5 @@
 """Module contains code for generating tasks and constructing a DAG"""
+
 # pylint: disable=ungrouped-imports
 import os
 import re
@@ -197,34 +198,34 @@ def get_dag_params(self) -> Dict[str, Any]:
 
         if utils.check_dict_key(dag_params["default_args"], "sla_miss_callback"):
             if isinstance(dag_params["default_args"]["sla_miss_callback"], str):
-                dag_params["default_args"][
-                    "sla_miss_callback"
-                ]: Callable = import_string(
-                    dag_params["default_args"]["sla_miss_callback"]
+                dag_params["default_args"]["sla_miss_callback"]: Callable = (
+                    import_string(dag_params["default_args"]["sla_miss_callback"])
+                )
+
+        if utils.check_dict_key(
+            dag_params["default_args"], "on_execute_callback"
+        ) and version.parse(AIRFLOW_VERSION) >= version.parse("2.0.0"):
+            if isinstance(dag_params["default_args"]["on_execute_callback"], str):
+                dag_params["default_args"]["on_execute_callback"]: Callable = (
+                    import_string(dag_params["default_args"]["on_execute_callback"])
                 )
 
         if utils.check_dict_key(dag_params["default_args"], "on_success_callback"):
             if isinstance(dag_params["default_args"]["on_success_callback"], str):
-                dag_params["default_args"][
-                    "on_success_callback"
-                ]: Callable = import_string(
-                    dag_params["default_args"]["on_success_callback"]
+                dag_params["default_args"]["on_success_callback"]: Callable = (
+                    import_string(dag_params["default_args"]["on_success_callback"])
                 )
 
         if utils.check_dict_key(dag_params["default_args"], "on_failure_callback"):
             if isinstance(dag_params["default_args"]["on_failure_callback"], str):
-                dag_params["default_args"][
-                    "on_failure_callback"
-                ]: Callable = import_string(
-                    dag_params["default_args"]["on_failure_callback"]
+                dag_params["default_args"]["on_failure_callback"]: Callable = (
+                    import_string(dag_params["default_args"]["on_failure_callback"])
                 )
 
         if utils.check_dict_key(dag_params["default_args"], "on_retry_callback"):
             if isinstance(dag_params["default_args"]["on_retry_callback"], str):
-                dag_params["default_args"][
-                    "on_retry_callback"
-                ]: Callable = import_string(
-                    dag_params["default_args"]["on_retry_callback"]
+                dag_params["default_args"]["on_retry_callback"]: Callable = (
+                    import_string(dag_params["default_args"]["on_retry_callback"])
                 )
 
         if utils.check_dict_key(dag_params, "sla_miss_callback"):
@@ -351,11 +352,11 @@ def make_task(operator: str, task_params: Dict[str, Any]) -> BaseOperator:
                         "  python_callable_file: !!python/name:my_module.my_func"
                     )
                 if not task_params.get("python_callable"):
-                    task_params[
-                        "python_callable"
-                    ]: Callable = utils.get_python_callable(
-                        task_params["python_callable_name"],
-                        task_params["python_callable_file"],
+                    task_params["python_callable"]: Callable = (
+                        utils.get_python_callable(
+                            task_params["python_callable_name"],
+                            task_params["python_callable_file"],
+                        )
                     )
                     # remove dag-factory specific parameters
                     # Airflow 2.0 doesn't allow these to be passed to operator
@@ -419,10 +420,10 @@ def make_task(operator: str, task_params: Dict[str, Any]) -> BaseOperator:
                     del task_params["response_check_name"]
                     del task_params["response_check_file"]
                 else:
-                    task_params[
-                        "response_check"
-                    ]: Callable = utils.get_python_callable_lambda(
-                        task_params["response_check_lambda"]
+                    task_params["response_check"]: Callable = (
+                        utils.get_python_callable_lambda(
+                            task_params["response_check_lambda"]
+                        )
                     )
                     # remove dag-factory specific parameters
                     # Airflow 2.0 doesn't allow these to be passed to operator
@@ -636,6 +637,63 @@ def make_task_groups(
             for task_group_name, task_group_conf in task_groups.items():
                 task_group_conf["group_id"] = task_group_name
                 task_group_conf["dag"] = dag
+
+                if version.parse(AIRFLOW_VERSION) >= version.parse(
+                    "2.2.0"
+                ) and isinstance(task_group_conf.get("default_args"), dict):
+                    # https://github.com/apache/airflow/pull/16557
+                    if utils.check_dict_key(
+                        task_group_conf["default_args"], "on_success_callback"
+                    ):
+                        if isinstance(
+                            task_group_conf["default_args"]["on_success_callback"],
+                            str,
+                        ):
+                            task_group_conf["default_args"][
+                                "on_success_callback"
+                            ]: Callable = import_string(
+                                task_group_conf["default_args"]["on_success_callback"]
+                            )
+
+                    if utils.check_dict_key(
+                        task_group_conf["default_args"], "on_execute_callback"
+                    ):
+                        if isinstance(
+                            task_group_conf["default_args"]["on_execute_callback"],
+                            str,
+                        ):
+                            task_group_conf["default_args"][
+                                "on_execute_callback"
+                            ]: Callable = import_string(
+                                task_group_conf["default_args"]["on_execute_callback"]
+                            )
+
+                    if utils.check_dict_key(
+                        task_group_conf["default_args"], "on_failure_callback"
+                    ):
+                        if isinstance(
+                            task_group_conf["default_args"]["on_failure_callback"],
+                            str,
+                        ):
+                            task_group_conf["default_args"][
+                                "on_failure_callback"
+                            ]: Callable = import_string(
+                                task_group_conf["default_args"]["on_failure_callback"]
+                            )
+
+                    if utils.check_dict_key(
+                        task_group_conf["default_args"], "on_retry_callback"
+                    ):
+                        if isinstance(
+                            task_group_conf["default_args"]["on_retry_callback"],
+                            str,
+                        ):
+                            task_group_conf["default_args"][
+                                "on_retry_callback"
+                            ]: Callable = import_string(
+                                task_group_conf["default_args"]["on_retry_callback"]
+                            )
+
                 task_group = TaskGroup(
                     **{
                         k: v
@@ -669,18 +727,18 @@ def set_dependencies(
                 group_id = conf["task_group"].group_id
                 name = f"{group_id}.{name}"
             if conf.get("dependencies"):
-                source: Union[
-                    BaseOperator, "TaskGroup"
-                ] = tasks_and_task_groups_instances[name]
+                source: Union[BaseOperator, "TaskGroup"] = (
+                    tasks_and_task_groups_instances[name]
+                )
                 for dep in conf["dependencies"]:
                     if tasks_and_task_groups_config[dep].get("task_group"):
                         group_id = tasks_and_task_groups_config[dep][
                             "task_group"
                         ].group_id
                         dep = f"{group_id}.{dep}"
-                    dep: Union[
-                        BaseOperator, "TaskGroup"
-                    ] = tasks_and_task_groups_instances[dep]
+                    dep: Union[BaseOperator, "TaskGroup"] = (
+                        tasks_and_task_groups_instances[dep]
+                    )
                     source.set_upstream(dep)
 
     @staticmethod

dagfactory/dagfactory.py

@@ -1,4 +1,5 @@
 """Module contains code for loading a DagFactory config and generating DAGs"""
+
 import logging
 import os
 from itertools import chain

dagfactory/utils.py

@@ -1,4 +1,5 @@
 """Module contains various utilities used by dag-factory"""
+
 import ast
 import importlib.util
 import os
@@ -212,9 +213,7 @@ def check_template_searchpath(template_searchpath: Union[str, List[str]]) -> boo
     return False
 
 
-def get_expand_partial_kwargs(
-    task_params: Dict[str, Any]
-) -> Tuple[
+def get_expand_partial_kwargs(task_params: Dict[str, Any]) -> Tuple[
     Dict[str, Any],
     Dict[str, Union[Dict[str, Any], Any]],
     Dict[str, Union[Dict[str, Any], Any]],

tests/test_dagbuilder.py

@@ -134,6 +134,44 @@
         },
     },
 }
+DAG_CONFIG_TASK_GROUP_WITH_CALLBACKS = {
+    "default_args": {"owner": "custom_owner"},
+    "schedule_interval": "0 3 * * *",
+    "task_groups": {
+        "task_group_1": {
+            "tooltip": "this is a task group",
+            "default_args": {
+                "on_failure_callback": f"{__name__}.print_context_callback",
+                "on_success_callback": f"{__name__}.print_context_callback",
+                "on_execute_callback": f"{__name__}.print_context_callback",
+                "on_retry_callback": f"{__name__}.print_context_callback",
+            }
+        },
+    },
+    "tasks": {
+        "task_1": {
+            "operator": "airflow.operators.bash_operator.BashOperator",
+            "bash_command": "echo 1",
+            "task_group_name": "task_group_1",
+        },
+        "task_2": {
+            "operator": "airflow.operators.bash_operator.BashOperator",
+            "bash_command": "echo 2",
+            "task_group_name": "task_group_1",
+        },
+        "task_3": {
+            "operator": "airflow.operators.bash_operator.BashOperator",
+            "bash_command": "echo 3",
+            "task_group_name": "task_group_1",
+            "dependencies": ["task_2"],
+        },
+        "task_4": {
+            "operator": "airflow.operators.bash_operator.BashOperator",
+            "bash_command": "echo 4",
+            "dependencies": ["task_group_1"],
+        }
+    },
+}
 DAG_CONFIG_DYNAMIC_TASK_MAPPING = {
     "default_args": {"owner": "custom_owner"},
     "description": "This is an example dag with dynamic task mapping",
@@ -562,6 +600,20 @@ def test_build_task_groups():
         assert {"task_group_1.task_2", "task_group_1.task_3"} == task_group_1
         assert {"task_group_2.task_5", "task_group_2.task_6"} == task_group_2
 
+def test_build_task_groups_with_callbacks():
+    td = dagbuilder.DagBuilder("test_dag", DAG_CONFIG_TASK_GROUP_WITH_CALLBACKS, DEFAULT_CONFIG)
+    if version.parse(AIRFLOW_VERSION) < version.parse("2.2.0"):
+        error_message = "`task_groups` key can only be used with Airflow 2.x.x"
+        with pytest.raises(Exception, match=error_message):
+            td.build()
+    else:
+        actual = td.build()
+        assert actual["dag_id"] == "test_dag"
+        assert isinstance(actual["dag"], DAG)
+        assert callable(actual["dag"].task_group.get_task_group_dict()["task_group_1"].default_args["on_failure_callback"])
+        assert callable(actual["dag"].task_group.get_task_group_dict()["task_group_1"].default_args["on_execute_callback"])
+        assert callable(actual["dag"].task_group.get_task_group_dict()["task_group_1"].default_args["on_success_callback"])
+        assert callable(actual["dag"].task_group.get_task_group_dict()["task_group_1"].default_args["on_retry_callback"])
 
 @patch("dagfactory.dagbuilder.TaskGroup", new=MockTaskGroup)
 def test_make_task_groups():