Skip to content

Conversation

@yanshil
Copy link

@yanshil yanshil commented Sep 30, 2025

Airflow supports use of {decorated_task}[{xcom_key}] for xcoms that are not return_value link

Here I support for fetching upstream task's xcom from decorated task with multiple_outputs=True

The following example has already been tested in Airflow 3.1.0 (and should also works in Airflow 3)

example_multiple_outputs_taskflow:
  schedule: null
  render_template_as_native_obj: True
  catchup: false
  tags: [example, test]
  tasks:
    - task_id: collect
      multiple_outputs: true
      decorator: airflow.decorators.task
      python_callable: sample.collect
    - task_id: echo1
      decorator: airflow.decorators.task
      python_callable: sample.echo
      value: "+collect['key1']"
    - task_id: echo2
      decorator: airflow.decorators.task
      python_callable: sample.echo
      value: "+collect['key2']"

@yanshil yanshil requested a review from a team as a code owner September 30, 2025 11:16
@yanshil yanshil changed the title feat: support key: {task}[{xcom_key}] use feat: support key: +{task}[{xcom_key}] use Sep 30, 2025
@codecov-commenter
Copy link

Codecov Report

❌ Patch coverage is 87.50000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.65%. Comparing base (294bc47) to head (6f8177a).
⚠️ Report is 1 commits behind head on main.

Files with missing lines Patch % Lines
dagfactory/dagbuilder.py 87.50% 2 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #593      +/-   ##
==========================================
- Coverage   93.77%   93.65%   -0.13%     
==========================================
  Files          13       13              
  Lines        1125     1135      +10     
==========================================
+ Hits         1055     1063       +8     
- Misses         70       72       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copy link
Collaborator

@pankajastro pankajastro left a comment

Choose a reason for hiding this comment

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

Could you please add an example DAG here and document this feature?

@pankajastro pankajastro requested a review from Copilot October 3, 2025 12:21
Copy link

Copilot AI left a comment

Choose a reason for hiding this comment

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

Pull Request Overview

This PR adds support for accessing XCom values from decorated tasks with multiple_outputs=True using the syntax +{task}[{xcom_key}]. This extends the existing +{task} functionality to allow fetching specific keys from tasks that return multiple outputs.

  • Refactored XCom value replacement logic into a dedicated helper method
  • Added regex pattern matching to support the new +{task}[{xcom_key}] syntax
  • Updated existing functionality to use the new consolidated approach

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File Description
dev/dags/sample.py Added sample functions to demonstrate the new XCom key syntax
dagfactory/dagbuilder.py Implemented regex-based parsing for XCom references and consolidated replacement logic

Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.

return decorator(**decorator_kwargs)(**callable_kwargs)

@staticmethod
def _replace_kwargs_values_as_xcom(kwargs: dict(str, Any), key: str, value: Any, tasks_dict: dict(str, Any)):
Copy link

Copilot AI Oct 3, 2025

Choose a reason for hiding this comment

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

The method signature uses dict(str, Any) which is incorrect syntax. Should be Dict[str, Any] and requires importing Dict from typing.

Copilot uses AI. Check for mistakes.
@staticmethod
def _replace_kwargs_values_as_xcom(kwargs: dict(str, Any), key: str, value: Any, tasks_dict: dict(str, Any)):
# Match with multiple_outputs=True case with {key}: +{task}["{xcom_key}"]
_PATTERN = re.compile(r"^\+(?P<task>\w*)(\[['\"](?P<xcom_key>.*)['\"]\])?$")
Copy link

Copilot AI Oct 3, 2025

Choose a reason for hiding this comment

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

The regex pattern allows empty task names with \w*. This should be \w+ to require at least one word character for the task name, preventing matches on invalid references like +[] or +['key'].

Suggested change
_PATTERN = re.compile(r"^\+(?P<task>\w*)(\[['\"](?P<xcom_key>.*)['\"]\])?$")
_PATTERN = re.compile(r"^\+(?P<task>\w+)(\[['\"](?P<xcom_key>.*)['\"]\])?$")

Copilot uses AI. Check for mistakes.
@staticmethod
def _replace_kwargs_values_as_xcom(kwargs: dict(str, Any), key: str, value: Any, tasks_dict: dict(str, Any)):
# Match with multiple_outputs=True case with {key}: +{task}["{xcom_key}"]
_PATTERN = re.compile(r"^\+(?P<task>\w*)(\[['\"](?P<xcom_key>.*)['\"]\])?$")
Copy link

Copilot AI Oct 3, 2025

Choose a reason for hiding this comment

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

The regex pattern (?P<xcom_key>.*) allows any characters including newlines and control characters in the xcom_key. Consider using [^'\"]* or similar to restrict allowed characters and prevent potential injection issues.

Suggested change
_PATTERN = re.compile(r"^\+(?P<task>\w*)(\[['\"](?P<xcom_key>.*)['\"]\])?$")
_PATTERN = re.compile(r"^\+(?P<task>\w*)(\[['\"](?P<xcom_key>[^'\"]*)['\"]\])?$")

Copilot uses AI. Check for mistakes.
@yanshil
Copy link
Author

yanshil commented Oct 9, 2025

I also add a simple feature docs. You can reorder it if you want.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants