You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by el...@apache.org on 2025/08/22 08:24:01 UTC
(airflow) branch main updated: Allow rendering of template fields with start from trigger (#53071)
This is an automated email from the ASF dual-hosted git repository.
eladkal 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 4362c25e5cb Allow rendering of template fields with start from trigger (#53071)
4362c25e5cb is described below
commit 4362c25e5cb96e3b6c99497b682d4dfe23a86aba
Author: David Blain <in...@dabla.be>
AuthorDate: Fri Aug 22 10:23:55 2025 +0200
Allow rendering of template fields with start from trigger (#53071)
* refactor: Render templated fields before starting from trigger
* refactor: Moved StartTriggerArgs from airflow-core to task-sdk so it can be part of the AbstractOperator and moved rendering logic from expand_start_trigger_args to _do_render_template_fields so that all the rendering logic is common in on method
* refactor: Import missing datetime
* refactor: Fixed import of datetime
* refactor: Reformatted files
* refactor: Changed import of StartTriggerArgs from new location in airflow-core
* refactor: Moved call of render_template_fields from DagRun to defer_task method of TaskInstance
* refactor: Moved call of render_template_fields from TaskInstance to expand_start_trigger_args method of BaseOperator
* refactor: Do the template rendering in update_triggers of TriggerRunnerSupervisor
* refactor: Fixed imports tests
* refactor: Fixed imports
* refactor: Reformatted update_triggers
* refactor: Reformatted expand_start_trigger_args
* refactor: Revert move of StartTriggerArgs to airflow-sdk
* Revert "refactor: Revert move of StartTriggerArgs to airflow-sdk"
This reverts commit 5918847dd5c8ee3701d3da2045bd45cc89d8138a.
* refactor: Defined StartTriggerArgs as deprecated classes in old import location
* refactor: Moved definition of StartTriggerArgs as deprecated classes in old import location to base instead of init
* refactor: Moved import of StartTriggerArgs in type checking block
* refactor: Added StartTriggerArgs in __all__
* Revert "refactor: Added StartTriggerArgs in __all__"
This reverts commit 2e2721bd6e2988fc64e3edb8133233b8722930c8.
* Revert "refactor: Moved definition of StartTriggerArgs as deprecated classes in old import location to base instead of init"
This reverts commit 50a8a918cf98acd63bd7484bdb8a0b7f31421903.
* refactor: Import StartTriggerArgs in type checking block
* refactor: Moved definition of StartTriggerArgs as deprecated classes in old import location to base instead of init
* refactor: Fixed inner import of StartTriggerArgs in expand_start_trigger_args
* refactor: Use explicit re-export for StartTriggerArgs
* Revert "refactor: Moved definition of StartTriggerArgs as deprecated classes in old import location to base instead of init"
This reverts commit 8a4e0c64d2ed11dadc4cb790da24f2e7c4b6cc22.
* Revert "refactor: Use explicit re-export for StartTriggerArgs"
This reverts commit b14cb93362c2d15007a1c5ca898d970472ebb326.
* Revert "refactor: Import StartTriggerArgs in type checking block"
This reverts commit c1533ac261e0e94a16519435fde96736559636b6.
* refactor: Removed StartTriggerArgs from deprecated_classes and override getattr in base module instead
* refactor: Updated warning message
* refactor: Reformatted warning message
* refactor: Reformatted base
* refactor: Reformatted getattr to not do inner import to avoid static check error
* refactor: Try mocking DagBag for test_trigger_lifecycle
* refactor: Reorganized imports
* refactor: Try to fix test_trigger_create_race_condition_38599
* refactor: Try to fix has not been assigned to a DAG yet error in test_trigger_create_race_condition_38599
* refactor: Try fixing DagRun validation errors for start_date and run_type in test_trigger_create_race_condition_38599
* refactor: expand_start_trigger_args method doesn't have session parameter anymore
* refactor: Try fixing remaining tests
* refactor: Reformatted test trigger job
* refactor: Check if start_trigger_args is not None
* refactor: Removed old imports
* refactor: Fixed import issues
* refactor: Reformatted test_trigger_create_race_condition_38599
* refactor: Try to fix test_update_triggers_prevents_duplicate_creation_queue_entries
* refactor: Reformatted test_update_triggers_prevents_duplicate_creation_queue_entries
* refactor: Added StartTriggerArgs to api.rst
* refactor: Added StartTriggerArgs in __init__.pyi
* refactor: Added lazy import for StartTriggerArgs
* refactor: Changed order in __all__ for StartTriggerArgs
* refactor: Reorganized imports for MappedOperator
---------
Co-authored-by: David Blain <da...@infrabel.be>
Co-authored-by: Tzu-ping Chung <ur...@gmail.com>
---
.../src/airflow/jobs/triggerer_job_runner.py | 61 ++++++++++++-------
.../airflow/serialization/serialized_objects.py | 3 +-
airflow-core/src/airflow/triggers/base.py | 23 ++++----
airflow-core/tests/unit/jobs/test_triggerer_job.py | 68 ++++++++++++++++++----
airflow-core/tests/unit/models/test_dagrun.py | 2 +-
.../unit/serialization/test_dag_serialization.py | 2 +-
task-sdk/docs/api.rst | 2 +
task-sdk/src/airflow/sdk/__init__.py | 3 +
task-sdk/src/airflow/sdk/__init__.pyi | 2 +
task-sdk/src/airflow/sdk/bases/operator.py | 6 +-
task-sdk/src/airflow/sdk/bases/trigger.py | 33 +++++++++++
.../sdk/definitions/_internal/abstractoperator.py | 12 ++++
.../src/airflow/sdk/definitions/mappedoperator.py | 14 ++++-
13 files changed, 180 insertions(+), 51 deletions(-)
diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
index dac78debe00..43e9cee40ab 100644
--- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
+++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
@@ -43,6 +43,7 @@ from airflow.configuration import conf
from airflow.executors import workloads
from airflow.jobs.base_job_runner import BaseJobRunner
from airflow.jobs.job import perform_heartbeat
+from airflow.models import DagBag
from airflow.models.trigger import Trigger
from airflow.sdk.api.datamodels._generated import HITLDetailResponse
from airflow.sdk.execution_time.comms import (
@@ -604,6 +605,45 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
trigger set.
"""
render_log_fname = log_filename_template_renderer()
+ dag_bag = DagBag(collect_dags=False)
+
+ def expand_start_trigger_args(trigger: Trigger) -> Trigger:
+ task = dag_bag.get_dag(trigger.task_instance.dag_id).get_task(trigger.task_instance.task_id)
+ if task.template_fields:
+ trigger.task_instance.refresh_from_task(task)
+ context = trigger.task_instance.get_template_context()
+ task.render_template_fields(context=context)
+ start_trigger_args = task.expand_start_trigger_args(context=context)
+ if start_trigger_args:
+ trigger.kwargs = start_trigger_args.trigger_kwargs
+ return trigger
+
+ def create_workload(trigger: Trigger) -> workloads.RunTrigger:
+ if trigger.task_instance:
+ log_path = render_log_fname(ti=trigger.task_instance)
+
+ trigger = expand_start_trigger_args(trigger)
+
+ ser_ti = workloads.TaskInstance.model_validate(trigger.task_instance, from_attributes=True)
+ # When producing logs from TIs, include the job id producing the logs to disambiguate it.
+ self.logger_cache[new_id] = TriggerLoggingFactory(
+ log_path=f"{log_path}.trigger.{self.job.id}.log",
+ ti=ser_ti, # type: ignore
+ )
+
+ return workloads.RunTrigger(
+ classpath=trigger.classpath,
+ id=new_id,
+ encrypted_kwargs=trigger.encrypted_kwargs,
+ ti=ser_ti,
+ timeout_after=trigger.task_instance.trigger_timeout,
+ )
+ return workloads.RunTrigger(
+ classpath=trigger.classpath,
+ id=new_id,
+ encrypted_kwargs=trigger.encrypted_kwargs,
+ ti=None,
+ )
known_trigger_ids = (
self.running_triggers.union(x[0] for x in self.events)
@@ -641,26 +681,7 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
)
continue
- workload = workloads.RunTrigger(
- classpath=new_trigger_orm.classpath,
- id=new_id,
- encrypted_kwargs=new_trigger_orm.encrypted_kwargs,
- ti=None,
- )
- if new_trigger_orm.task_instance:
- log_path = render_log_fname(ti=new_trigger_orm.task_instance)
-
- ser_ti = workloads.TaskInstance.model_validate(
- new_trigger_orm.task_instance, from_attributes=True
- )
- # When producing logs from TIs, include the job id producing the logs to disambiguate it.
- self.logger_cache[new_id] = TriggerLoggingFactory(
- log_path=f"{log_path}.trigger.{self.job.id}.log",
- ti=ser_ti, # type: ignore
- )
-
- workload.ti = ser_ti
- workload.timeout_after = new_trigger_orm.task_instance.trigger_timeout
+ workload = create_workload(new_trigger_orm)
to_create.append(workload)
diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py
index e370526e6da..6a8476ba7a3 100644
--- a/airflow-core/src/airflow/serialization/serialized_objects.py
+++ b/airflow-core/src/airflow/serialization/serialized_objects.py
@@ -52,6 +52,7 @@ from airflow.models.xcom import XComModel
from airflow.models.xcom_arg import SchedulerXComArg, deserialize_xcom_arg
from airflow.sdk import Asset, AssetAlias, AssetAll, AssetAny, AssetWatcher, BaseOperator, XComArg
from airflow.sdk.bases.operator import OPERATOR_DEFAULTS # TODO: Copy this into the scheduler?
+from airflow.sdk.bases.trigger import StartTriggerArgs
from airflow.sdk.definitions._internal.expandinput import EXPAND_INPUT_EMPTY
from airflow.sdk.definitions._internal.node import DAGNode
from airflow.sdk.definitions.asset import (
@@ -83,7 +84,7 @@ from airflow.ti_deps.deps.not_in_retry_period_dep import NotInRetryPeriodDep
from airflow.ti_deps.deps.not_previously_skipped_dep import NotPreviouslySkippedDep
from airflow.ti_deps.deps.prev_dagrun_dep import PrevDagrunDep
from airflow.ti_deps.deps.trigger_rule_dep import TriggerRuleDep
-from airflow.triggers.base import BaseTrigger, StartTriggerArgs
+from airflow.triggers.base import BaseTrigger
from airflow.utils.code_utils import get_python_source
from airflow.utils.context import (
ConnectionAccessor,
diff --git a/airflow-core/src/airflow/triggers/base.py b/airflow-core/src/airflow/triggers/base.py
index 490423da5fd..dcb43be9add 100644
--- a/airflow-core/src/airflow/triggers/base.py
+++ b/airflow-core/src/airflow/triggers/base.py
@@ -19,8 +19,6 @@ from __future__ import annotations
import abc
import json
from collections.abc import AsyncIterator
-from dataclasses import dataclass
-from datetime import timedelta
from typing import Annotated, Any
import structlog
@@ -33,20 +31,25 @@ from pydantic import (
)
from airflow.utils.log.logging_mixin import LoggingMixin
+from airflow.utils.module_loading import import_string
from airflow.utils.state import TaskInstanceState
log = structlog.get_logger(logger_name=__name__)
-@dataclass
-class StartTriggerArgs:
- """Arguments required for start task execution from triggerer."""
+def __getattr__(name: str):
+ if name == "StartTriggerArgs":
+ import warnings
- trigger_cls: str
- next_method: str
- trigger_kwargs: dict[str, Any] | None = None
- next_kwargs: dict[str, Any] | None = None
- timeout: timedelta | None = None
+ warnings.warn(
+ "airflow.triggers.base.StartTriggerArgs is deprecated. "
+ "Use airflow.sdk.bases.trigger.StartTriggerArgs instead.",
+ DeprecationWarning,
+ stacklevel=2,
+ )
+ return import_string(f"airflow.sdk.bases.trigger.{name}")
+
+ raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
class BaseTrigger(abc.ABC, LoggingMixin):
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index 9ad062ca1a4..68637caae36 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -42,7 +42,7 @@ from airflow.jobs.triggerer_job_runner import (
TriggerRunnerSupervisor,
messages,
)
-from airflow.models import DagModel, DagRun, TaskInstance, Trigger
+from airflow.models import DagBag, DagModel, DagRun, TaskInstance, Trigger
from airflow.models.baseoperator import BaseOperator
from airflow.models.connection import Connection
from airflow.models.dag import DAG
@@ -129,6 +129,15 @@ def create_trigger_in_db(session, trigger, operator=None):
return dag_model, run, trigger_orm, task_instance
+def mock_dag_bag(mock_dag_bag_cls, task_instance: TaskInstance):
+ mock_dag = MagicMock(spec=DAG)
+ mock_dag.get_task.return_value = task_instance.task
+
+ mock_dag_bag = MagicMock(spec=DagBag)
+ mock_dag_bag.get_dag.return_value = mock_dag
+ mock_dag_bag_cls.return_value = mock_dag_bag
+
+
def test_is_needed(session):
"""Checks the triggerer-is-needed logic"""
# No triggers, no need
@@ -207,7 +216,8 @@ def supervisor_builder(mocker, session):
return builder
-def test_trigger_lifecycle(spy_agency: SpyAgency, session, testing_dag_bundle):
+@patch("airflow.jobs.triggerer_job_runner.DagBag")
+def test_trigger_lifecycle(mock_dag_bag_cls, spy_agency: SpyAgency, session, testing_dag_bundle):
"""
Checks that the triggerer will correctly see a new Trigger in the database
and send it to the trigger runner, and then delete it when it vanishes.
@@ -216,6 +226,8 @@ def test_trigger_lifecycle(spy_agency: SpyAgency, session, testing_dag_bundle):
# (we want to avoid it firing and deleting itself)
trigger = TimeDeltaTrigger(datetime.timedelta(days=7))
dag_model, run, trigger_orm, task_instance = create_trigger_in_db(session, trigger)
+ mock_dag_bag(mock_dag_bag_cls, task_instance)
+
# Make a TriggererJobRunner and have it retrieve DB tasks
trigger_runner_supervisor = TriggerRunnerSupervisor.start(job=Job(id=12345), capacity=10)
@@ -398,7 +410,10 @@ class TestTriggerRunner:
@pytest.mark.asyncio
-async def test_trigger_create_race_condition_38599(session, supervisor_builder, testing_dag_bundle):
+@patch("airflow.jobs.triggerer_job_runner.DagBag")
+async def test_trigger_create_race_condition_38599(
+ mock_dag_bag_cls, session, supervisor_builder, testing_dag_bundle
+):
"""
This verifies the resolution of race condition documented in github issue #38599.
More details in the issue description.
@@ -427,10 +442,14 @@ async def test_trigger_create_race_condition_38599(session, supervisor_builder,
dm = DagModel(dag_id="test-dag", bundle_name=bundle_name)
session.add(dm)
SerializedDagModel.write_dag(dag, bundle_name=bundle_name)
- dag_run = DagRun(dag.dag_id, run_id="abc", run_type="none", run_after=timezone.utcnow())
+ dag_run = DagRun(
+ dag.dag_id, run_id="abc", run_type="manual", start_date=timezone.utcnow(), run_after=timezone.utcnow()
+ )
dag_version = DagVersion.get_latest_version(dag.dag_id)
+ task = PythonOperator(task_id="dummy-task", python_callable=print)
+ task.dag = dag
ti = TaskInstance(
- PythonOperator(task_id="dummy-task", python_callable=print),
+ task,
run_id=dag_run.run_id,
state=TaskInstanceState.DEFERRED,
dag_version_id=dag_version.id,
@@ -447,6 +466,8 @@ async def test_trigger_create_race_condition_38599(session, supervisor_builder,
session.commit()
+ mock_dag_bag(mock_dag_bag_cls, ti)
+
supervisor1 = supervisor_builder(job1)
supervisor2 = supervisor_builder(job2)
@@ -580,7 +601,8 @@ async def test_trigger_failing():
info["task"].cancel()
-def test_failed_trigger(session, dag_maker, supervisor_builder):
+@patch("airflow.jobs.triggerer_job_runner.DagBag")
+def test_failed_trigger(mock_dag_bag_cls, session, dag_maker, supervisor_builder):
"""
Checks that the triggerer will correctly fail task instances that depend on
triggers that can't even be loaded.
@@ -603,6 +625,8 @@ def test_failed_trigger(session, dag_maker, supervisor_builder):
task_instance.trigger_id = trigger_orm.id
session.commit()
+ mock_dag_bag(mock_dag_bag_cls, task_instance)
+
supervisor: TriggerRunnerSupervisor = supervisor_builder()
supervisor.load_triggers()
@@ -746,7 +770,8 @@ class DummyTriggerRunnerSupervisor(TriggerRunnerSupervisor):
@pytest.mark.asyncio
@pytest.mark.execution_timeout(20)
-async def test_trigger_can_call_variables_connections_and_xcoms_methods(session, dag_maker):
+@patch("airflow.jobs.triggerer_job_runner.DagBag")
+async def test_trigger_can_call_variables_connections_and_xcoms_methods(mock_dag_bag_cls, session, dag_maker):
"""Checks that the trigger will successfully call Variables, Connections and XComs methods."""
# Create the test DAG and task
with dag_maker(dag_id="trigger_accessing_variable_connection_and_xcom", session=session):
@@ -808,6 +833,8 @@ async def test_trigger_can_call_variables_connections_and_xcoms_methods(session,
session.add(job)
session.commit()
+ mock_dag_bag(mock_dag_bag_cls, task_instance)
+
supervisor = DummyTriggerRunnerSupervisor.start(job=job, capacity=1, logger=None)
supervisor.run()
@@ -878,7 +905,10 @@ class CustomTriggerDagRun(BaseTrigger):
@pytest.mark.asyncio
@pytest.mark.execution_timeout(10)
-async def test_trigger_can_fetch_trigger_dag_run_count_and_state_in_deferrable(session, dag_maker):
+@patch("airflow.jobs.triggerer_job_runner.DagBag")
+async def test_trigger_can_fetch_trigger_dag_run_count_and_state_in_deferrable(
+ mock_dag_bag_cls, session, dag_maker
+):
"""Checks that the trigger will successfully fetch the count of trigger DAG runs."""
# Create the test DAG and task
with dag_maker(dag_id="trigger_can_fetch_trigger_dag_run_count_and_state_in_deferrable", session=session):
@@ -909,6 +939,8 @@ async def test_trigger_can_fetch_trigger_dag_run_count_and_state_in_deferrable(s
session.add(job)
session.commit()
+ mock_dag_bag(mock_dag_bag_cls, task_instance)
+
supervisor = DummyTriggerRunnerSupervisor.start(job=job, capacity=1, logger=None)
supervisor.run()
@@ -969,7 +1001,8 @@ class CustomTriggerWorkflowStateTrigger(BaseTrigger):
@pytest.mark.asyncio
@pytest.mark.execution_timeout(10)
-async def test_trigger_can_fetch_dag_run_count_ti_count_in_deferrable(session, dag_maker):
+@patch("airflow.jobs.triggerer_job_runner.DagBag")
+async def test_trigger_can_fetch_dag_run_count_ti_count_in_deferrable(mock_dag_bag_cls, session, dag_maker):
"""Checks that the trigger will successfully fetch the count of DAG runs, Task count and task states."""
# Create the test DAG and task
with dag_maker(dag_id="parent_dag", session=session):
@@ -1010,6 +1043,8 @@ async def test_trigger_can_fetch_dag_run_count_ti_count_in_deferrable(session, d
session.add(job)
session.commit()
+ mock_dag_bag(mock_dag_bag_cls, task_instance)
+
supervisor = DummyTriggerRunnerSupervisor.start(job=job, capacity=1, logger=None)
supervisor.run()
@@ -1022,7 +1057,10 @@ async def test_trigger_can_fetch_dag_run_count_ti_count_in_deferrable(session, d
}
-def test_update_triggers_prevents_duplicate_creation_queue_entries(session, supervisor_builder):
+@patch("airflow.jobs.triggerer_job_runner.DagBag")
+def test_update_triggers_prevents_duplicate_creation_queue_entries(
+ mock_dag_bag_cls, session, supervisor_builder
+):
"""
Test that update_triggers prevents adding triggers to the creation queue
if they are already queued for creation.
@@ -1030,6 +1068,8 @@ def test_update_triggers_prevents_duplicate_creation_queue_entries(session, supe
trigger = TimeDeltaTrigger(datetime.timedelta(days=7))
dag_model, run, trigger_orm, task_instance = create_trigger_in_db(session, trigger)
+ mock_dag_bag(mock_dag_bag_cls, task_instance)
+
supervisor = supervisor_builder()
# First call to update_triggers should add the trigger to creating_triggers
@@ -1051,8 +1091,9 @@ def test_update_triggers_prevents_duplicate_creation_queue_entries(session, supe
assert not any(trigger_id == trigger_orm.id for trigger_id, _ in supervisor.failed_triggers)
+@patch("airflow.jobs.triggerer_job_runner.DagBag")
def test_update_triggers_prevents_duplicate_creation_queue_entries_with_multiple_triggers(
- session, supervisor_builder, dag_maker
+ mock_dag_bag_cls, session, supervisor_builder, dag_maker
):
"""
Test that update_triggers prevents adding multiple triggers to the creation queue
@@ -1063,6 +1104,8 @@ def test_update_triggers_prevents_duplicate_creation_queue_entries_with_multiple
dag_model1, run1, trigger_orm1, task_instance1 = create_trigger_in_db(session, trigger1)
+ mock_dag_bag(mock_dag_bag_cls, task_instance1)
+
with dag_maker("test_dag_2"):
EmptyOperator(task_id="test_ti_2")
@@ -1071,6 +1114,9 @@ def test_update_triggers_prevents_duplicate_creation_queue_entries_with_multiple
ti2 = run2.task_instances[0]
session.add(trigger_orm2)
session.flush()
+
+ mock_dag_bag(mock_dag_bag_cls, ti2)
+
ti2.trigger_id = trigger_orm2.id
session.merge(ti2)
session.flush()
diff --git a/airflow-core/tests/unit/models/test_dagrun.py b/airflow-core/tests/unit/models/test_dagrun.py
index 9c85dc041e6..9dfe096b193 100644
--- a/airflow-core/tests/unit/models/test_dagrun.py
+++ b/airflow-core/tests/unit/models/test_dagrun.py
@@ -44,11 +44,11 @@ from airflow.providers.standard.operators.bash import BashOperator
from airflow.providers.standard.operators.empty import EmptyOperator
from airflow.providers.standard.operators.python import PythonOperator, ShortCircuitOperator
from airflow.sdk import BaseOperator, setup, task, task_group, teardown
+from airflow.sdk.bases.trigger import StartTriggerArgs
from airflow.sdk.definitions.deadline import AsyncCallback, DeadlineAlert, DeadlineReference
from airflow.serialization.serialized_objects import SerializedDAG
from airflow.stats import Stats
from airflow.task.trigger_rule import TriggerRule
-from airflow.triggers.base import StartTriggerArgs
from airflow.utils.span_status import SpanStatus
from airflow.utils.state import DagRunState, State, TaskInstanceState
from airflow.utils.thread_safe_dict import ThreadSafeDict
diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py
index 2959b5af458..3420bae225a 100644
--- a/airflow-core/tests/unit/serialization/test_dag_serialization.py
+++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py
@@ -66,6 +66,7 @@ from airflow.providers.standard.sensors.bash import BashSensor
from airflow.sdk import AssetAlias, BaseHook, teardown
from airflow.sdk.bases.decorator import DecoratedOperator
from airflow.sdk.bases.operator import BaseOperator
+from airflow.sdk.bases.trigger import StartTriggerArgs
from airflow.sdk.definitions._internal.expandinput import EXPAND_INPUT_EMPTY
from airflow.sdk.definitions.asset import Asset, AssetUniqueKey
from airflow.sdk.definitions.operator_resources import Resources
@@ -83,7 +84,6 @@ from airflow.serialization.serialized_objects import (
from airflow.task.priority_strategy import _DownstreamPriorityWeightStrategy
from airflow.ti_deps.deps.ready_to_reschedule import ReadyToRescheduleDep
from airflow.timetables.simple import NullTimetable, OnceTimetable
-from airflow.triggers.base import StartTriggerArgs
from airflow.utils.module_loading import qualname
from tests_common.test_utils.config import conf_vars
diff --git a/task-sdk/docs/api.rst b/task-sdk/docs/api.rst
index 40ea8e17287..fc375ab48c5 100644
--- a/task-sdk/docs/api.rst
+++ b/task-sdk/docs/api.rst
@@ -72,6 +72,8 @@ Bases
.. autoapiclass:: airflow.sdk.BaseHook
+.. autoapiclass:: airflow.sdk.StartTriggerArgs
+
Connections & Variables
-----------------------
.. autoapiclass:: airflow.sdk.Connection
diff --git a/task-sdk/src/airflow/sdk/__init__.py b/task-sdk/src/airflow/sdk/__init__.py
index cc29ef4057b..7c4cfacc227 100644
--- a/task-sdk/src/airflow/sdk/__init__.py
+++ b/task-sdk/src/airflow/sdk/__init__.py
@@ -39,6 +39,7 @@ __all__ = [
"ObjectStoragePath",
"Param",
"PokeReturnValue",
+ "StartTriggerArgs",
"TaskGroup",
"TriggerRule",
"Variable",
@@ -67,6 +68,7 @@ if TYPE_CHECKING:
from airflow.sdk.bases.operator import BaseOperator, chain, chain_linear, cross_downstream
from airflow.sdk.bases.operatorlink import BaseOperatorLink
from airflow.sdk.bases.sensor import BaseSensorOperator, PokeReturnValue
+ from airflow.sdk.bases.trigger import StartTriggerArgs
from airflow.sdk.definitions.asset import Asset, AssetAlias, AssetAll, AssetAny, AssetWatcher
from airflow.sdk.definitions.asset.decorators import asset
from airflow.sdk.definitions.asset.metadata import Metadata
@@ -104,6 +106,7 @@ __lazy_imports: dict[str, str] = {
"Param": ".definitions.param",
"PokeReturnValue": ".bases.sensor",
"SecretCache": ".execution_time.cache",
+ "StartTriggerArgs": ".bases.trigger",
"TaskGroup": ".definitions.taskgroup",
"TriggerRule": ".api.datamodels._generated",
"Variable": ".definitions.variable",
diff --git a/task-sdk/src/airflow/sdk/__init__.pyi b/task-sdk/src/airflow/sdk/__init__.pyi
index c1112d26c04..5fe59a9f194 100644
--- a/task-sdk/src/airflow/sdk/__init__.pyi
+++ b/task-sdk/src/airflow/sdk/__init__.pyi
@@ -29,6 +29,7 @@ from airflow.sdk.bases.sensor import (
BaseSensorOperator as BaseSensorOperator,
PokeReturnValue as PokeReturnValue,
)
+from airflow.sdk.bases.trigger import StartTriggerArgs as StartTriggerArgs
from airflow.sdk.definitions.asset import (
Asset as Asset,
AssetAlias as AssetAlias,
@@ -78,6 +79,7 @@ __all__ = [
"Param",
"PokeReturnValue",
"SecretCache",
+ "StartTriggerArgs",
"TaskGroup",
"TriggerRule",
"Variable",
diff --git a/task-sdk/src/airflow/sdk/bases/operator.py b/task-sdk/src/airflow/sdk/bases/operator.py
index 3ece4b99372..762e2118500 100644
--- a/task-sdk/src/airflow/sdk/bases/operator.py
+++ b/task-sdk/src/airflow/sdk/bases/operator.py
@@ -89,6 +89,7 @@ if TYPE_CHECKING:
import jinja2
from airflow.sdk.bases.operatorlink import BaseOperatorLink
+ from airflow.sdk.bases.trigger import StartTriggerArgs
from airflow.sdk.definitions.context import Context
from airflow.sdk.definitions.dag import DAG
from airflow.sdk.definitions.operator_resources import Resources
@@ -96,7 +97,7 @@ if TYPE_CHECKING:
from airflow.sdk.definitions.xcom_arg import XComArg
from airflow.serialization.enums import DagAttributeTypes
from airflow.task.priority_strategy import PriorityWeightStrategy
- from airflow.triggers.base import BaseTrigger, StartTriggerArgs
+ from airflow.triggers.base import BaseTrigger
from airflow.typing_compat import Self
TaskPreExecuteHook = Callable[[Context], None]
@@ -945,9 +946,6 @@ class BaseOperator(AbstractOperator, metaclass=BaseOperatorMeta):
# Set to True for an operator instantiated by a mapped operator.
__from_mapped: bool = False
- start_trigger_args: StartTriggerArgs | None = None
- start_from_trigger: bool = False
-
# base list which includes all the attrs that don't need deep copy.
_base_operator_shallow_copy_attrs: Final[tuple[str, ...]] = (
"user_defined_macros",
diff --git a/task-sdk/src/airflow/sdk/bases/trigger.py b/task-sdk/src/airflow/sdk/bases/trigger.py
new file mode 100644
index 00000000000..5429b4c6ade
--- /dev/null
+++ b/task-sdk/src/airflow/sdk/bases/trigger.py
@@ -0,0 +1,33 @@
+# 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
+
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Any
+
+
+@dataclass
+class StartTriggerArgs:
+ """Arguments required for start task execution from triggerer."""
+
+ trigger_cls: str
+ next_method: str
+ trigger_kwargs: dict[str, Any] | None = None
+ next_kwargs: dict[str, Any] | None = None
+ timeout: timedelta | None = None
diff --git a/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py b/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
index 7629c06f795..f7db152a5cd 100644
--- a/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
+++ b/task-sdk/src/airflow/sdk/definitions/_internal/abstractoperator.py
@@ -43,6 +43,7 @@ if TYPE_CHECKING:
from airflow.sdk.bases.operator import BaseOperator
from airflow.sdk.bases.operatorlink import BaseOperatorLink
+ from airflow.sdk.bases.trigger import StartTriggerArgs
from airflow.sdk.definitions.dag import DAG
from airflow.sdk.definitions.mappedoperator import MappedOperator
from airflow.sdk.definitions.taskgroup import MappedTaskGroup
@@ -119,6 +120,9 @@ class AbstractOperator(Templater, DAGNode):
is_setup: bool = False
is_teardown: bool = False
+ start_trigger_args: StartTriggerArgs | None = None
+ start_from_trigger: bool = False
+
HIDE_ATTRS_FROM_UI: ClassVar[frozenset[str]] = frozenset(
(
"log",
@@ -315,6 +319,14 @@ class AbstractOperator(Templater, DAGNode):
else:
setattr(parent, attr_name, rendered_content)
+ if (
+ self.start_from_trigger
+ and self.start_trigger_args
+ and self.start_trigger_args.trigger_kwargs
+ ):
+ if attr_name in self.start_trigger_args.trigger_kwargs:
+ self.start_trigger_args.trigger_kwargs[attr_name] = rendered_content
+
def _iter_all_mapped_downstreams(self) -> Iterator[MappedOperator | MappedTaskGroup]:
"""
Return mapped nodes that are direct dependencies of the current task.
diff --git a/task-sdk/src/airflow/sdk/definitions/mappedoperator.py b/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
index febc922ba26..a2e7e7bf11c 100644
--- a/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
+++ b/task-sdk/src/airflow/sdk/definitions/mappedoperator.py
@@ -64,11 +64,19 @@ if TYPE_CHECKING:
OperatorExpandArgument,
OperatorExpandKwargsArgument,
)
- from airflow.sdk import DAG, BaseOperator, BaseOperatorLink, Context, TaskGroup, TriggerRule, XComArg
+ from airflow.sdk import (
+ DAG,
+ BaseOperator,
+ BaseOperatorLink,
+ Context,
+ StartTriggerArgs,
+ TaskGroup,
+ TriggerRule,
+ XComArg,
+ )
from airflow.sdk.definitions._internal.expandinput import ExpandInput
from airflow.sdk.definitions.operator_resources import Resources
from airflow.sdk.definitions.param import ParamsDict
- from airflow.triggers.base import StartTriggerArgs
ValidationSource = Literal["expand"] | Literal["partial"]
@@ -817,7 +825,7 @@ class MappedOperator(AbstractOperator):
This method is for allowing mapped operator to start execution from triggerer.
"""
- from airflow.triggers.base import StartTriggerArgs
+ from airflow.sdk.bases.trigger import StartTriggerArgs
if not self.start_trigger_args:
return None