Use + as prefix for taskflow references

astronomer · Dec 5, 2024 · 2a2f978 · 2a2f978
1 parent 0b07eed
commit 2a2f978
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/dagfactory/dagbuilder.py b/dagfactory/dagbuilder.py
@@ -987,15 +987,15 @@ def make_decorator(
                 else operator_obj.partial(**task_params).expand(**expand_kwargs)
             )
         """
-
         callable_args_keys = inspect.getfullargspec(python_callable).args
         callable_kwargs = {}
         decorator_kwargs = dict(**task_params)
         for arg_key, arg_value in task_params.items():
             if arg_key in callable_args_keys:
                 decorator_kwargs.pop(arg_key)
-                if arg_value in tasks_dict:
-                    callable_kwargs[arg_key] = tasks_dict[arg_value]
+                if arg_value.startswith("+"):
+                    upstream_task_name = arg_value.split("+")[0]
+                    callable_kwargs[arg_key] = tasks_dict[upstream_task_name]
                 else:
                     callable_kwargs[arg_key] = arg_value
 

diff --git a/dev/dags/example_taskflow.yml b/dev/dags/example_taskflow.yml
@@ -8,19 +8,19 @@ example_taskflow:
   tasks:
     some_number:
       decorator: airflow.decorators.task
-      python_callable: "sample.some_number"
+      python_callable: sample.some_number
     numbers_list:
       decorator: airflow.decorators.task
       python_callable_name: build_numbers_list
       python_callable_file: $CONFIG_ROOT_DIR/sample.py
     double_number_from_arg:
       decorator: airflow.decorators.task
-      python_callable: "sample.double"
+      python_callable: sample.double
       number: 2
     double_number_from_task:
       decorator: airflow.decorators.task
-      python_callable: "sample.double"
-      number: some_number  # this is a task previously defined
+      python_callable: sample.double
+      number: +some_number  # this is a task previously defined
       #dependencies: [some_number]
     #double_list_dynamically:
     #  python_callable: !!python/name:sample.double"