This is an automated email from the ASF dual-hosted git repository.
jason810496 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 1300e566658 Add run_after to TriggerDagRunOperator (#62259)
1300e566658 is described below
commit 1300e56665824c27ac019f9705d1042041d2b0b3
Author: arnoldmr01 <[email protected]>
AuthorDate: Wed Apr 29 21:37:44 2026 +0800
Add run_after to TriggerDagRunOperator (#62259)
* feat: add run_after to DagRunTriggerException
* feat: add support to TriggerDagRunOperator
feat: add support to TriggerDagRunOperator
* feat: add run_after to TriggerDagRunPayload
* feat: add run_after support to task_runner and supervisor
* feat: Add _validate_datetime_param
* feat: Add _parse_datetime_param
* feat: add versioning for TriggerDagRunPayload
* fix: add overload to _parse and _validate def
* test: update unit test
* fix: import is_arg_set
* chore: typo modification
* fix: fix unit test
* fix: fix static check
* feat: Add comment for unsupported attrs in AF2
---------
Co-authored-by: Arnold Lin <[email protected]>
---
.../api_fastapi/execution_api/datamodels/dagrun.py | 1 +
.../api_fastapi/execution_api/routes/dag_runs.py | 1 +
.../api_fastapi/execution_api/versions/__init__.py | 2 +
.../execution_api/versions/v2026_04_06.py | 15 ++-
.../providers/standard/operators/trigger_dagrun.py | 113 ++++++++++++++++-----
.../airflow/providers/standard/version_compat.py | 13 ++-
.../unit/standard/operators/test_trigger_dagrun.py | 61 ++++++++++-
task-sdk/src/airflow/sdk/api/client.py | 7 +-
.../src/airflow/sdk/api/datamodels/_generated.py | 1 +
task-sdk/src/airflow/sdk/exceptions.py | 2 +
.../src/airflow/sdk/execution_time/supervisor.py | 2 +-
.../src/airflow/sdk/execution_time/task_runner.py | 1 +
task-sdk/tests/task_sdk/api/test_client.py | 5 +-
.../task_sdk/execution_time/test_supervisor.py | 41 +++++++-
14 files changed, 229 insertions(+), 36 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dagrun.py
b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dagrun.py
index 8619901f717..dd1f1aeb866 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dagrun.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/dagrun.py
@@ -28,6 +28,7 @@ class TriggerDAGRunPayload(StrictBaseModel):
"""Schema for Trigger DAG Run API request."""
logical_date: UtcDateTime | None = None
+ run_after: UtcDateTime | None = None
conf: dict = Field(default_factory=dict)
reset_dag_run: bool = False
partition_key: str | None = None
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py
b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py
index 448d0b6945a..d4a0d24024b 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/dag_runs.py
@@ -128,6 +128,7 @@ def trigger_dag_run(
run_type=DagRunType.OPERATOR_TRIGGERED,
conf=payload.conf,
logical_date=payload.logical_date,
+ run_after=payload.run_after,
triggered_by=DagRunTriggeredByType.OPERATOR,
replace_microseconds=False,
partition_key=payload.partition_key,
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
index 6863c64efe5..0a33e2e3af5 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py
@@ -33,6 +33,7 @@ from airflow.api_fastapi.execution_api.versions.v2026_04_06
import (
AddDagRunDetailEndpoint,
AddNoteField,
AddPartitionKeyField,
+ AddRunAfterField,
AddTaskInstanceStartDateField,
MakeDagRunStartDateNullable,
ModifyDeferredTaskKwargsToJsonValue,
@@ -57,6 +58,7 @@ bundle = VersionBundle(
RemoveUpstreamMapIndexesField,
AddNoteField,
AddTaskInstanceStartDateField,
+ AddRunAfterField,
AddDagEndpoint,
),
Version("2025-11-05", AddTriggeringUserNameField),
diff --git
a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py
b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py
index be9e167e654..5a3b0f0f5fc 100644
--- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py
+++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_04_06.py
@@ -169,7 +169,10 @@ class AddNoteField(VersionChange):
description = __doc__
- instructions_to_migrate_to_previous_version =
(schema(DagRun).field("note").didnt_exist,)
+ instructions_to_migrate_to_previous_version = (
+ schema(DagRun).field("note").didnt_exist,
+ schema(TriggerDAGRunPayload).field("note").didnt_exist,
+ )
@convert_response_to_previous_version_for(TIRunContext) # type:
ignore[arg-type]
def remove_note_field(response: ResponseInfo) -> None: # type:
ignore[misc]
@@ -197,3 +200,13 @@ class AddDagEndpoint(VersionChange):
description = __doc__
instructions_to_migrate_to_previous_version = (endpoint("/dags/{dag_id}",
["GET"]).didnt_exist,)
+
+
+class AddRunAfterField(VersionChange):
+ """Add run_after parameter to TriggerDAGRunPayload Model."""
+
+ description = __doc__
+
+ instructions_to_migrate_to_previous_version = (
+ schema(TriggerDAGRunPayload).field("run_after").didnt_exist,
+ )
diff --git
a/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py
b/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py
index 26b990f9c9e..73884433af5 100644
---
a/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py
+++
b/providers/standard/src/airflow/providers/standard/operators/trigger_dagrun.py
@@ -23,7 +23,7 @@ import json
import time
from collections.abc import Sequence
from json import JSONDecodeError
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast, overload
from sqlalchemy import select
from sqlalchemy.orm.exc import NoResultFound
@@ -43,7 +43,7 @@ from airflow.providers.common.compat.sdk import (
)
from airflow.providers.standard.triggers.external_task import DagStateTrigger
from airflow.providers.standard.utils.openlineage import
safe_inject_openlineage_properties_into_dagrun_conf
-from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS,
BaseOperator
+from airflow.providers.standard.version_compat import AIRFLOW_V_3_0_PLUS,
BaseOperator, is_arg_set
from airflow.utils.state import DagRunState
from airflow.utils.types import DagRunType
@@ -127,6 +127,7 @@ class TriggerDagRunOperator(BaseOperator):
If not provided, a run ID will be automatically generated.
:param conf: Configuration for the DAG run (templated).
:param logical_date: Logical date for the triggered DAG (templated).
+ :param run_after: The date before which the triggered DAG should not run.
:param reset_dag_run: Whether clear existing DAG run if already exists.
This is useful when backfill or rerun an existing DAG run.
This only resets (not recreates) the DAG run.
@@ -162,6 +163,13 @@ class TriggerDagRunOperator(BaseOperator):
"wait_for_completion",
"skip_when_already_exists",
)
+
+ attributes_not_supported_in_airflow_2 = {
+ # `run_after` uses NOTSET here so we can detect whether the user
+ # explicitly provided it and warn in Airflow 2.
+ "run_after": NOTSET,
+ "note": None,
+ }
template_fields_renderers = {"conf": "py"}
ui_color = "#ffefeb"
operator_extra_links = [TriggerDagRunLink()]
@@ -173,6 +181,7 @@ class TriggerDagRunOperator(BaseOperator):
trigger_run_id: str | None = None,
conf: dict | None = None,
logical_date: str | datetime.datetime | None | ArgNotSet = NOTSET,
+ run_after: str | datetime.datetime | None | ArgNotSet = NOTSET,
reset_dag_run: bool = False,
wait_for_completion: bool = False,
poke_interval: int = 60,
@@ -205,27 +214,29 @@ class TriggerDagRunOperator(BaseOperator):
self.openlineage_inject_parent_info = openlineage_inject_parent_info
self.note = note
self.deferrable = deferrable
+ logical_date = _validate_datetime_param("logical_date", logical_date)
+ run_after = _validate_datetime_param("run_after", run_after)
self.logical_date = logical_date
- if logical_date is NOTSET:
- self.logical_date = NOTSET
- elif logical_date is None or isinstance(logical_date, (str,
datetime.datetime)):
- self.logical_date = logical_date
- else:
- raise TypeError(
- f"Expected str, datetime.datetime, or None for parameter
'logical_date'. Got {type(logical_date).__name__}"
- )
-
+ self.run_after = run_after
if fail_when_dag_is_paused and AIRFLOW_V_3_0_PLUS:
raise NotImplementedError("Setting `fail_when_dag_is_paused` not
yet supported for Airflow 3.x")
def execute(self, context: Context):
if self.logical_date is NOTSET:
- # If no logical_date is provided we will set utcnow()
- parsed_logical_date = timezone.utcnow()
- elif self.logical_date is None or isinstance(self.logical_date,
datetime.datetime):
- parsed_logical_date = self.logical_date # type: ignore
- elif isinstance(self.logical_date, str):
- parsed_logical_date = timezone.parse(self.logical_date)
+ if self.run_after is not NOTSET:
+ parsed_logical_date = None
+ else:
+ # If no logical_date is provided we will set utcnow()
+ parsed_logical_date = timezone.utcnow()
+ else:
+ logical_date = cast("str | datetime.datetime | None",
self.logical_date)
+ parsed_logical_date = _parse_datetime_param(logical_date)
+
+ if self.run_after is NOTSET:
+ parsed_run_after = parsed_logical_date
+ else:
+ run_after = cast("str | datetime.datetime | None", self.run_after)
+ parsed_run_after = _parse_datetime_param(run_after)
try:
if self.conf and isinstance(self.conf, str):
@@ -247,7 +258,7 @@ class TriggerDagRunOperator(BaseOperator):
run_id = DagRun.generate_run_id(
run_type=DagRunType.MANUAL,
logical_date=parsed_logical_date,
- run_after=parsed_logical_date or timezone.utcnow(),
+ run_after=parsed_run_after or timezone.utcnow(),
)
else:
run_id = DagRun.generate_run_id(DagRunType.MANUAL,
parsed_logical_date or timezone.utcnow()) # type: ignore[misc,call-arg]
@@ -267,14 +278,17 @@ class TriggerDagRunOperator(BaseOperator):
if AIRFLOW_V_3_0_PLUS:
self._trigger_dag_af_3(
- context=context, run_id=self.trigger_run_id,
parsed_logical_date=parsed_logical_date
+ context=context,
+ run_id=self.trigger_run_id,
+ parsed_logical_date=parsed_logical_date,
+ parsed_run_after=parsed_run_after if self.run_after is not
NOTSET else None,
)
else:
self._trigger_dag_af_2(
context=context, run_id=self.trigger_run_id,
parsed_logical_date=parsed_logical_date
)
- def _trigger_dag_af_3(self, context, run_id, parsed_logical_date):
+ def _trigger_dag_af_3(self, context, run_id, parsed_logical_date,
parsed_run_after=None):
from airflow.providers.common.compat.sdk import DagRunTriggerException
kwargs_accepted = dict(
@@ -291,16 +305,28 @@ class TriggerDagRunOperator(BaseOperator):
deferrable=self.deferrable,
)
- if self.note and "note" in
inspect.signature(DagRunTriggerException.__init__).parameters:
+ parameters =
inspect.signature(DagRunTriggerException.__init__).parameters
+ if self.note and "note" in parameters:
kwargs_accepted["note"] = self.note
+ if parsed_run_after and "run_after" in parameters:
+ kwargs_accepted["run_after"] = parsed_run_after
+
raise DagRunTriggerException(**kwargs_accepted)
def _trigger_dag_af_2(self, context, run_id, parsed_logical_date):
try:
- if self.note:
- self.log.warning("Parameter 'note' is not supported in Airflow
2.x and will be ignored.")
-
+ unsupported_parameters = []
+ for attr, default_value in
self.attributes_not_supported_in_airflow_2.items():
+ value = getattr(self, attr, default_value)
+ if value is not default_value:
+ unsupported_parameters.append(attr)
+
+ if unsupported_parameters:
+ self.log.warning(
+ "The following parameters are not supported in Airflow 2.x
and will be ignored: %s",
+ ", ".join(unsupported_parameters),
+ )
dag_run = trigger_dag(
dag_id=self.trigger_dag_id,
run_id=run_id,
@@ -453,3 +479,42 @@ class TriggerDagRunOperator(BaseOperator):
f"{self.trigger_dag_id} return {state} which is not in
{self.failed_states}"
f" or {self.allowed_states}"
)
+
+
+@overload
+def _validate_datetime_param(name: str, value: ArgNotSet) -> ArgNotSet: ...
+@overload
+def _validate_datetime_param(name: str, value: None) -> None: ...
+@overload
+def _validate_datetime_param(name: str, value: str) -> str: ...
+@overload
+def _validate_datetime_param(name: str, value: datetime.datetime) ->
datetime.datetime: ...
+
+
+def _validate_datetime_param(
+ name: str,
+ value: str | datetime.datetime | None | ArgNotSet,
+) -> str | datetime.datetime | None | ArgNotSet:
+ if not is_arg_set(value):
+ return NOTSET
+ if value is None or isinstance(value, (str, datetime.datetime)):
+ return value
+ raise TypeError(
+ f"Expected str, datetime.datetime, or None for parameter '{name}'. Got
{type(value).__name__}"
+ )
+
+
+@overload
+def _parse_datetime_param(value: None) -> None: ...
+@overload
+def _parse_datetime_param(value: datetime.datetime) -> datetime.datetime: ...
+@overload
+def _parse_datetime_param(value: str) -> datetime.datetime: ...
+
+
+def _parse_datetime_param(
+ value: str | datetime.datetime | None,
+) -> datetime.datetime | None:
+ if value is None or isinstance(value, datetime.datetime):
+ return value
+ return timezone.parse(value)
diff --git
a/providers/standard/src/airflow/providers/standard/version_compat.py
b/providers/standard/src/airflow/providers/standard/version_compat.py
index db58809b127..f99a5589378 100644
--- a/providers/standard/src/airflow/providers/standard/version_compat.py
+++ b/providers/standard/src/airflow/providers/standard/version_compat.py
@@ -42,10 +42,18 @@ AIRFLOW_V_3_3_PLUS: bool = get_base_airflow_version_tuple()
>= (3, 3, 0)
# This is needed for DecoratedOperator compatibility
if AIRFLOW_V_3_1_PLUS:
from airflow.sdk import BaseOperator
- from airflow.sdk.definitions._internal.types import ArgNotSet
+ from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet
else:
from airflow.models.baseoperator import BaseOperator # type:
ignore[no-redef]
- from airflow.utils.types import ArgNotSet # type:
ignore[attr-defined,no-redef]
+ from airflow.utils.types import NOTSET, ArgNotSet # type:
ignore[attr-defined,no-redef]
+
+try:
+ from airflow.sdk.definitions._internal.types import is_arg_set
+except ImportError:
+
+ def is_arg_set(value): # type: ignore[misc,no-redef]
+ return value is not NOTSET
+
__all__ = [
"AIRFLOW_V_3_0_PLUS",
@@ -54,4 +62,5 @@ __all__ = [
"AIRFLOW_V_3_3_PLUS",
"ArgNotSet",
"BaseOperator",
+ "is_arg_set",
]
diff --git
a/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py
b/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py
index 29b4132b235..bcc5ad6153c 100644
--- a/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py
+++ b/providers/standard/tests/unit/standard/operators/test_trigger_dagrun.py
@@ -104,6 +104,44 @@ class TestDagRunOperator:
session.execute(delete(DagBundleModel))
session.commit()
+ @pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is
different for Airflow 2 & 3")
+ def test_trigger_dagrun_with_run_after(self):
+ """
+ Test TriggerDagRunOperator.
+
+ We only verify that the operator runs and raises correct exception.
The actual execution logic
+ after the exception is in Task SDK code.
+ """
+ with time_machine.travel("2025-02-18T08:04:46Z", tick=False):
+ task = TriggerDagRunOperator(
+ task_id="test_task",
+ trigger_dag_id=TRIGGERED_DAG_ID,
+ conf={"foo": "bar"},
+ run_after=timezone.datetime(2025, 2, 19, 12, 0, 0),
+ )
+
+ # Ensure correct exception is raised
+ with pytest.raises(DagRunTriggerException) as exc_info:
+ task.execute(context={})
+
+ assert exc_info.value.trigger_dag_id == TRIGGERED_DAG_ID
+ assert exc_info.value.conf == {"foo": "bar"}
+ assert exc_info.value.logical_date is None
+ assert exc_info.value.reset_dag_run is False
+ assert exc_info.value.skip_when_already_exists is False
+ assert exc_info.value.wait_for_completion is False
+ assert exc_info.value.allowed_states == [DagRunState.SUCCESS]
+ assert exc_info.value.failed_states == [DagRunState.FAILED]
+ if getattr(exc_info, "note", None) is not None:
+ assert exc_info.value.note == "Test note"
+
+ expected_run_id = DagRun.generate_run_id(
+ run_type=DagRunType.MANUAL, run_after=task.run_after
+ ).rsplit("_", 1)[0]
+ # rsplit because last few characters are random.
+ assert exc_info.value.dag_run_id.rsplit("_", 1)[0] ==
expected_run_id
+ assert task.trigger_run_id.rsplit("_", 1)[0] == expected_run_id #
run_id is saved as attribute
+
@pytest.mark.skipif(not AIRFLOW_V_3_0_PLUS, reason="Implementation is
different for Airflow 2 & 3")
def test_trigger_dagrun(self):
"""
@@ -557,7 +595,7 @@ class TestDagRunOperatorAF2:
with time_machine.travel("2025-02-18T08:04:46Z", tick=False):
with dag_maker(TEST_DAG_ID, default_args={"start_date":
DEFAULT_DATE}, serialized=True):
task = TriggerDagRunOperator(
- task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID,
note="Test note"
+ task_id="test_task", trigger_dag_id=TRIGGERED_DAG_ID,
note="Test note", run_after=None
)
mock_warning = mock.patch.object(task.log, "warning").start()
dag_maker.sync_dagbag_to_db()
@@ -566,12 +604,31 @@ class TestDagRunOperatorAF2:
task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
ignore_ti_state=True)
dagrun =
dag_maker.session.scalar(select(DagRun).where(DagRun.dag_id ==
TRIGGERED_DAG_ID))
+ unsupported_params = ["run_after", "note"]
assert mock_warning.mock_calls == [
- mock.call("Parameter 'note' is not supported in Airflow 2.x
and will be ignored.")
+ mock.call(
+ "The following parameters are not supported in Airflow 2.x
and will be ignored: %s",
+ ", ".join(unsupported_params),
+ )
]
assert dagrun.run_type == DagRunType.MANUAL
assert dagrun.run_id == DagRun.generate_run_id(DagRunType.MANUAL,
dagrun.logical_date)
+ def test_trigger_dagrun_does_not_warn_for_default_unsupported_params(
+ self, dag_maker, mock_supervisor_comms
+ ):
+ """Test TriggerDagRunOperator does not warn for unsupported params
when they are not provided."""
+ with time_machine.travel("2025-02-18T08:04:46Z", tick=False):
+ with dag_maker(TEST_DAG_ID, default_args={"start_date":
DEFAULT_DATE}, serialized=True):
+ task = TriggerDagRunOperator(task_id="test_task",
trigger_dag_id=TRIGGERED_DAG_ID)
+ mock_warning = mock.patch.object(task.log, "warning").start()
+ dag_maker.sync_dagbag_to_db()
+ parse_and_sync_to_db(self.f_name)
+ dag_maker.create_dagrun()
+ task.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE,
ignore_ti_state=True)
+
+ mock_warning.assert_not_called()
+
def test_explicitly_provided_trigger_run_id_is_saved_as_attr(self,
dag_maker, session):
with dag_maker(TEST_DAG_ID, default_args={"start_date": DEFAULT_DATE},
serialized=True):
task = TriggerDagRunOperator(
diff --git a/task-sdk/src/airflow/sdk/api/client.py
b/task-sdk/src/airflow/sdk/api/client.py
index d0362ce12aa..a662b900664 100644
--- a/task-sdk/src/airflow/sdk/api/client.py
+++ b/task-sdk/src/airflow/sdk/api/client.py
@@ -711,12 +711,17 @@ class DagRunOperations:
run_id: str,
conf: dict | None = None,
logical_date: datetime | None = None,
+ run_after: datetime | None = None,
reset_dag_run: bool = False,
note: str | None = None,
) -> OKResponse | ErrorResponse:
"""Trigger a Dag run via the API server."""
body = TriggerDAGRunPayload(
- logical_date=logical_date, conf=conf or {},
reset_dag_run=reset_dag_run, note=note
+ logical_date=logical_date,
+ conf=conf or {},
+ reset_dag_run=reset_dag_run,
+ note=note,
+ run_after=run_after,
)
try:
diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
index 138fe8a9fd8..bafd6890f46 100644
--- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
+++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py
@@ -371,6 +371,7 @@ class TriggerDAGRunPayload(BaseModel):
extra="forbid",
)
logical_date: Annotated[AwareDatetime | None, Field(title="Logical Date")]
= None
+ run_after: Annotated[AwareDatetime | None, Field(title="Run After")] = None
conf: Annotated[dict[str, Any] | None, Field(title="Conf")] = None
reset_dag_run: Annotated[bool | None, Field(title="Reset Dag Run")] = False
partition_key: Annotated[str | None, Field(title="Partition Key")] = None
diff --git a/task-sdk/src/airflow/sdk/exceptions.py
b/task-sdk/src/airflow/sdk/exceptions.py
index c7698ab75ff..7d42dad5d85 100644
--- a/task-sdk/src/airflow/sdk/exceptions.py
+++ b/task-sdk/src/airflow/sdk/exceptions.py
@@ -241,6 +241,7 @@ class DagRunTriggerException(AirflowException):
dag_run_id: str,
conf: dict | None,
logical_date=None,
+ run_after=None,
reset_dag_run: bool,
skip_when_already_exists: bool,
wait_for_completion: bool,
@@ -255,6 +256,7 @@ class DagRunTriggerException(AirflowException):
self.dag_run_id = dag_run_id
self.conf = conf
self.logical_date = logical_date
+ self.run_after = run_after
self.reset_dag_run = reset_dag_run
self.skip_when_already_exists = skip_when_already_exists
self.wait_for_completion = wait_for_completion
diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
index aa70bd99675..abc0260e213 100644
--- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py
+++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py
@@ -1523,7 +1523,7 @@ class ActivitySubprocess(WatchedSubprocess):
dump_opts = {"exclude_unset": True}
elif isinstance(msg, TriggerDagRun):
resp = self.client.dag_runs.trigger(
- msg.dag_id, msg.run_id, msg.conf, msg.logical_date,
msg.reset_dag_run, msg.note
+ msg.dag_id, msg.run_id, msg.conf, msg.logical_date,
msg.run_after, msg.reset_dag_run, msg.note
)
elif isinstance(msg, GetDagRun):
dr_resp = self.client.dag_runs.get_detail(msg.dag_id, msg.run_id)
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index c2398c0b295..fe0aa81f077 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -1458,6 +1458,7 @@ def _handle_trigger_dag_run(
dag_id=drte.trigger_dag_id,
run_id=drte.dag_run_id,
logical_date=drte.logical_date,
+ run_after=drte.run_after,
conf=drte.conf,
reset_dag_run=drte.reset_dag_run,
note=drte.note,
diff --git a/task-sdk/tests/task_sdk/api/test_client.py
b/task-sdk/tests/task_sdk/api/test_client.py
index 0df8839c55f..26e2a7e66bb 100644
--- a/task-sdk/tests/task_sdk/api/test_client.py
+++ b/task-sdk/tests/task_sdk/api/test_client.py
@@ -1238,7 +1238,10 @@ class TestDagRunOperations:
client = make_client(transport=httpx.MockTransport(handle_request))
result = client.dag_runs.trigger(
- dag_id="test_trigger", run_id="test_run_id",
logical_date=timezone.datetime(2025, 1, 1)
+ dag_id="test_trigger",
+ run_id="test_run_id",
+ logical_date=timezone.datetime(2025, 1, 1),
+ run_after=timezone.datetime(2025, 1, 2),
)
assert result == OKResponse(ok=True)
diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
index 334dfee4e10..c7fc804dcea 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py
@@ -2156,15 +2156,24 @@ REQUEST_TEST_CASES = [
run_id="test_run",
conf={"key": "value"},
logical_date=timezone.datetime(2025, 1, 1),
+ run_after=timezone.datetime(2025, 1, 1, 12, 0, 0),
reset_dag_run=True,
),
expected_body={"ok": True, "type": "OKResponse"},
client_mock=ClientMock(
method_path="dag_runs.trigger",
- args=("test_dag", "test_run", {"key": "value"},
timezone.datetime(2025, 1, 1), True, None),
+ args=(
+ "test_dag",
+ "test_run",
+ {"key": "value"},
+ timezone.datetime(2025, 1, 1),
+ timezone.datetime(2025, 1, 1, 12, 0, 0),
+ True,
+ None,
+ ),
response=OKResponse(ok=True),
),
- test_id="dag_run_trigger",
+ test_id="dag_run_trigger_with_run_after",
),
RequestTestCase(
message=TriggerDagRun(
@@ -2178,7 +2187,31 @@ REQUEST_TEST_CASES = [
expected_body={"ok": True, "type": "OKResponse"},
client_mock=ClientMock(
method_path="dag_runs.trigger",
- args=("test_dag", "test_run", {"key": "value"},
timezone.datetime(2025, 1, 1), True, "Test Note"),
+ args=(
+ "test_dag",
+ "test_run",
+ {"key": "value"},
+ timezone.datetime(2025, 1, 1),
+ None,
+ True,
+ "Test Note",
+ ),
+ response=OKResponse(ok=True),
+ ),
+ test_id="dag_run_trigger",
+ ),
+ RequestTestCase(
+ message=TriggerDagRun(
+ dag_id="test_dag",
+ run_id="test_run",
+ conf={"key": "value"},
+ logical_date=timezone.datetime(2025, 1, 1),
+ reset_dag_run=True,
+ ),
+ expected_body={"ok": True, "type": "OKResponse"},
+ client_mock=ClientMock(
+ method_path="dag_runs.trigger",
+ args=("test_dag", "test_run", {"key": "value"},
timezone.datetime(2025, 1, 1), None, True, None),
response=OKResponse(ok=True),
),
test_id="dag_run_trigger",
@@ -2188,7 +2221,7 @@ REQUEST_TEST_CASES = [
expected_body={"error": "DAGRUN_ALREADY_EXISTS", "detail": None,
"type": "ErrorResponse"},
client_mock=ClientMock(
method_path="dag_runs.trigger",
- args=("test_dag", "test_run", None, None, False, None),
+ args=("test_dag", "test_run", None, None, None, False, None),
response=ErrorResponse(error=ErrorType.DAGRUN_ALREADY_EXISTS),
),
test_id="dag_run_trigger_already_exists",
