This is an automated email from the ASF dual-hosted git repository.
uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new fb51c4d9bd3 Add decorator to mark task as returning dag result (#64563)
fb51c4d9bd3 is described below
commit fb51c4d9bd30146748bc5acd0ab16aacfb9f3446
Author: Tzu-ping Chung <[email protected]>
AuthorDate: Tue Jun 23 16:38:00 2026 +0800
Add decorator to mark task as returning dag result (#64563)
This does not mean result flag works for mapped tasks yet. Something
will probably be returned, but I'm not sure what exactly. We'll still
need to figure out what the exact behavior should be. This might involve
special handling during return-value-XCom-push, or when the result value
is aggregated in /wait API, or both. Not sure.
---
airflow-core/docs/public-airflow-interface.rst | 3 +-
.../unit/serialization/test_dag_serialization.py | 4 ++
task-sdk/docs/api.rst | 1 +
task-sdk/src/airflow/sdk/__init__.py | 4 +-
task-sdk/src/airflow/sdk/__init__.pyi | 7 ++-
task-sdk/src/airflow/sdk/bases/decorator.py | 7 +++
task-sdk/src/airflow/sdk/definitions/dag.py | 24 ++++++--
.../airflow/sdk/definitions/decorators/__init__.py | 26 +++++++-
.../sdk/definitions/decorators/__init__.pyi | 2 +
.../src/airflow/sdk/definitions/mappedoperator.py | 5 +-
task-sdk/src/airflow/sdk/definitions/xcom_arg.py | 4 +-
.../task_sdk/definitions/decorators/test_result.py | 69 ++++++++++++++++++++++
task-sdk/tests/task_sdk/definitions/test_dag.py | 36 +++++++++--
13 files changed, 175 insertions(+), 17 deletions(-)
diff --git a/airflow-core/docs/public-airflow-interface.rst
b/airflow-core/docs/public-airflow-interface.rst
index 4f4c09d66d1..3d9dccf577c 100644
--- a/airflow-core/docs/public-airflow-interface.rst
+++ b/airflow-core/docs/public-airflow-interface.rst
@@ -116,10 +116,11 @@ API Server, etc.), providing a version-agnostic, stable
interface for writing an
* :func:`airflow.sdk.asset`
* :func:`airflow.sdk.dag`
-* :func:`airflow.sdk.setup`
* :func:`airflow.sdk.task`
* :func:`airflow.sdk.task_group`
+* :func:`airflow.sdk.setup`
* :func:`airflow.sdk.teardown`
+* :func:`airflow.sdk.result`
* :func:`airflow.sdk.chain`
* :func:`airflow.sdk.chain_linear`
* :func:`airflow.sdk.cross_downstream`
diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py
b/airflow-core/tests/unit/serialization/test_dag_serialization.py
index 5bc5b0ae847..61a66c60fb6 100644
--- a/airflow-core/tests/unit/serialization/test_dag_serialization.py
+++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py
@@ -866,12 +866,16 @@ class TestStringifiedDAGs:
"_is_sensor",
# trigger_kwargs is kept as raw JSON after deserialization;
checked separately
"start_trigger_args",
+ # Only needed at execution time; intentionally excluded
+ "returns_dag_result",
}
else: # Promised to be mapped by the assert above.
assert isinstance(serialized_task, SerializedMappedOperator)
fields_to_check = {f.name for f in attrs.fields(MappedOperator)}
fields_to_check -= {
+ # Only needed at execution time; intentionally excluded
"map_index_template",
+ "returns_dag_result",
# Matching logic in BaseOperator.get_serialized_fields().
"dag",
"task_group",
diff --git a/task-sdk/docs/api.rst b/task-sdk/docs/api.rst
index d289628097e..d31695f486e 100644
--- a/task-sdk/docs/api.rst
+++ b/task-sdk/docs/api.rst
@@ -72,6 +72,7 @@ Task Decorators:
.. autoapifunction:: airflow.sdk.teardown
.. autofunction:: airflow.sdk.task
+.. autofunction:: airflow.sdk.result
.. autofunction:: airflow.sdk.setup
.. autofunction:: airflow.sdk.teardown
.. autofunction:: airflow.sdk.asset
diff --git a/task-sdk/src/airflow/sdk/__init__.py
b/task-sdk/src/airflow/sdk/__init__.py
index 9f7fd1b0484..8e8953a8aca 100644
--- a/task-sdk/src/airflow/sdk/__init__.py
+++ b/task-sdk/src/airflow/sdk/__init__.py
@@ -110,6 +110,7 @@ __all__ = [
"literal",
"lineage",
"macros",
+ "result",
"setup",
"task",
"task_group",
@@ -151,7 +152,7 @@ if TYPE_CHECKING:
from airflow.sdk.definitions.context import Context, get_current_context,
get_parsing_context
from airflow.sdk.definitions.dag import DAG, dag
from airflow.sdk.definitions.deadline import DeadlineAlert,
DeadlineReference
- from airflow.sdk.definitions.decorators import setup, task, teardown
+ from airflow.sdk.definitions.decorators import result, setup, task,
teardown
from airflow.sdk.definitions.decorators.task_group import task_group
from airflow.sdk.definitions.edges import EdgeModifier, Label
from airflow.sdk.definitions.param import Param, ParamsDict
@@ -313,6 +314,7 @@ __lazy_imports: dict[str, str] = {
"literal": ".definitions.template",
"lineage": ".lineage",
"macros": ".execution_time",
+ "result": ".definitions.decorators",
"setup": ".definitions.decorators",
"task": ".definitions.decorators",
"task_group": ".definitions.decorators",
diff --git a/task-sdk/src/airflow/sdk/__init__.pyi
b/task-sdk/src/airflow/sdk/__init__.pyi
index 78e37465379..273b0579f3e 100644
--- a/task-sdk/src/airflow/sdk/__init__.pyi
+++ b/task-sdk/src/airflow/sdk/__init__.pyi
@@ -60,7 +60,12 @@ from airflow.sdk.definitions.context import (
get_parsing_context as get_parsing_context,
)
from airflow.sdk.definitions.dag import DAG as DAG, dag as dag
-from airflow.sdk.definitions.decorators import setup as setup, task as task,
teardown as teardown
+from airflow.sdk.definitions.decorators import (
+ result as result,
+ setup as setup,
+ task as task,
+ teardown as teardown,
+)
from airflow.sdk.definitions.decorators.task_group import task_group as
task_group
from airflow.sdk.definitions.edges import EdgeModifier as EdgeModifier, Label
as Label
from airflow.sdk.definitions.param import Param as Param
diff --git a/task-sdk/src/airflow/sdk/bases/decorator.py
b/task-sdk/src/airflow/sdk/bases/decorator.py
index 8634fbe9964..e45778725cc 100644
--- a/task-sdk/src/airflow/sdk/bases/decorator.py
+++ b/task-sdk/src/airflow/sdk/bases/decorator.py
@@ -468,6 +468,7 @@ class _TaskDecorator(ExpandableFactory, Generic[FParams,
FReturn, OperatorSubcla
_airflow_is_task_decorator: ClassVar[bool] = True
is_setup: bool = False
is_teardown: bool = False
+ returns_dag_result: bool = False
on_failure_fail_dagrun: bool = False
# This is set in __attrs_post_init__ by update_wrapper. Provided here for
type hints.
@@ -524,6 +525,7 @@ class _TaskDecorator(ExpandableFactory, Generic[FParams,
FReturn, OperatorSubcla
op.is_setup = self.is_setup
op.is_teardown = self.is_teardown
op.on_failure_fail_dagrun = on_failure_fail_dagrun
+ op.returns_dag_result = self.returns_dag_result
op_doc_attrs = [op.doc, op.doc_json, op.doc_md, op.doc_rst,
op.doc_yaml]
# Set the task's doc_md to the function's docstring if it exists and
no other doc* args are set.
if self.function.__doc__ and not any(op_doc_attrs):
@@ -685,6 +687,7 @@ class _TaskDecorator(ExpandableFactory, Generic[FParams,
FReturn, OperatorSubcla
expand_input_attr="op_kwargs_expand_input",
start_trigger_args=self.operator_class.start_trigger_args,
start_from_trigger=self.operator_class.start_from_trigger,
+ returns_dag_result=self.returns_dag_result,
)
return XComArg(operator=operator)
@@ -834,3 +837,7 @@ def task_decorator_factory(
)
return cast("TaskDecorator", decorator_factory)
+
+
+def is_decorated_task(f: Callable) -> typing_extensions.TypeIs[_TaskDecorator]:
+ return getattr(f, "_airflow_is_task_decorator", False)
diff --git a/task-sdk/src/airflow/sdk/definitions/dag.py
b/task-sdk/src/airflow/sdk/definitions/dag.py
index fd47de467fe..46794edbe4a 100644
--- a/task-sdk/src/airflow/sdk/definitions/dag.py
+++ b/task-sdk/src/airflow/sdk/definitions/dag.py
@@ -74,6 +74,7 @@ if TYPE_CHECKING:
from airflow.sdk.definitions.edges import EdgeInfoType
from airflow.sdk.definitions.mappedoperator import MappedOperator
from airflow.sdk.definitions.taskgroup import TaskGroup
+ from airflow.sdk.definitions.xcom_arg import PlainXComArg
from airflow.sdk.execution_time.supervisor import TaskRunResult
from airflow.timetables.base import DataInterval, Timetable as
CoreTimetable
@@ -303,6 +304,13 @@ def _default_task_group(instance: DAG) -> TaskGroup:
return TaskGroup.create_root(dag=instance)
+def _is_valid_dag_result(value: Any) -> TypeIs[PlainXComArg]:
+ from airflow.sdk.bases.xcom import BaseXCom
+ from airflow.sdk.definitions.xcom_arg import PlainXComArg
+
+ return isinstance(value, PlainXComArg) and value.key ==
BaseXCom.XCOM_RETURN_KEY
+
+
# TODO: Task-SDK: look at re-enabling slots after we remove pickling
@attrs.define(repr=False, field_transformer=_all_after_dag_id_to_kw_only,
slots=False)
class DAG:
@@ -1116,12 +1124,8 @@ class DAG:
tg._remove(task)
def add_result(self, xcom_arg: X) -> X:
- from airflow.sdk.bases.xcom import BaseXCom
- from airflow.sdk.definitions.xcom_arg import PlainXComArg
-
- if not isinstance(xcom_arg, PlainXComArg) or xcom_arg.key !=
BaseXCom.XCOM_RETURN_KEY:
+ if not _is_valid_dag_result(xcom_arg):
raise ValueError("Only plain return value can be used as dag
result")
-
xcom_arg.operator.returns_dag_result = True
return xcom_arg
@@ -1674,7 +1678,15 @@ def dag(dag_id_or_func=None, __DAG_class=DAG,
__warnings_stacklevel_delta=2, **d
dag_obj.fileloc = back.f_code.co_filename if back else ""
# Invoke function to create operators in the Dag scope.
- f(**f_kwargs)
+ r = f(**f_kwargs)
+
+ if _is_valid_dag_result(r):
+ log.debug(
+ "Automatically adding function return value %r as
result for dag %s",
+ r,
+ dag_obj.dag_id,
+ )
+ dag_obj.add_result(r)
# Return dag object such that it's accessible in Globals.
return dag_obj
diff --git a/task-sdk/src/airflow/sdk/definitions/decorators/__init__.py
b/task-sdk/src/airflow/sdk/definitions/decorators/__init__.py
index d3d99951609..44c4bd717ca 100644
--- a/task-sdk/src/airflow/sdk/definitions/decorators/__init__.py
+++ b/task-sdk/src/airflow/sdk/definitions/decorators/__init__.py
@@ -18,7 +18,7 @@ from __future__ import annotations
from typing import TYPE_CHECKING
-from airflow.sdk.bases.decorator import TaskDecorator
+from airflow.sdk.bases.decorator import TaskDecorator, is_decorated_task
from airflow.sdk.definitions.dag import dag
from airflow.sdk.definitions.decorators.condition import run_if, skip_if
from airflow.sdk.definitions.decorators.setup_teardown import setup_task,
teardown_task
@@ -27,12 +27,16 @@ from airflow.sdk.providers_manager_runtime import
ProvidersManagerTaskRuntime
if TYPE_CHECKING:
from collections.abc import Callable
+ from typing import TypeVar
+
+ C = TypeVar("C", bound=Callable)
# Please keep this in sync with the .pyi's __all__.
__all__ = [
"TaskDecorator",
"TaskDecoratorCollection",
"dag",
+ "result",
"task",
"task_group",
"setup",
@@ -63,3 +67,23 @@ class TaskDecoratorCollection:
task = TaskDecoratorCollection()
setup: Callable = setup_task
teardown: Callable = teardown_task
+
+
+def result(t: C) -> C:
+ """
+ Mark a task as returning the dag's result.
+
+ This must be used *on top of* a ``@task`` decorator like this::
+
+ @result
+ @task
+ def emit_values():
+ something = ...
+ return something
+
+ This causes :func:`DAG.add_result` to later be called internally.
+ """
+ if not is_decorated_task(t):
+ raise TypeError("@result must be used on top of a @task-decorated
function")
+ t.returns_dag_result = True
+ return t
diff --git a/task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi
b/task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi
index 8e84b33a421..7b7c354a9a6 100644
--- a/task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi
+++ b/task-sdk/src/airflow/sdk/definitions/decorators/__init__.pyi
@@ -38,6 +38,7 @@ __all__ = [
"TaskDecorator",
"TaskDecoratorCollection",
"dag",
+ "result",
"task",
"task_group",
"setup",
@@ -948,3 +949,4 @@ class TaskDecoratorCollection:
task: TaskDecoratorCollection
setup: Callable
teardown: Callable
+result: Callable
diff --git a/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
b/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
index 0faa2ab6f18..4c030be4beb 100644
--- a/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
+++ b/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
@@ -315,6 +315,7 @@ class MappedOperator(AbstractOperator):
start_trigger_args: StartTriggerArgs | None
start_from_trigger: bool
_needs_expansion: bool = True
+ returns_dag_result: bool = False
dag: DAG | None
task_group: TaskGroup | None
@@ -361,7 +362,7 @@ class MappedOperator(AbstractOperator):
@classmethod
def get_serialized_fields(cls):
# Not using 'cls' here since we only want to serialize base fields.
- return (frozenset(attrs.fields_dict(MappedOperator))) - {
+ return frozenset(attrs.fields_dict(MappedOperator)) - {
"_is_empty",
"_can_skip_downstream",
"dag",
@@ -376,6 +377,7 @@ class MappedOperator(AbstractOperator):
"_needs_expansion",
"partial_kwargs",
"operator_extra_links",
+ "returns_dag_result",
}
@property
@@ -792,6 +794,7 @@ class MappedOperator(AbstractOperator):
op.is_setup = is_setup
op.is_teardown = is_teardown
op.on_failure_fail_dagrun = on_failure_fail_dagrun
+ op.returns_dag_result = self.returns_dag_result
op.downstream_task_ids = self.downstream_task_ids
op.upstream_task_ids = self.upstream_task_ids
return op
diff --git a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
index 88db9009242..59549083259 100644
--- a/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
+++ b/task-sdk/src/airflow/sdk/definitions/xcom_arg.py
@@ -420,8 +420,10 @@ class MapXComArg(XComArg):
callables: MapCallables
def __attrs_post_init__(self) -> None:
+ from airflow.sdk.bases.decorator import is_decorated_task
+
for c in self.callables:
- if getattr(c, "_airflow_is_task_decorator", False):
+ if is_decorated_task(c):
raise ValueError("map() argument must be a plain function, not
a @task operator")
def __repr__(self) -> str:
diff --git a/task-sdk/tests/task_sdk/definitions/decorators/test_result.py
b/task-sdk/tests/task_sdk/definitions/decorators/test_result.py
new file mode 100644
index 00000000000..b68d1f98621
--- /dev/null
+++ b/task-sdk/tests/task_sdk/definitions/decorators/test_result.py
@@ -0,0 +1,69 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import re
+
+import pytest
+
+from airflow.sdk import DAG, result, task
+
+
+def test_result_error_if_not_task():
+ with pytest.raises(
+ TypeError,
+ match=re.escape("@result must be used on top of a @task-decorated
function"),
+ ):
+
+ @result
+ def f():
+ pass
+
+
+def test_result_marks_returns_dag_result():
+ @result
+ @task
+ def foo():
+ pass
+
+ t = foo()
+ assert t.operator.returns_dag_result is True
+
+
+def test_retain_returns_dag_result_when_other_attrs_are_overridden():
+ @result
+ @task
+ def foo():
+ pass
+
+ with DAG("test_retain_returns_dag_result_when_other_attrs_are_overridden")
as dag:
+ foo.override(task_id="foo2")()
+ assert dag.get_task("foo2").returns_dag_result is True
+
+
+def test_retain_returns_dag_result_when_task_is_expanded():
+ @result
+ @task
+ def foo(x):
+ return x
+
+ with DAG("test_retain_returns_dag_result_when_task_is_expanded") as dag:
+ foo.expand(x=[1, 2, 3])
+
+ assert dag.get_task("foo").returns_dag_result is True
diff --git a/task-sdk/tests/task_sdk/definitions/test_dag.py
b/task-sdk/tests/task_sdk/definitions/test_dag.py
index c42eb8dfc80..0e9f6c78b38 100644
--- a/task-sdk/tests/task_sdk/definitions/test_dag.py
+++ b/task-sdk/tests/task_sdk/definitions/test_dag.py
@@ -21,15 +21,14 @@ import warnings
import weakref
from datetime import datetime, timedelta, timezone
from typing import Any
+from unittest import mock
import pytest
-from airflow.sdk import Context, Label, PartitionAtRuntime, TaskGroup
+from airflow.sdk import DAG, Context, Label, Param, PartitionAtRuntime,
TaskGroup, dag as dag_decorator, task
from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.bases.timetable import BaseTimetable
-from airflow.sdk.definitions.dag import DAG, dag as dag_decorator
-from airflow.sdk.definitions.param import DagParam, Param, ParamsDict
-from airflow.sdk.definitions.timetables import assets, events, interval,
simple, trigger # noqa: F401
+from airflow.sdk.definitions.param import DagParam, ParamsDict
from airflow.sdk.exceptions import AirflowDagCycleException,
DuplicateTaskIdFound, RemovedInAirflow4Warning
from airflow.utils.types import DagRunType
@@ -788,7 +787,6 @@ class TestDagDecorator:
def test_dag_param_resolves(self):
"""Test that dag param is correctly resolved by operator"""
- from airflow.decorators import task
@dag_decorator(schedule=None, default_args=self.DEFAULT_ARGS)
def xcom_pass_to_op(value=self.VALUE):
@@ -805,6 +803,34 @@ class TestDagDecorator:
self.operator.render_template_fields({})
assert self.operator.op_args[0] == 42
+ def test_ignore_function_result(self, monkeypatch):
+ monkeypatch.setattr(DAG, "add_result",
mock.create_autospec(DAG.add_result))
+
+ @dag_decorator
+ def d():
+ @task
+ def return_num(num):
+ return num
+
+ return_num(123)
+ return 123
+
+ dag = d()
+ assert dag.get_task("return_num").returns_dag_result is False
+ assert DAG.add_result.mock_calls == []
+
+ def test_function_result_set_to_xcom_arg(self):
+ @dag_decorator
+ def d():
+ @task
+ def return_num(num):
+ return num
+
+ return return_num(123)
+
+ dag = d()
+ assert dag.get_task("return_num").returns_dag_result is True
+
class DoNothingOperator(BaseOperator):
"""
