From bb251ac2f6d9fbbc37819d5e45efbe5c827ab42f Mon Sep 17 00:00:00 2001 From: Wei Lee Date: Mon, 10 Aug 2026 17:12:30 +0800 Subject: [PATCH] Pre-populate a suggested partition key when manually triggering cron-partitioned Dags Manually triggering a Dag whose timetable is CronPartitionTimetable leaves the Partition Key field blank. If the user submits without typing one, the Dag run is created unpartitioned and any downstream asset-partition-aware Dag can never be satisfied by it. Add a new Timetable.suggest_partition_key(now) hook, overridden by CronPartitionTimetable to suggest the most recently elapsed partition tick. Thread the suggestion through a new suggested_partition_key field on DAGResponse to pre-populate the trigger form's Partition Key field. The field stays editable and still submits null if the user clears it. The Dags list view builds its payload by reading every DAGResponse field off a bare DagModel, so a field that only the single-Dag endpoints supply has to be declared as such or that endpoint fails with AttributeError. The trigger form applies the suggestion again whenever a newer one arrives, because a response served from cache can name a partition that has since elapsed. --- .../api_fastapi/core_api/datamodels/dags.py | 1 + .../openapi/v2-rest-api-generated.yaml | 5 + .../core_api/routes/public/dags.py | 4 + airflow-core/src/airflow/timetables/base.py | 29 ++++ .../src/airflow/timetables/trigger.py | 12 ++ .../ui/openapi-gen/requests/schemas.gen.ts | 11 ++ .../ui/openapi-gen/requests/types.gen.ts | 1 + .../TriggerDag/TriggerDAGForm.test.tsx | 147 ++++++++++++++++++ .../components/TriggerDag/TriggerDAGForm.tsx | 27 +++- .../components/TriggerDag/TriggerDAGModal.tsx | 8 +- .../core_api/routes/public/test_dags.py | 26 ++++ .../unit/timetables/test_base_timetable.py | 8 + .../unit/timetables/test_trigger_timetable.py | 21 +++ .../airflowctl/api/datamodels/generated.py | 1 + 14 files changed, 295 insertions(+), 6 deletions(-) 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") ]