Skip to content

Add decorator to mark task as returning dag result#64563

Merged
uranusjr merged 3 commits into
apache:mainfrom
astronomer:result-task-integration
Jun 23, 2026
Merged

Add decorator to mark task as returning dag result#64563
uranusjr merged 3 commits into
apache:mainfrom
astronomer:result-task-integration

Conversation

@uranusjr

@uranusjr uranusjr commented Apr 1, 2026

Copy link
Copy Markdown
Member

Continuing #51711 (comment)

This adds

  • A @result decorator that can be used on top of any taskflow task instead of an explicit add_result call (useful when you don’t have a reference to the dag object)
  • If a task’s result XComArg is returned by a @dag-decorated function, add_result is called automatically on the created dag object.

Comment thread task-sdk/src/airflow/sdk/bases/decorator.py
Comment thread task-sdk/tests/task_sdk/definitions/test_dag.py Outdated
Comment thread task-sdk/tests/task_sdk/definitions/test_dag.py Outdated

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds Task SDK ergonomics for marking a single TaskFlow task as the “DAG result”, supporting synchronous/inference-style execution patterns by allowing authors to either explicitly mark a task via a decorator or implicitly infer the result when returning an XComArg from a @dag function.

Changes:

  • Introduces @result decorator to mark a TaskFlow task as returning the DAG result (sets returns_dag_result).
  • Updates @dag decorator factory to automatically call dag.add_result() when the decorated function returns a plain return-value XComArg.
  • Adds/updates exports, docs, and tests for the new behavior; small refactor to centralize DAG-result validation.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 4 comments.

Show a summary per file
File Description
task-sdk/tests/task_sdk/definitions/test_dag.py Adds tests for ignoring non-XCom returns and inferring DAG result from returned XComArg.
task-sdk/tests/task_sdk/definitions/decorators/test_result.py New tests validating @result usage and behavior.
task-sdk/src/airflow/sdk/definitions/xcom_arg.py Uses shared is_decorated_task() helper to validate map() callables.
task-sdk/src/airflow/sdk/definitions/decorators/init.pyi Exposes result in typing stubs for decorators module.
task-sdk/src/airflow/sdk/definitions/decorators/init.py Implements result() decorator and exports it.
task-sdk/src/airflow/sdk/definitions/dag.py Adds _is_valid_dag_result() helper; auto-adds DAG result when @dag function returns a valid XComArg.
task-sdk/src/airflow/sdk/bases/decorator.py Adds returns_dag_result flag propagation from decorator to created operator; adds is_decorated_task().
task-sdk/src/airflow/sdk/init.pyi Exposes result from airflow.sdk in typing stubs.
task-sdk/src/airflow/sdk/init.py Adds result to public exports and lazy-import map.
airflow-core/docs/public-airflow-interface.rst Documents airflow.sdk.result as part of the public interface.

Comment thread task-sdk/src/airflow/sdk/definitions/decorators/__init__.py Outdated
Comment thread task-sdk/src/airflow/sdk/definitions/decorators/__init__.py Outdated
Comment thread task-sdk/src/airflow/sdk/definitions/dag.py Outdated
Comment thread task-sdk/tests/task_sdk/definitions/test_dag.py Outdated

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.

Comment thread task-sdk/src/airflow/sdk/definitions/decorators/__init__.py Outdated
@uranusjr uranusjr force-pushed the result-task-integration branch from ed11b50 to 1de7834 Compare April 8, 2026 06:25
@kaxil kaxil requested a review from Copilot April 10, 2026 19:55

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 10 out of 10 changed files in this pull request and generated 1 comment.

Comment thread task-sdk/src/airflow/sdk/bases/decorator.py
@uranusjr uranusjr force-pushed the result-task-integration branch from 1de7834 to d35ec16 Compare June 14, 2026 01:48
This ensures the flag is forwarded through .expand() and later .unmap()
to the dynamically-created BaseOperator used on the runtime TI.

This still does not mean result flag works for mapped tasks. 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.
@uranusjr uranusjr force-pushed the result-task-integration branch from d35ec16 to a8e9626 Compare June 14, 2026 02:17
@uranusjr uranusjr requested a review from kaxil June 15, 2026 01:19
@uranusjr uranusjr added this to the Airflow 3.3.0 milestone Jun 15, 2026
@uranusjr uranusjr linked an issue Jun 15, 2026 that may be closed by this pull request
Comment thread task-sdk/tests/task_sdk/definitions/decorators/test_result.py
@uranusjr uranusjr merged commit fb51c4d into apache:main Jun 23, 2026
102 checks passed
@uranusjr uranusjr deleted the result-task-integration branch June 23, 2026 08:38
cetingokhan pushed a commit to cetingokhan/airflow that referenced this pull request Jun 24, 2026
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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Enable Inference Execution / Synchronous DAG Execution

3 participants