diff --git a/airflow-core/src/airflow/api/common/trigger_dag.py b/airflow-core/src/airflow/api/common/trigger_dag.py index 4f8ca4c25e718..8dfd06c1ac9e8 100644 --- a/airflow-core/src/airflow/api/common/trigger_dag.py +++ b/airflow-core/src/airflow/api/common/trigger_dag.py @@ -118,6 +118,8 @@ def _trigger_dag( if dag_run := DagRun.find_duplicate(dag_id=dag_id, run_id=run_id): raise DagRunAlreadyExists(dag_run) + partition_date = dag.timetable.resolve_partition_date(partition_key) + run_conf = None if is_arg_set(conf): run_conf = _normalize_conf(conf) @@ -133,6 +135,7 @@ def _trigger_dag( note=note, state=DagRunState.QUEUED, partition_key=partition_key, + partition_date=partition_date, session=session, ) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py index 88ffceca6d7c7..033437e2ad98b 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dag_run.py @@ -191,6 +191,9 @@ def validate_context(self, dag: SerializedDAG) -> dict: run_after=timezone.coerce_datetime(run_after), data_interval=data_interval, ) + + partition_date = dag.timetable.resolve_partition_date(self.partition_key) + return { "run_id": run_id, "logical_date": coerced_logical_date, @@ -199,6 +202,7 @@ def validate_context(self, dag: SerializedDAG) -> dict: "conf": self.conf, "note": self.note, "partition_key": self.partition_key, + "partition_date": partition_date, } diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py index 0860b29da9424..30ae090289c75 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -657,6 +657,7 @@ def trigger_dag_run( triggering_user_name=user.get_name(), state=DagRunState.QUEUED, partition_key=params["partition_key"], + partition_date=params["partition_date"], session=session, ) except (ParamValidationError, ValueError) as e: 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 872c74957527e..1092369e913bc 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 @@ -34,7 +34,7 @@ from airflow.api_fastapi.execution_api.datamodels.taskinstance import DagRun from airflow.api_fastapi.execution_api.datamodels.token import TIToken from airflow.api_fastapi.execution_api.security import CurrentTIToken -from airflow.exceptions import DagNotPartitionedError, DagRunAlreadyExists +from airflow.exceptions import DagNotPartitionedError, DagRunAlreadyExists, InvalidPartitionKeyError from airflow.models.dag import DagModel from airflow.models.dagrun import DagRun as DagRunModel from airflow.models.taskinstance import TaskInstance @@ -158,6 +158,11 @@ def trigger_dag_run( status.HTTP_400_BAD_REQUEST, detail={"reason": "not_partitioned", "message": str(e)}, ) + except InvalidPartitionKeyError as e: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, + detail={"reason": "invalid_partition_key", "message": str(e)}, + ) from e @router.post( diff --git a/airflow-core/src/airflow/exceptions.py b/airflow-core/src/airflow/exceptions.py index ff45daebc02e2..addd4f8f2a547 100644 --- a/airflow-core/src/airflow/exceptions.py +++ b/airflow-core/src/airflow/exceptions.py @@ -146,7 +146,12 @@ class DagNotPartitionedError(ValueError): class InvalidPartitionKeyError(ValueError): - """Raise when a partition_key value is empty or exceeds the maximum allowed length.""" + """ + Raise when a partition_key value is invalid. + + 1. empty or exceeds the maximum allowed length + 2. cannot be decoded to a partition_date by the timetable + """ class DagRunAlreadyExists(AirflowBadRequest): diff --git a/airflow-core/src/airflow/timetables/base.py b/airflow-core/src/airflow/timetables/base.py index b07e5c98c9352..400b63f3d7090 100644 --- a/airflow-core/src/airflow/timetables/base.py +++ b/airflow-core/src/airflow/timetables/base.py @@ -292,6 +292,45 @@ def resolve_day_bound(self, day: datetime.date) -> DateTime: datetime.datetime(day.year, day.month, day.day, tzinfo=datetime.timezone.utc) ) + def resolve_partition_date(self, partition_key: str | None) -> datetime.datetime | None: + """ + Decode *partition_key* into the period-start datetime it represents. + + Returns the timezone-aware datetime that was used to format *partition_key* + when the timetable originally created the run, or ``None`` when no temporal + meaning can be derived from the key. ``None`` is returned without decoding + when *partition_key* is ``None``, when this timetable is not ``partitioned``, + or when it defers partition selection to runtime (``partitioned_at_runtime``). + + Partitioned timetables whose keys carry a temporal structure override + :meth:`_decode_partition_date`: + + - :class:`~airflow.timetables.trigger.CronPartitionTimetable` parses the + key with ``strptime`` using its ``key_format`` and localizes with its + timezone. + - :class:`~airflow.timetables.simple.PartitionedAssetTimetable` delegates + to each asset's partition mapper; when the mappers agree on the same + instant it is returned, otherwise ``None`` is returned. + + :param partition_key: The partition key string to decode, or ``None``. + :returns: The period-start datetime, or ``None`` if not resolvable. + :raises InvalidPartitionKeyError: When *partition_key* is syntactically + invalid for this timetable's key format (e.g. ``strptime`` fails). + """ + if partition_key is None or not self.partitioned or self.partitioned_at_runtime: + return None + return self._decode_partition_date(partition_key) + + def _decode_partition_date(self, partition_key: str) -> datetime.datetime | None: + """ + Decode a non-empty *partition_key* into its period-start datetime. + + Called by :meth:`resolve_partition_date` only after the partitioned-state + guards pass. The default returns ``None``; partitioned timetables whose + keys carry temporal structure override this. + """ + return None + @property def partition_mapper_info(self) -> list[PartitionMapperInfo]: """ diff --git a/airflow-core/src/airflow/timetables/simple.py b/airflow-core/src/airflow/timetables/simple.py index 496ba9da156ed..a8b499dbe50ba 100644 --- a/airflow-core/src/airflow/timetables/simple.py +++ b/airflow-core/src/airflow/timetables/simple.py @@ -17,11 +17,13 @@ from __future__ import annotations from contextlib import suppress +from datetime import datetime from typing import TYPE_CHECKING, Any, TypeAlias import structlog from airflow._shared.timezones import timezone +from airflow.exceptions import InvalidPartitionKeyError from airflow.partition_mappers.identity import IdentityMapper from airflow.serialization.definitions.assets import ( SerializedAsset, @@ -355,6 +357,46 @@ def partition_mapper_info(self) -> list[PartitionMapperInfo]: entries.append(PartitionMapperInfo(uri=s_asset_ref.uri, is_rollup=mapper.is_rollup)) return entries + def _decode_partition_date(self, partition_key: str) -> datetime | None: + """ + Decode *partition_key* into the period-start datetime shared by all asset mappers. + + Iterates every asset (and asset ref) reachable from the asset condition, asks + each mapper for the temporal anchor of *partition_key*, and returns it when all + temporal mappers agree. Returns ``None`` when no mapper is temporal or when the + mappers disagree — consistent with how the scheduler resolves ``partition_date`` + for asset-triggered runs. + """ + anchors: set[datetime] = set() + for unique_key, _ in self.asset_condition.iter_assets(): + mapper = self.get_partition_mapper(name=unique_key.name, uri=unique_key.uri) + try: + anchor = mapper.to_partition_date(partition_key) + except ValueError as exc: + raise InvalidPartitionKeyError( + f"Partition key {partition_key!r} is invalid for this timetable's mappers: {exc}" + ) from exc + if anchor is not None: + anchors.add(anchor) + for s_asset_ref in self.asset_condition.iter_asset_refs(): + if isinstance(s_asset_ref, SerializedAssetNameRef): + mapper = self.get_partition_mapper(name=s_asset_ref.name) + elif isinstance(s_asset_ref, SerializedAssetUriRef): + mapper = self.get_partition_mapper(uri=s_asset_ref.uri) + else: + continue + try: + anchor = mapper.to_partition_date(partition_key) + except ValueError as exc: + raise InvalidPartitionKeyError( + f"Partition key {partition_key!r} is invalid for this timetable's mappers: {exc}" + ) from exc + if anchor is not None: + anchors.add(anchor) + if len(anchors) == 1: + return anchors.pop() + return None + def serialize(self) -> dict[str, Any]: from airflow.serialization.serialized_objects import encode_asset_like diff --git a/airflow-core/src/airflow/timetables/trigger.py b/airflow-core/src/airflow/timetables/trigger.py index b3ac7cbf2e071..a1de05ac76e65 100644 --- a/airflow-core/src/airflow/timetables/trigger.py +++ b/airflow-core/src/airflow/timetables/trigger.py @@ -27,7 +27,8 @@ import structlog -from airflow._shared.timezones.timezone import coerce_datetime, parse_timezone, utcnow +from airflow._shared.timezones.timezone import coerce_datetime, make_aware, parse_timezone, utcnow +from airflow.exceptions import InvalidPartitionKeyError from airflow.timetables._cron import CronMixin from airflow.timetables._delta import DeltaMixin from airflow.timetables.base import DagRunInfo, DataInterval, Timetable @@ -518,6 +519,26 @@ def _format_key(self, partition_date: DateTime) -> str: # midnight partition keys as "...T00:00:00", not the prior UTC day's "...T16:00:00"). return partition_date.in_timezone(self._timezone).strftime(self._key_format) + def _decode_partition_date(self, partition_key: str) -> datetime.datetime: + """ + Decode *partition_key* back to the period-start datetime. + + Parses the key with ``strptime`` using this timetable's ``key_format`` + and localizes with the timetable's timezone, mirroring the forward + direction in :meth:`_format_key`. + + :raises InvalidPartitionKeyError: When *partition_key* does not match + the timetable's ``key_format``. + """ + try: + naive = datetime.datetime.strptime(partition_key, self._key_format) + except ValueError as exc: + raise InvalidPartitionKeyError( + f"Partition key {partition_key!r} does not match the timetable's " + f"key_format {self._key_format!r}: {exc}" + ) from exc + return make_aware(naive, self._timezone) + def next_dagrun_info_v2( self, *, diff --git a/airflow-core/tests/unit/api/common/test_trigger_dag.py b/airflow-core/tests/unit/api/common/test_trigger_dag.py index 49a385865f409..6fc48bd081411 100644 --- a/airflow-core/tests/unit/api/common/test_trigger_dag.py +++ b/airflow-core/tests/unit/api/common/test_trigger_dag.py @@ -16,12 +16,17 @@ # under the License. from __future__ import annotations +from datetime import datetime, timezone + import pytest from sqlalchemy import select from airflow.api.common.trigger_dag import trigger_dag +from airflow.exceptions import InvalidPartitionKeyError from airflow.models import DagModel from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.timetables.simple import PartitionAtRuntime +from airflow.timetables.trigger import CronPartitionTimetable from airflow.utils.types import DagRunTriggeredByType, DagRunType from tests_common.test_utils.db import ( @@ -83,6 +88,69 @@ def test_trigger_dag_with_custom_run_type(dag_maker, session): assert dag_run.run_type == DagRunType.ASSET_MATERIALIZATION +def test_trigger_dag_populates_partition_date_for_cron_partition_timetable(dag_maker, session): + """Manually triggering a CronPartitionTimetable Dag with a partition_key populates partition_date.""" + with dag_maker( + session=session, + dag_id="TEST_CRON_PARTITION_DAG", + schedule=CronPartitionTimetable("0 0 * * *", timezone="UTC"), + ): + EmptyOperator(task_id="mytask") + session.commit() + + dag_run = trigger_dag( + dag_id="TEST_CRON_PARTITION_DAG", + triggered_by=DagRunTriggeredByType.REST_API, + partition_key="2025-06-01T00:00:00", + session=session, + ) + + assert dag_run is not None + assert dag_run.partition_key == "2025-06-01T00:00:00" + assert dag_run.partition_date == datetime(2025, 6, 1, 0, 0, 0, tzinfo=timezone.utc) + + +def test_trigger_dag_raises_invalid_partition_key_for_cron_partition_timetable(dag_maker, session): + """A malformed partition_key for CronPartitionTimetable raises InvalidPartitionKeyError.""" + with dag_maker( + session=session, + dag_id="TEST_CRON_PARTITION_BAD_KEY", + schedule=CronPartitionTimetable("0 0 * * *", timezone="UTC"), + ): + EmptyOperator(task_id="mytask") + session.commit() + + with pytest.raises(InvalidPartitionKeyError, match="does not match"): + trigger_dag( + dag_id="TEST_CRON_PARTITION_BAD_KEY", + triggered_by=DagRunTriggeredByType.REST_API, + partition_key="not-a-valid-date", + session=session, + ) + + +def test_trigger_dag_partition_at_runtime_leaves_partition_date_none(dag_maker, session): + """PartitionAtRuntime Dags accept arbitrary keys; partition_date stays None.""" + with dag_maker( + session=session, + dag_id="TEST_PARTITION_AT_RUNTIME", + schedule=PartitionAtRuntime(), + ): + EmptyOperator(task_id="mytask") + session.commit() + + dag_run = trigger_dag( + dag_id="TEST_PARTITION_AT_RUNTIME", + triggered_by=DagRunTriggeredByType.REST_API, + partition_key="arbitrary-runtime-key", + session=session, + ) + + assert dag_run is not None + assert dag_run.partition_key == "arbitrary-runtime-key" + assert dag_run.partition_date is None + + def test_trigger_dag_operator_denied_when_only_manual_allowed(dag_maker, session): with dag_maker(session=session, dag_id="TEST_DAG_1", schedule="0 * * * *"): EmptyOperator(task_id="mytask") diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py index 0d0478f8d01e6..ad11f764d4839 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py @@ -46,7 +46,7 @@ from airflow.sdk.definitions.param import Param from airflow.settings import _configure_async_session from airflow.timetables.interval import CronDataIntervalTimetable -from airflow.timetables.simple import PartitionAtRuntime +from airflow.timetables.simple import PartitionAtRuntime, PartitionedAssetTimetable from airflow.timetables.trigger import CronPartitionTimetable from airflow.utils.session import provide_session from airflow.utils.state import DagRunState, State @@ -2874,7 +2874,7 @@ def test_should_respond_200_when_exactly_max_length_partition_key(self, dag_make partitioned_dag_id = "test_max_length_partition_key" with dag_maker( dag_id=partitioned_dag_id, - schedule=CronPartitionTimetable("0 * * * *", timezone="UTC"), + schedule=PartitionedAssetTimetable(assets=Asset("test")), start_date=START_DATE1, session=session, serialized=True, @@ -2889,6 +2889,86 @@ def test_should_respond_200_when_exactly_max_length_partition_key(self, dag_make ) assert response.status_code == 200 + @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") + def test_trigger_partitioned_dag_populates_partition_date(self, dag_maker, test_client, session): + """Triggering a CronPartitionTimetable Dag with a valid key populates partition_date on the run. + + Regression guard: before this fix partition_date was NULL for manually triggered runs even + when partition_key was supplied, making partition-date-based filtering (e.g. + ``airflow dags clear --partition-date-*``) silently skip those runs. + """ + partitioned_dag_id = "test_trigger_populates_partition_date" + with dag_maker( + dag_id=partitioned_dag_id, + schedule=CronPartitionTimetable("0 0 * * *", timezone="UTC"), + start_date=START_DATE1, + session=session, + serialized=True, + ): + EmptyOperator(task_id="task") + + session.commit() + + response = test_client.post( + f"/dags/{partitioned_dag_id}/dagRuns", + json={"logical_date": None, "partition_key": "2025-06-01T00:00:00"}, + ) + assert response.status_code == 200 + + dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == partitioned_dag_id)) + assert dag_run is not None + assert dag_run.partition_key == "2025-06-01T00:00:00" + assert dag_run.partition_date == datetime(2025, 6, 1, 0, 0, 0, tzinfo=timezone.utc) + + @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") + def test_trigger_partitioned_dag_invalid_key_returns_400(self, dag_maker, test_client, session): + """An invalid partition_key for a CronPartitionTimetable Dag must return HTTP 400.""" + partitioned_dag_id = "test_trigger_invalid_partition_key" + with dag_maker( + dag_id=partitioned_dag_id, + schedule=CronPartitionTimetable("0 0 * * *", timezone="UTC"), + start_date=START_DATE1, + session=session, + serialized=True, + ): + EmptyOperator(task_id="task") + + session.commit() + + response = test_client.post( + f"/dags/{partitioned_dag_id}/dagRuns", + json={"logical_date": None, "partition_key": "not-a-valid-date"}, + ) + assert response.status_code == 400 + + @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") + def test_trigger_partition_at_runtime_dag_leaves_partition_date_none( + self, dag_maker, test_client, session + ): + """PartitionAtRuntime Dag with an arbitrary key must produce partition_date=None.""" + runtime_dag_id = "test_trigger_partition_at_runtime_none_date" + with dag_maker( + dag_id=runtime_dag_id, + schedule=PartitionAtRuntime(), + start_date=START_DATE1, + session=session, + serialized=True, + ): + EmptyOperator(task_id="task") + + session.commit() + + response = test_client.post( + f"/dags/{runtime_dag_id}/dagRuns", + json={"logical_date": None, "partition_key": "arbitrary-key"}, + ) + assert response.status_code == 200 + + dag_run = session.scalar(select(DagRun).where(DagRun.dag_id == runtime_dag_id)) + assert dag_run is not None + assert dag_run.partition_key == "arbitrary-key" + assert dag_run.partition_date is None + class TestResolveRunOnLatestVersion: @pytest.mark.parametrize("explicit_value", [True, False]) diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py index 500e1ba4de4c3..5d3c37c9b521d 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_dag_runs.py @@ -71,7 +71,7 @@ def test_trigger_dag_run_with_partition_key(self, client, session, dag_maker): dag_id = "test_trigger_dag_run_partition_key" run_id = "test_run_id" logical_date = timezone.datetime(2025, 2, 20) - partition_key = "2025-02-20" + partition_key = "2025-02-20T00:00:00" with dag_maker( dag_id=dag_id, @@ -126,6 +126,31 @@ def test_trigger_dag_run_partition_key_for_non_partitioned_dag(self, client, ses } } + def test_trigger_dag_run_invalid_partition_key(self, client, session, dag_maker): + """partition_key that the timetable cannot decode must return 400.""" + dag_id = "test_trigger_dag_run_invalid_partition_key" + run_id = "test_run_id_invalid_pk" + logical_date = timezone.datetime(2025, 2, 20) + + with dag_maker( + dag_id=dag_id, + schedule=CronPartitionTimetable("0 * * * *", timezone="UTC"), + session=session, + serialized=True, + ): + EmptyOperator(task_id="test_task") + session.commit() + + response = client.post( + f"/execution/dag-runs/{dag_id}/{run_id}", + json={"logical_date": logical_date.isoformat(), "partition_key": "not-a-date"}, + ) + + assert response.status_code == 400 + detail = response.json()["detail"] + assert detail["reason"] == "invalid_partition_key" + assert "does not match the timetable's key_format" in detail["message"] + def test_trigger_dag_run_dag_not_found(self, client): """Test that a DAG that does not exist cannot be triggered.""" dag_id = "dag_not_found" diff --git a/airflow-core/tests/unit/timetables/test_partitioned_timetable.py b/airflow-core/tests/unit/timetables/test_partitioned_timetable.py index 23caf06e79d0e..cb5d882179421 100644 --- a/airflow-core/tests/unit/timetables/test_partitioned_timetable.py +++ b/airflow-core/tests/unit/timetables/test_partitioned_timetable.py @@ -27,6 +27,7 @@ import pytest from airflow._shared.module_loading import qualname +from airflow.exceptions import InvalidPartitionKeyError from airflow.partition_mappers.base import RollupMapper from airflow.partition_mappers.identity import IdentityMapper as IdentityMapper from airflow.partition_mappers.temporal import StartOfDayMapper @@ -258,6 +259,28 @@ def test_deserialize(self): assert isinstance(timetable.default_partition_mapper, IdentityMapper) assert isinstance(timetable.partition_mapper_config[ser_asset], IdentityMapper) + @pytest.mark.parametrize( + "asset_like", + [ + Asset(name="daily", uri="s3://bucket/daily"), + Asset.ref(name="daily"), + ], + ) + def test_decode_partition_date_raises_on_invalid_key(self, asset_like): + timetable = PartitionedAssetTimetable( + assets=asset_like, + default_partition_mapper=StartOfDayMapper(), + ) + with pytest.raises(InvalidPartitionKeyError, match="is invalid for this timetable's mappers"): + timetable._decode_partition_date("not-a-date") + + def test_decode_partition_date_returns_period_start_for_valid_key(self): + timetable = PartitionedAssetTimetable( + assets=Asset(name="daily", uri="s3://bucket/daily"), + default_partition_mapper=StartOfDayMapper(), + ) + assert timetable._decode_partition_date("2025-01-01") == pendulum.datetime(2025, 1, 1, tz="UTC") + def test_partitioned_asset_timetable_resolve_day_bound_returns_midnight_utc(self): """PartitionedAssetTimetable has no local timezone; resolve_day_bound uses the base default.