diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dags.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dags.py index e67777d6ffa4e..b2058765650bb 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dags.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/dags.py @@ -208,6 +208,7 @@ class DAGDetailsResponse(DAGResponse): is_favorite: bool = False active_runs_count: int = 0 team_name: str | None = None + suggested_partition_key: str | None = None @field_validator("timezone", mode="before") @classmethod diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 5200635ec9e4f..ec9f150eb1b13 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -13110,6 +13110,11 @@ components: - type: string - type: 'null' title: Team Name + suggested_partition_key: + anyOf: + - type: string + - type: 'null' + title: Suggested Partition Key is_backfillable: type: boolean title: Is Backfillable diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dags.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dags.py index 112e6c9839190..3f20567cde612 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dags.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dags.py @@ -24,6 +24,7 @@ from pydantic import ValidationError from sqlalchemy import delete, func, insert, select, update +from airflow._shared.timezones.timezone import coerce_datetime, utcnow from airflow.api.common import delete_dag as delete_dag_module from airflow.api_fastapi.common.dagbag import DagBagDep, get_latest_version_of_dag from airflow.api_fastapi.common.db.common import SessionDep, apply_filters_to_select, paginated_select @@ -258,6 +259,9 @@ def get_dag_details( # Add is_favorite and active_runs_count fields to the Dag model setattr(dag_model, "is_favorite", is_favorite) setattr(dag_model, "active_runs_count", active_runs_count) + setattr( + dag_model, "suggested_partition_key", dag.timetable.suggest_partition_key(coerce_datetime(utcnow())) + ) return DAGDetailsResponse.model_validate(dag_model) diff --git a/airflow-core/src/airflow/timetables/base.py b/airflow-core/src/airflow/timetables/base.py index 365f980b93274..f1d8455fb98c9 100644 --- a/airflow-core/src/airflow/timetables/base.py +++ b/airflow-core/src/airflow/timetables/base.py @@ -338,6 +338,35 @@ def _decode_partition_date(self, partition_key: str) -> datetime.datetime | None """ return None + def suggest_partition_key(self, now: DateTime) -> str | None: + """ + Suggest a partition key for a manual trigger happening at *now*. + + Used to pre-populate the partition key field when a user manually triggers a Dag + run through the UI, so they are not required to know the exact key format by + heart. Returns ``None`` when this timetable is not ``partitioned``, when it defers + partition selection to runtime (``partitioned_at_runtime``), or when the concrete + timetable does not implement a suggestion (the default). + + :param now: The instant to compute the suggestion for — normally "right now", the + moment the trigger form is rendered. + :returns: A partition key string ready to pre-fill the trigger form, or ``None`` if + no suggestion is available. + """ + if not self.partitioned or self.partitioned_at_runtime: + return None + return self._suggest_partition_key(now) + + def _suggest_partition_key(self, now: DateTime) -> str | None: + """ + Compute a suggested partition key for *now*. + + Called by :meth:`suggest_partition_key` only after the partitioned-state guards + pass. The default returns ``None``; partitioned timetables that can derive a + "current" partition from a point in time override this. + """ + return None + @property def partition_mapper_info(self) -> list[PartitionMapperInfo]: """ diff --git a/airflow-core/src/airflow/timetables/trigger.py b/airflow-core/src/airflow/timetables/trigger.py index af3f1c23c4cf1..dd9d09fe26788 100644 --- a/airflow-core/src/airflow/timetables/trigger.py +++ b/airflow-core/src/airflow/timetables/trigger.py @@ -467,6 +467,18 @@ def _get_partition_info(self, run_date: DateTime) -> tuple[DateTime, str]: partition_key = self._format_key(partition_date) return partition_date, partition_key + def _suggest_partition_key(self, now: DateTime) -> str: + """ + Suggest the key of the most recently elapsed partition as of *now*. + + The tick is derived from the cron expression alone, so the suggestion is not + clamped by the Dag's ``start_date`` — a Dag that has not started scheduling yet + still gets the tick preceding *now*, which is a valid key for a manual run. + """ + aligned = self._align_to_prev(coerce_datetime(now)) + _, partition_key = self._get_partition_info(run_date=aligned) + return partition_key + def iter_partition_dagrun_infos( self, *, diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index d9afa239aa2e9..36b9f8dd9d4c4 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -3258,6 +3258,17 @@ export const $DAGDetailsResponse = { ], title: 'Team Name' }, + suggested_partition_key: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Suggested Partition Key' + }, is_backfillable: { type: 'boolean', title: 'Is Backfillable', diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index cf8c86e997cf4..34c993381db99 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -902,6 +902,7 @@ export type DAGDetailsResponse = { is_favorite?: boolean; active_runs_count?: number; team_name?: string | null; + suggested_partition_key?: string | null; /** * Whether this Dag's schedule supports backfilling. */ diff --git a/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGForm.test.tsx b/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGForm.test.tsx index 2181842f68154..2479fee35bc1b 100644 --- a/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGForm.test.tsx +++ b/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGForm.test.tsx @@ -214,4 +214,151 @@ describe("TriggerDAGForm", () => { await waitFor(() => expect(screen.getByText("dagRun.partitionKey")).toBeInTheDocument()); expect(screen.getByText("components:triggerDag.partitionKeyHelp")).toBeInTheDocument(); }); + + it("pre-fills the partition key field with the suggested value", async () => { + const { container } = render( + , + { wrapper: Wrapper }, + ); + + fireEvent.click(screen.getByText("Advanced Options")); + + await waitFor(() => + expect(container.querySelector('input[name="partitionKey"]')).toBeInTheDocument(), + ); + const partitionKeyField = container.querySelector('input[name="partitionKey"]'); + + expect(partitionKeyField?.value).toBe("2026-02-18T00:00:00"); + }); + + it("leaves the partition key field blank when no suggestion is provided", async () => { + const { container } = render( + , + { wrapper: Wrapper }, + ); + + fireEvent.click(screen.getByText("Advanced Options")); + + await waitFor(() => + expect(container.querySelector('input[name="partitionKey"]')).toBeInTheDocument(), + ); + const partitionKeyField = container.querySelector('input[name="partitionKey"]'); + + expect(partitionKeyField?.value).toBe(""); + }); + + it("applies a fresher suggestion arriving after the form is shown", async () => { + const { container, rerender } = render( + , + { wrapper: Wrapper }, + ); + + fireEvent.click(screen.getByText("Advanced Options")); + + await waitFor(() => + expect(container.querySelector('input[name="partitionKey"]')).toBeInTheDocument(), + ); + + rerender( + , + ); + + await waitFor(() => + expect(container.querySelector('input[name="partitionKey"]')?.value).toBe( + "2026-02-18T01:00:00", + ), + ); + }); + + it("keeps an edited partition key when a fresher suggestion arrives", async () => { + const { container, rerender } = render( + , + { wrapper: Wrapper }, + ); + + fireEvent.click(screen.getByText("Advanced Options")); + + await waitFor(() => + expect(container.querySelector('input[name="partitionKey"]')).toBeInTheDocument(), + ); + + const partitionKeyField = container.querySelector('input[name="partitionKey"]'); + + fireEvent.change(partitionKeyField as HTMLInputElement, { target: { value: "2026-01-01T00:00:00" } }); + + rerender( + , + ); + + await waitFor(() => + expect(container.querySelector('input[name="partitionKey"]')?.value).toBe( + "2026-01-01T00:00:00", + ), + ); + }); }); diff --git a/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGForm.tsx b/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGForm.tsx index 185c3da46f90b..e246218efca3d 100644 --- a/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGForm.tsx +++ b/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGForm.tsx @@ -53,6 +53,7 @@ type TriggerDAGFormProps = { runId: string; } | undefined; + readonly suggestedPartitionKey?: string | undefined; }; const TriggerDAGForm = ({ @@ -66,6 +67,7 @@ const TriggerDAGForm = ({ onSubmitTrigger, open, prefillConfig, + suggestedPartitionKey, }: TriggerDAGFormProps) => { const { t: translate } = useTranslation(["common", "components"]); const [errors, setErrors] = useState<{ conf?: string; date?: unknown }>({}); @@ -76,7 +78,14 @@ const TriggerDAGForm = ({ const [hasAppliedPrefill, setHasAppliedPrefill] = useState(false); const { mutate: togglePause } = useTogglePause({ dagId }); - const { control, handleSubmit, reset, watch } = useForm({ + const { + control, + formState: { dirtyFields }, + handleSubmit, + reset, + setValue, + watch, + } = useForm({ defaultValues: { conf, dagRunId: "", @@ -87,12 +96,13 @@ const TriggerDAGForm = ({ // For partitioned Dags, logical date is not applicable. logicalDate: isPartitioned ? "" : dayjs().format(DEFAULT_DATETIME_FORMAT), note: "", - partitionKey: undefined, + partitionKey: suggestedPartitionKey, }, }); // Pre-fill form when prefillConfig is provided (priority over conf) - // Only restore 'conf' (parameters), not logicalDate, runId, or partitionKey to avoid 409 conflicts + // Only restore 'conf' (parameters) from the previous run. Its logicalDate, runId and partition key + // would collide with that run (409), so the partition key is seeded from the current suggestion. useEffect(() => { if (prefillConfig && open) { const confString = prefillConfig.conf ? JSON.stringify(prefillConfig.conf, undefined, 2) : ""; @@ -105,7 +115,7 @@ const TriggerDAGForm = ({ dataIntervalStart: "", logicalDate: isPartitioned ? "" : dayjs().format(DEFAULT_DATETIME_FORMAT), note: "", - partitionKey: undefined, + partitionKey: suggestedPartitionKey, }); // Also update the param store to keep it in sync. Seed the initial params (for stable // section ordering) only once they are available, but always push the conf so a run's @@ -132,8 +142,17 @@ const TriggerDAGForm = ({ initialParamDict, setInitialParamDict, isPartitioned, + suggestedPartitionKey, ]); + // The suggestion is computed per request, so a response served from cache can seed the field with + // a partition that has already elapsed. Apply a newer suggestion unless the user edited the field. + useEffect(() => { + if (suggestedPartitionKey !== undefined && !Boolean(dirtyFields.partitionKey)) { + setValue("partitionKey", suggestedPartitionKey); + } + }, [dirtyFields.partitionKey, setValue, suggestedPartitionKey]); + // Automatically reset form when conf is fetched (only if no prefillConfig) useEffect(() => { if (conf && open && (!prefillConfig || hasAppliedPrefill)) { diff --git a/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGModal.tsx b/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGModal.tsx index 3aa26a29ebe62..34f1d28e3fc23 100644 --- a/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGModal.tsx +++ b/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGModal.tsx @@ -20,7 +20,7 @@ import { Heading, VStack, HStack, Spinner, Center, Text } from "@chakra-ui/react import React, { useState } from "react"; import { useTranslation } from "react-i18next"; -import { useDagServiceGetDag } from "openapi/queries"; +import { useDagServiceGetDagDetails } from "openapi/queries"; import { Dialog, Tooltip } from "src/components/ui"; import { RadioCardItem, RadioCardRoot } from "src/components/ui/RadioCard"; import { useTrigger } from "src/queries/useTrigger"; @@ -62,19 +62,22 @@ const TriggerDAGModal: React.FC = ({ data: dag, isError, isLoading, - } = useDagServiceGetDag( + } = useDagServiceGetDagDetails( { dagId, }, undefined, { enabled: open, + // suggested_partition_key is computed per request; a cached one may name an elapsed partition. + staleTime: 0, }, ); const isBackfillable = dag?.is_backfillable ?? false; const hasSchedule = dag?.timetable_summary !== null; const isPartitioned = dag ? dag.timetable_partitioned : false; + const suggestedPartitionKey = dag?.suggested_partition_key ?? undefined; const { error, isPending, triggerDagRun } = useTrigger({ dagId, onSuccessConfirm: onClose }); const maxDisplayLength = 59; // hard-coded length to prevent dag name overflowing the modal const nameOverflowing = dagDisplayName.length > maxDisplayLength; @@ -149,6 +152,7 @@ const TriggerDAGModal: React.FC = ({ onSubmitTrigger={triggerDagRun} open={open} prefillConfig={prefillConfig} + suggestedPartitionKey={suggestedPartitionKey} /> ) : ( isBackfillable && dag && diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py index 8873c82e27311..f0e91a1d3ffe0 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dags.py @@ -21,6 +21,7 @@ import pendulum import pytest +import time_machine from sqlalchemy import delete, insert, select, update from airflow.models.asset import AssetModel, DagScheduleAssetReference @@ -30,6 +31,7 @@ from airflow.models.dagrun import DagRun from airflow.models.team import Team from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.timetables.trigger import CronPartitionTimetable from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType @@ -1078,6 +1080,7 @@ def test_dag_details( "render_template_as_native_obj": False, "rerun_with_latest_version": None, "start_date": start_date, + "suggested_partition_key": None, "tags": [], "template_search_path": None, "timetable_description": "Never, external triggers only", @@ -1181,6 +1184,7 @@ def test_dag_details_with_view_url_template( "render_template_as_native_obj": False, "rerun_with_latest_version": None, "start_date": start_date, + "suggested_partition_key": None, "tags": [], "template_search_path": None, "timetable_summary": None, @@ -1426,6 +1430,28 @@ def test_get_dag_should_response_403(self, unauthorized_test_client): assert response.status_code == 403 +class TestSuggestedPartitionKey(TestDagEndpoint): + """Unit tests for the ``suggested_partition_key`` field of the Dag details endpoint.""" + + @time_machine.travel(datetime(2026, 2, 18, 9, 30, tzinfo=timezone.utc), tick=False) + def test_details_endpoint_suggests_last_elapsed_partition(self, session, test_client, dag_maker): + dag_id = "test_cron_partitioned_suggested_key" + with dag_maker( + dag_id=dag_id, + schedule=CronPartitionTimetable("0 0 * * *", timezone="UTC"), + session=session, + serialized=True, + ): + EmptyOperator(task_id=TASK_ID) + dag_maker.sync_dagbag_to_db() + session.commit() + + response = test_client.get(f"/dags/{dag_id}/details") + + assert response.status_code == 200 + assert response.json()["suggested_partition_key"] == "2026-02-18T00:00:00" + + class TestDagWithoutFileloc(TestDagEndpoint): def _make_dag_without_fileloc(self, dag_maker, session, dag_id="test_dag_no_fileloc"): with dag_maker(dag_id=dag_id, schedule=None): diff --git a/airflow-core/tests/unit/timetables/test_base_timetable.py b/airflow-core/tests/unit/timetables/test_base_timetable.py index 8eb3ca2871b33..b21fbc27566ea 100644 --- a/airflow-core/tests/unit/timetables/test_base_timetable.py +++ b/airflow-core/tests/unit/timetables/test_base_timetable.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import pendulum import pytest from airflow._shared.module_loading import qualname @@ -46,6 +47,13 @@ class CustomTimetable(Timetable): assert inst.type_name == expected +def test_suggest_partition_key_non_partitioned_returns_none(): + """Non-partitioned timetables never suggest a partition key.""" + + tt = NullTimetable() + assert tt.suggest_partition_key(pendulum.datetime(2026, 2, 18, tz="UTC")) is None + + # --------------------------------------------------------------------------- # compute_rollup_fingerprint # --------------------------------------------------------------------------- diff --git a/airflow-core/tests/unit/timetables/test_trigger_timetable.py b/airflow-core/tests/unit/timetables/test_trigger_timetable.py index f4718f8d04bb9..c253e188d7cd5 100644 --- a/airflow-core/tests/unit/timetables/test_trigger_timetable.py +++ b/airflow-core/tests/unit/timetables/test_trigger_timetable.py @@ -1081,6 +1081,27 @@ def test_iter_partition_dagrun_infos_inclusive_endpoint_pair() -> None: assert not any(i.partition_key == "2026-02-20T00:00:00" for i in infos_two) +def test_suggest_partition_key_after_tick_returns_most_recently_elapsed() -> None: + """A *now* strictly after today's tick suggests today's already-elapsed partition.""" + timetable = CoreCronPartitionTimetable("0 0 * * *", timezone="Asia/Taipei", run_offset=0) + now = pendulum.datetime(2026, 2, 18, 12, tz="Asia/Taipei") + assert timetable.suggest_partition_key(now) == "2026-02-18T00:00:00" + + +def test_suggest_partition_key_exactly_on_tick_returns_that_tick() -> None: + """A *now* landing exactly on a tick (the inclusive boundary) suggests that tick itself.""" + timetable = CoreCronPartitionTimetable("0 0 * * *", timezone="Asia/Taipei", run_offset=0) + now = pendulum.datetime(2026, 2, 18, tz="Asia/Taipei") + assert timetable.suggest_partition_key(now) == "2026-02-18T00:00:00" + + +def test_suggest_partition_key_reflects_run_offset() -> None: + """A non-zero run_offset shifts the suggested partition the same way it shifts scheduled runs.""" + timetable = CoreCronPartitionTimetable("0 0 * * *", timezone="Asia/Taipei", run_offset=-1) + now = pendulum.datetime(2026, 2, 18, 12, tz="Asia/Taipei") + assert timetable.suggest_partition_key(now) == "2026-02-17T00:00:00" + + def test_iter_partition_dagrun_infos_subday_window_does_not_expand_to_whole_day() -> None: """A sub-day window of an hourly timetable yields only the ticks inside it — not the whole day. diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index 3e9681b4c185f..4b9c6f886dc76 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -2655,6 +2655,7 @@ class DAGDetailsResponse(BaseModel): is_favorite: Annotated[bool | None, Field(title="Is Favorite")] = False active_runs_count: Annotated[int | None, Field(title="Active Runs Count")] = 0 team_name: Annotated[str | None, Field(title="Team Name")] = None + suggested_partition_key: Annotated[str | None, Field(title="Suggested Partition Key")] = None is_backfillable: Annotated[ bool, Field(description="Whether this Dag's schedule supports backfilling.", title="Is Backfillable") ]