diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index b6f7fbcf87933..44c822c0b1b04 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``76c46545c91e`` (head) | ``3c525f44bea8`` | ``3.4.0`` | Add new index for trigger. | +| ``b8f1c0a4d276`` (head) | ``76c46545c91e`` | ``3.4.0`` | Add timetable_partitioned_at_runtime to dag. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``76c46545c91e`` | ``3c525f44bea8`` | ``3.4.0`` | Add new index for trigger. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``3c525f44bea8`` | ``b2f1a9c7d4e0`` | ``3.4.0`` | Add indexes on serialized_dag and dag_code. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/src/airflow/api_fastapi/common/partition_helpers.py b/airflow-core/src/airflow/api_fastapi/common/partition_helpers.py index 88812a7014d21..1bdb99d5586a9 100644 --- a/airflow-core/src/airflow/api_fastapi/common/partition_helpers.py +++ b/airflow-core/src/airflow/api_fastapi/common/partition_helpers.py @@ -19,14 +19,21 @@ from typing import TYPE_CHECKING import structlog +from sqlalchemy import select from airflow.exceptions import DeserializationError +from airflow.models.asset import AssetPartitionDagRun +from airflow.models.dagrun import DagRun from airflow.models.serialized_dag import SerializedDagModel from airflow.timetables.simple import PartitionedAssetTimetable +from airflow.utils.state import DagRunState if TYPE_CHECKING: + from pendulum import DateTime from sqlalchemy.orm import Session + from airflow.timetables.base import Timetable + log = structlog.get_logger(logger_name=__name__) @@ -87,3 +94,73 @@ def load_partitioned_timetables( serdag.dag_id: _extract_partitioned_timetable(serdag) for serdag in SerializedDagModel.get_latest_serialized_dags(dag_ids=dag_ids, session=session) } + + +def suggest_partition_key_for_dag( + *, dag_id: str, timetable: Timetable, now: DateTime, session: Session +) -> str | None: + """ + Suggest a partition key to pre-fill the manual-trigger form for *dag_id*. + + This is a **guess** used only to pre-fill the UI form field; it is not a + validation and callers must still accept ``None`` from + ``validate_partition_key``. Returns ``None`` immediately for a timetable + that is neither ``partitioned`` nor ``partitioned_at_runtime`` (no query is + issued in that case). Otherwise, tries in order and returns the first + non-``None`` result: + + 1. The oldest pending ``AssetPartitionDagRun`` for *dag_id* + (``created_dag_run_id IS NULL``) — an asset event has already arrived + for this partition and is only waiting for the run to be created. Only + asset-driven timetables have any rows here. The ordering mirrors the + scheduler's FIFO claim order in + ``SchedulerJobRunner._create_partition_dag_runs``, so the suggestion + names the same partition the scheduler would create next rather than a + newer one that will not run until the backlog drains. + 2. ``timetable.suggest_partition_key(now)`` — a purely time-based guess + that needs no history, so even a brand-new Dag with no runs or asset + events gets a suggestion (asset-driven timetables only; see + :meth:`~airflow.timetables.base.Timetable.suggest_partition_key`). + 3. The ``partition_key`` of the most recent successful ``DagRun`` for + *dag_id*. This is the only source available to + ``partitioned_at_runtime`` timetables, which have no asset and no + temporal anchor of their own. Only the single most recent successful + run is consulted — if its ``partition_key`` is ``None`` (an + unpartitioned run), this returns ``None`` rather than searching + further back. + + Source 3 can name a partition that has already run, and nothing rejects a + second ``DagRun`` for the same ``(dag_id, partition_key)`` — re-running a + partition is a legitimate re-materialization, so this deliberately neither + de-duplicates nor blocks it. The trigger form's help text warns the user + that the pre-filled key may already have run. + + ``models.dag.get_last_dagrun`` is intentionally not reused here: it + filters ``logical_date.is_not(None)`` and excludes ``MANUAL`` runs by + default, which would drop exactly the partitioned runs this needs. + """ + if not timetable.partitioned and not timetable.partitioned_at_runtime: + return None + + pending_key = session.execute( + select(AssetPartitionDagRun.partition_key) + .where( + AssetPartitionDagRun.target_dag_id == dag_id, + AssetPartitionDagRun.created_dag_run_id.is_(None), + ) + .order_by(AssetPartitionDagRun.created_at, AssetPartitionDagRun.id) + .limit(1) + ).scalar_one_or_none() + if pending_key is not None: + return pending_key + + suggested_key = timetable.suggest_partition_key(now) + if suggested_key is not None: + return suggested_key + + return session.execute( + select(DagRun.partition_key) + .where(DagRun.dag_id == dag_id, DagRun.state == DagRunState.SUCCESS) + .order_by(DagRun.id.desc()) + .limit(1) + ).scalar_one_or_none() 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..debc413e0dcf0 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 @@ -71,6 +71,15 @@ def _get_file_token_serializer() -> URLSafeSerializer: "next_dagrun_run_after": "next_dagrun_create_after", } +DAG_RESPONSE_ROUTE_SUPPLIED_FIELDS: frozenset[str] = frozenset( + { + # Computed per request by the route rather than read off ``DagModel``, so + # routes that build the response by reading attributes off a bare + # ``DagModel`` must skip these instead of raising ``AttributeError``. + "suggested_partition_key", + } +) + class DAGResponse(BaseModel): """Dag serializer for responses.""" @@ -96,6 +105,7 @@ class DAGResponse(BaseModel): timetable_summary: str | None timetable_description: str | None timetable_partitioned: bool + timetable_partitioned_at_runtime: bool timetable_periodic: bool tags: list[DagTagResponse] max_active_tasks: int @@ -109,6 +119,7 @@ class DAGResponse(BaseModel): next_dagrun_run_after: datetime | None allowed_run_types: list[DagRunType] | None owners: list[str] + suggested_partition_key: str | None = None @field_serializer("tags") def serialize_tags(self, tags: list[DagTagResponse]) -> list[DagTagResponse]: diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml index 32f0f6f173511..00238daee6400 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/_private_ui.yaml @@ -2929,6 +2929,9 @@ components: timetable_partitioned: type: boolean title: Timetable Partitioned + timetable_partitioned_at_runtime: + type: boolean + title: Timetable Partitioned At Runtime timetable_periodic: type: boolean title: Timetable Periodic @@ -2990,6 +2993,11 @@ components: type: string type: array title: Owners + suggested_partition_key: + anyOf: + - type: string + - type: 'null' + title: Suggested Partition Key asset_expression: anyOf: - oneOf: @@ -3045,6 +3053,7 @@ components: - timetable_summary - timetable_description - timetable_partitioned + - timetable_partitioned_at_runtime - timetable_periodic - tags - max_active_tasks 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 b3ed8a3a2b83d..8dae37a10225a 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 @@ -12944,6 +12944,9 @@ components: timetable_partitioned: type: boolean title: Timetable Partitioned + timetable_partitioned_at_runtime: + type: boolean + title: Timetable Partitioned At Runtime timetable_periodic: type: boolean title: Timetable Periodic @@ -13005,6 +13008,11 @@ components: type: string type: array title: Owners + suggested_partition_key: + anyOf: + - type: string + - type: 'null' + title: Suggested Partition Key catchup: type: boolean title: Catchup @@ -13146,6 +13154,7 @@ components: - timetable_summary - timetable_description - timetable_partitioned + - timetable_partitioned_at_runtime - timetable_periodic - tags - max_active_tasks @@ -13258,6 +13267,9 @@ components: timetable_partitioned: type: boolean title: Timetable Partitioned + timetable_partitioned_at_runtime: + type: boolean + title: Timetable Partitioned At Runtime timetable_periodic: type: boolean title: Timetable Periodic @@ -13319,6 +13331,11 @@ components: type: string type: array title: Owners + suggested_partition_key: + anyOf: + - type: string + - type: 'null' + title: Suggested Partition Key is_backfillable: type: boolean title: Is Backfillable @@ -13346,6 +13363,7 @@ components: - timetable_summary - timetable_description - timetable_partitioned + - timetable_partitioned_at_runtime - timetable_periodic - tags - max_active_tasks 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..77b3f85c7f7fb 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 @@ -17,13 +17,14 @@ from __future__ import annotations -from typing import Annotated +from typing import TYPE_CHECKING, Annotated from fastapi import Depends, HTTPException, Query, Response, status from fastapi.exceptions import RequestValidationError from pydantic import ValidationError from sqlalchemy import delete, func, insert, select, update +from airflow._shared.timezones import timezone 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 @@ -56,6 +57,7 @@ datetime_range_filter_factory, filter_param_factory, ) +from airflow.api_fastapi.common.partition_helpers import suggest_partition_key_for_dag from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT from airflow.api_fastapi.core_api.datamodels.dags import ( @@ -79,9 +81,38 @@ from airflow.models.dagrun import DagRun from airflow.utils.state import DagRunState +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + from airflow.serialization.serialized_objects import SerializedDAG + + dags_router = AirflowRouter(tags=["DAG"], prefix="/dags") +def attach_suggested_partition_key(dag_model: DagModel, dag: SerializedDAG, *, session: Session) -> None: + """ + Attach the manual-trigger partition key suggestion to *dag_model* for serialization. + + ``suggested_partition_key`` is computed per request from live pending + partitions and run history, so it cannot be a ``DagModel`` column. Setting a + non-mapped attribute leaves the SQLAlchemy identity map untouched — nothing + is persisted by a later ``flush()``. Only single-Dag routes carry it; the + collection routes leave it at its ``None`` default rather than issue a query + per Dag. + """ + setattr( + dag_model, + "suggested_partition_key", + suggest_partition_key_for_dag( + dag_id=dag_model.dag_id, + timetable=dag.timetable, + now=timezone.coerce_datetime(timezone.utcnow()), + session=session, + ), + ) + + @dags_router.get("", dependencies=[Depends(requires_access_dag(method="GET"))]) def get_dags( limit: QueryLimit, @@ -209,6 +240,8 @@ def get_dag( if not key.startswith("_") and not hasattr(dag_model, key): setattr(dag_model, key, value) + attach_suggested_partition_key(dag_model, dag, session=session) + return dag_model @@ -236,6 +269,8 @@ def get_dag_details( if not key.startswith("_") and not hasattr(dag_model, key): setattr(dag_model, key, value) + attach_suggested_partition_key(dag_model, dag, session=session) + # Check if this Dag is marked as favorite by the current user user_id = str(user.get_id()) is_favorite = ( diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dags.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dags.py index 2700cb1e0b50b..dbeadbcfa2039 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dags.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/ui/dags.py @@ -58,7 +58,11 @@ filter_param_factory, ) from airflow.api_fastapi.common.router import AirflowRouter -from airflow.api_fastapi.core_api.datamodels.dags import DAG_ALIAS_MAPPING, DAGResponse +from airflow.api_fastapi.core_api.datamodels.dags import ( + DAG_ALIAS_MAPPING, + DAG_RESPONSE_ROUTE_SUPPLIED_FIELDS, + DAGResponse, +) from airflow.api_fastapi.core_api.datamodels.ui.dag_runs import DAGRunLightResponse from airflow.api_fastapi.core_api.datamodels.ui.dags import ( DAGRunStateCountsResponse, @@ -255,6 +259,9 @@ def get_dags( # aggregate rows by dag_id # Build the dict dynamically from DAGResponse.model_fields so that new fields # added to DAGResponse are picked up automatically without code changes here. + # ``getattr`` is deliberately left without a default so a field that is neither + # a ``DagModel`` attribute nor declared route-supplied fails loudly here + # instead of silently serializing as ``None``. dag_runs_by_dag_id: dict[str, DAGWithLatestDagRunsResponse] = {} for dag in dags: dag_data = { @@ -262,6 +269,7 @@ def get_dags( dag, DAG_ALIAS_MAPPING.get(field_name, field_name) ) for field_name in DAGResponse.model_fields + if field_name not in DAG_RESPONSE_ROUTE_SUPPLIED_FIELDS } dag_data.update( { diff --git a/airflow-core/src/airflow/cli/commands/dag_command.py b/airflow-core/src/airflow/cli/commands/dag_command.py index 3b5f450d0e262..9d8f031a1e16c 100644 --- a/airflow-core/src/airflow/cli/commands/dag_command.py +++ b/airflow-core/src/airflow/cli/commands/dag_command.py @@ -389,6 +389,7 @@ def _get_dagbag_dag_details(dag: DAG) -> dict: "timetable_summary": core_timetable.summary, "timetable_description": core_timetable.description, "timetable_partitioned": core_timetable.partitioned, + "timetable_partitioned_at_runtime": core_timetable.partitioned_at_runtime, "timetable_periodic": core_timetable.periodic, "tags": dag.tags, "max_active_tasks": dag.max_active_tasks, @@ -403,6 +404,8 @@ def _get_dagbag_dag_details(dag: DAG) -> dict: "next_dagrun_logical_date": None, "next_dagrun_run_after": None, "allowed_run_types": dag.allowed_run_types, + # No DB session here, so the trigger-form suggestion sources are unavailable. + "suggested_partition_key": None, "is_backfillable": core_timetable.periodic and (dag.allowed_run_types is None or DagRunType.BACKFILL_JOB in dag.allowed_run_types), } diff --git a/airflow-core/src/airflow/dag_processing/collection.py b/airflow-core/src/airflow/dag_processing/collection.py index 8fa5209ff15fb..f0a2c36c4f5eb 100644 --- a/airflow-core/src/airflow/dag_processing/collection.py +++ b/airflow-core/src/airflow/dag_processing/collection.py @@ -682,6 +682,7 @@ def update_dags( dm.timetable_summary = dag.timetable.summary dm.timetable_description = dag.timetable.description dm.timetable_partitioned = dag.timetable.partitioned + dm.timetable_partitioned_at_runtime = dag.timetable.partitioned_at_runtime dm.timetable_periodic = dag.timetable.periodic dm.partition_mapper_info = dag.timetable.partition_mapper_info dm.fail_fast = dag.fail_fast if dag.fail_fast is not None else False diff --git a/airflow-core/src/airflow/migrations/versions/0131_3_4_0_add_timetable_partitioned_at_runtime_.py b/airflow-core/src/airflow/migrations/versions/0131_3_4_0_add_timetable_partitioned_at_runtime_.py new file mode 100644 index 0000000000000..ba25c5dbc91b4 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0131_3_4_0_add_timetable_partitioned_at_runtime_.py @@ -0,0 +1,52 @@ +# +# 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. + +""" +Add timetable_partitioned_at_runtime to dag. + +Revision ID: b8f1c0a4d276 +Revises: 76c46545c91e +Create Date: 2026-08-12 18:40:11.402913 + +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "b8f1c0a4d276" +down_revision = "76c46545c91e" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + + +def upgrade(): + """Add timetable_partitioned_at_runtime to dag.""" + with op.batch_alter_table("dag", schema=None) as batch_op: + batch_op.add_column( + sa.Column("timetable_partitioned_at_runtime", sa.Boolean, nullable=False, server_default="0") + ) + + +def downgrade(): + """Remove timetable_partitioned_at_runtime from dag.""" + with op.batch_alter_table("dag", schema=None) as batch_op: + batch_op.drop_column("timetable_partitioned_at_runtime") diff --git a/airflow-core/src/airflow/models/dag.py b/airflow-core/src/airflow/models/dag.py index 4c8d4d3e680dd..5c8cdecbfcf84 100644 --- a/airflow-core/src/airflow/models/dag.py +++ b/airflow-core/src/airflow/models/dag.py @@ -358,6 +358,10 @@ class DagModel(Base): timetable_description: Mapped[str | None] = mapped_column(String(1000), nullable=True) # Whether the timetable do partitioning. timetable_partitioned: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default="0") + # Whether partition keys are supplied at trigger time instead of derived from a schedule. + timetable_partitioned_at_runtime: Mapped[bool] = mapped_column( + Boolean, nullable=False, server_default="0" + ) # Whether the timetable is periodic (supports backfilling). timetable_periodic: Mapped[bool] = mapped_column(Boolean, nullable=False, server_default="0") # Cached partition mapper metadata for partitioned timetables, populated diff --git a/airflow-core/src/airflow/timetables/base.py b/airflow-core/src/airflow/timetables/base.py index 365f980b93274..2b76bbe46f3b6 100644 --- a/airflow-core/src/airflow/timetables/base.py +++ b/airflow-core/src/airflow/timetables/base.py @@ -338,6 +338,28 @@ 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 manually triggering this Dag at *now*. + + This is a **suggestion, not a derivation**: it is used to pre-fill the + partition key field in the manual-trigger UI so the field is not left + empty, but the caller is not required to use the returned value and + validation still accepts ``None``. Returning ``None`` means this + timetable has no purely time-based suggestion to offer; the caller + should leave the field blank or fall back to another source. + + This method must not access the database or open a session — it runs + on a deserialized timetable instance shared across the scheduler, Dag + processor, and API layers. Sources that require a database query + (e.g. the most recent successful run's partition key) belong in the + API layer, not here. + + :param now: The current time, used to compute a candidate key. + :returns: A suggested partition key, or ``None`` if none can be derived. + """ + 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 87741a8eb0d7e..9b820b52082d5 100644 --- a/airflow-core/src/airflow/timetables/simple.py +++ b/airflow-core/src/airflow/timetables/simple.py @@ -397,6 +397,51 @@ def _decode_partition_date(self, partition_key: str) -> datetime | None: return anchors.pop() return None + def suggest_partition_key(self, now: DateTime) -> str | None: + """ + Suggest a partition key for the period containing *now*, purely from the asset mappers. + + This is a **guess**, not a derivation: it asks every asset (and asset + ref) reachable from ``asset_condition`` for the downstream key of the + period that contains *now* — ``mapper.format(mapper.normalize(now))`` — + and returns it only when every mapper that can answer agrees on the same + key. It does not consult any actual asset events or ``DagRun`` history, + so it can suggest a value even for a Dag that has never run. + + Only mappers exposing both a callable ``normalize`` and ``format`` + (duck-typed, not ``isinstance``) can answer; enumerated against the + current ``airflow.partition_mappers`` set: + + - Can answer: the ``_BaseTemporalMapper`` family — ``StartOfHourMapper``, + ``StartOfDayMapper``, ``StartOfWeekMapper``, ``StartOfMonthMapper``, + ``StartOfQuarterMapper``, ``StartOfYearMapper``. + - Cannot answer (no purely time-based guess to offer): ``RollupMapper``, + ``FanOutMapper``, ``ChainMapper``, ``ProductMapper``, ``IdentityMapper``, + ``FixedKeyMapper``, ``AllowedKeyMapper``. + + Zero mappers able to answer, or answers that disagree, both return + ``None`` — same convention as :meth:`_decode_partition_date`. + """ + keys: set[str] = set() + for unique_key, _ in self.asset_condition.iter_assets(): + mapper = self.get_partition_mapper(name=unique_key.name, uri=unique_key.uri) + key = _suggest_partition_key_with(mapper, now) + if key is not None: + keys.add(key) + 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 + key = _suggest_partition_key_with(mapper, now) + if key is not None: + keys.add(key) + if len(keys) == 1: + return keys.pop() + return None + def serialize(self) -> dict[str, Any]: from airflow.serialization.serialized_objects import encode_asset_like @@ -426,3 +471,30 @@ def deserialize(cls, data: dict[str, Any]) -> PartitionedAssetTimetable: }, ) return timetable + + +def _suggest_partition_key_with(mapper: PartitionMapper, now: DateTime) -> str | None: + """ + Return *mapper*'s guess for the downstream key of the period containing *now*, or ``None``. + + Duck-types on ``normalize``/``format`` (mirrors the ``format`` check in + :func:`~airflow.partition_mappers.temporal._format_with`) rather than + ``isinstance``, so any mapper — builtin or third-party — that exposes both + callables participates. A mapper that raises (e.g. a misconfigured format) + must not turn a manual-trigger modal into a 500: the exception is logged + and swallowed, same handling as + :func:`~airflow.api_fastapi.core_api.routes.ui.partitioned_dag_runs._resolve_rollup_status`. + """ + normalize = getattr(mapper, "normalize", None) + format_ = getattr(mapper, "format", None) + if not callable(normalize) or not callable(format_): + return None + try: + return format_(normalize(now)) + except Exception: + log.warning( + "Failed to suggest partition key from mapper; ignoring", + mapper=type(mapper).__name__, + exc_info=True, + ) + return None 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 1a1f29b09515c..d9e152704167f 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 @@ -2951,6 +2951,10 @@ export const $DAGDetailsResponse = { type: 'boolean', title: 'Timetable Partitioned' }, + timetable_partitioned_at_runtime: { + type: 'boolean', + title: 'Timetable Partitioned At Runtime' + }, timetable_periodic: { type: 'boolean', title: 'Timetable Periodic' @@ -3058,6 +3062,17 @@ export const $DAGDetailsResponse = { type: 'array', title: 'Owners' }, + suggested_partition_key: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Suggested Partition Key' + }, catchup: { type: 'boolean', title: 'Catchup' @@ -3293,7 +3308,7 @@ Deprecated: Use max_active_tasks instead.`, } }, type: 'object', - required: ['dag_id', 'dag_display_name', 'is_paused', 'is_stale', 'last_parsed_time', 'last_parse_duration', 'last_expired', 'bundle_name', 'bundle_version', 'relative_fileloc', 'fileloc', 'description', 'timetable_summary', 'timetable_description', 'timetable_partitioned', 'timetable_periodic', 'tags', 'max_active_tasks', 'max_active_runs', 'max_consecutive_failed_dag_runs', 'has_task_concurrency_limits', 'has_import_errors', 'next_dagrun_logical_date', 'next_dagrun_data_interval_start', 'next_dagrun_data_interval_end', 'next_dagrun_run_after', 'allowed_run_types', 'owners', 'catchup', 'dag_run_timeout', 'asset_expression', 'doc_md', 'start_date', 'end_date', 'is_paused_upon_creation', 'params', 'render_template_as_native_obj', 'template_search_path', 'timezone', 'last_parsed', 'default_args', 'is_backfillable', 'file_token', 'concurrency', 'latest_dag_version'], + required: ['dag_id', 'dag_display_name', 'is_paused', 'is_stale', 'last_parsed_time', 'last_parse_duration', 'last_expired', 'bundle_name', 'bundle_version', 'relative_fileloc', 'fileloc', 'description', 'timetable_summary', 'timetable_description', 'timetable_partitioned', 'timetable_partitioned_at_runtime', 'timetable_periodic', 'tags', 'max_active_tasks', 'max_active_runs', 'max_consecutive_failed_dag_runs', 'has_task_concurrency_limits', 'has_import_errors', 'next_dagrun_logical_date', 'next_dagrun_data_interval_start', 'next_dagrun_data_interval_end', 'next_dagrun_run_after', 'allowed_run_types', 'owners', 'catchup', 'dag_run_timeout', 'asset_expression', 'doc_md', 'start_date', 'end_date', 'is_paused_upon_creation', 'params', 'render_template_as_native_obj', 'template_search_path', 'timezone', 'last_parsed', 'default_args', 'is_backfillable', 'file_token', 'concurrency', 'latest_dag_version'], title: 'DAGDetailsResponse', description: 'Specific serializer for Dag Details responses.' } as const; @@ -3446,6 +3461,10 @@ export const $DAGResponse = { type: 'boolean', title: 'Timetable Partitioned' }, + timetable_partitioned_at_runtime: { + type: 'boolean', + title: 'Timetable Partitioned At Runtime' + }, timetable_periodic: { type: 'boolean', title: 'Timetable Periodic' @@ -3553,6 +3572,17 @@ export const $DAGResponse = { type: 'array', title: 'Owners' }, + suggested_partition_key: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Suggested Partition Key' + }, is_backfillable: { type: 'boolean', title: 'Is Backfillable', @@ -3567,7 +3597,7 @@ export const $DAGResponse = { } }, type: 'object', - required: ['dag_id', 'dag_display_name', 'is_paused', 'is_stale', 'last_parsed_time', 'last_parse_duration', 'last_expired', 'bundle_name', 'bundle_version', 'relative_fileloc', 'fileloc', 'description', 'timetable_summary', 'timetable_description', 'timetable_partitioned', 'timetable_periodic', 'tags', 'max_active_tasks', 'max_active_runs', 'max_consecutive_failed_dag_runs', 'has_task_concurrency_limits', 'has_import_errors', 'next_dagrun_logical_date', 'next_dagrun_data_interval_start', 'next_dagrun_data_interval_end', 'next_dagrun_run_after', 'allowed_run_types', 'owners', 'is_backfillable', 'file_token'], + required: ['dag_id', 'dag_display_name', 'is_paused', 'is_stale', 'last_parsed_time', 'last_parse_duration', 'last_expired', 'bundle_name', 'bundle_version', 'relative_fileloc', 'fileloc', 'description', 'timetable_summary', 'timetable_description', 'timetable_partitioned', 'timetable_partitioned_at_runtime', 'timetable_periodic', 'tags', 'max_active_tasks', 'max_active_runs', 'max_consecutive_failed_dag_runs', 'has_task_concurrency_limits', 'has_import_errors', 'next_dagrun_logical_date', 'next_dagrun_data_interval_start', 'next_dagrun_data_interval_end', 'next_dagrun_run_after', 'allowed_run_types', 'owners', 'is_backfillable', 'file_token'], title: 'DAGResponse', description: 'Dag serializer for responses.' } as const; @@ -9335,6 +9365,10 @@ export const $DAGWithLatestDagRunsResponse = { type: 'boolean', title: 'Timetable Partitioned' }, + timetable_partitioned_at_runtime: { + type: 'boolean', + title: 'Timetable Partitioned At Runtime' + }, timetable_periodic: { type: 'boolean', title: 'Timetable Periodic' @@ -9442,6 +9476,17 @@ export const $DAGWithLatestDagRunsResponse = { type: 'array', title: 'Owners' }, + suggested_partition_key: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Suggested Partition Key' + }, asset_expression: { anyOf: [ { @@ -9512,7 +9557,7 @@ export const $DAGWithLatestDagRunsResponse = { } }, type: 'object', - required: ['dag_id', 'dag_display_name', 'is_paused', 'is_stale', 'last_parsed_time', 'last_parse_duration', 'last_expired', 'bundle_name', 'bundle_version', 'relative_fileloc', 'fileloc', 'description', 'timetable_summary', 'timetable_description', 'timetable_partitioned', 'timetable_periodic', 'tags', 'max_active_tasks', 'max_active_runs', 'max_consecutive_failed_dag_runs', 'has_task_concurrency_limits', 'has_import_errors', 'next_dagrun_logical_date', 'next_dagrun_data_interval_start', 'next_dagrun_data_interval_end', 'next_dagrun_run_after', 'allowed_run_types', 'owners', 'asset_expression', 'latest_dag_runs', 'pending_actions', 'is_favorite', 'is_backfillable', 'file_token'], + required: ['dag_id', 'dag_display_name', 'is_paused', 'is_stale', 'last_parsed_time', 'last_parse_duration', 'last_expired', 'bundle_name', 'bundle_version', 'relative_fileloc', 'fileloc', 'description', 'timetable_summary', 'timetable_description', 'timetable_partitioned', 'timetable_partitioned_at_runtime', 'timetable_periodic', 'tags', 'max_active_tasks', 'max_active_runs', 'max_consecutive_failed_dag_runs', 'has_task_concurrency_limits', 'has_import_errors', 'next_dagrun_logical_date', 'next_dagrun_data_interval_start', 'next_dagrun_data_interval_end', 'next_dagrun_run_after', 'allowed_run_types', 'owners', 'asset_expression', 'latest_dag_runs', 'pending_actions', 'is_favorite', 'is_backfillable', 'file_token'], title: 'DAGWithLatestDagRunsResponse', description: 'DAG with latest dag runs response serializer.' } as const; 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 92b30f7719982..fc33961e97b9f 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 @@ -865,6 +865,7 @@ export type DAGDetailsResponse = { timetable_summary: string | null; timetable_description: string | null; timetable_partitioned: boolean; + timetable_partitioned_at_runtime: boolean; timetable_periodic: boolean; tags: Array; max_active_tasks: number; @@ -878,6 +879,7 @@ export type DAGDetailsResponse = { next_dagrun_run_after: string | null; allowed_run_types: Array | null; owners: Array<(string)>; + suggested_partition_key?: string | null; catchup: boolean; dag_run_timeout: string | null; asset_expression: AssetExpressionAsset | AssetExpressionAlias | AssetExpressionRef | AssetExpressionAny | AssetExpressionAll | null; @@ -949,6 +951,7 @@ export type DAGResponse = { timetable_summary: string | null; timetable_description: string | null; timetable_partitioned: boolean; + timetable_partitioned_at_runtime: boolean; timetable_periodic: boolean; tags: Array; max_active_tasks: number; @@ -962,6 +965,7 @@ export type DAGResponse = { next_dagrun_run_after: string | null; allowed_run_types: Array | null; owners: Array<(string)>; + suggested_partition_key?: string | null; /** * Whether this Dag's schedule supports backfilling. */ @@ -2366,6 +2370,7 @@ export type DAGWithLatestDagRunsResponse = { timetable_summary: string | null; timetable_description: string | null; timetable_partitioned: boolean; + timetable_partitioned_at_runtime: boolean; timetable_periodic: boolean; tags: Array; max_active_tasks: number; @@ -2379,6 +2384,7 @@ export type DAGWithLatestDagRunsResponse = { next_dagrun_run_after: string | null; allowed_run_types: Array | null; owners: Array<(string)>; + suggested_partition_key?: string | null; asset_expression: AssetExpressionAsset | AssetExpressionAlias | AssetExpressionRef | AssetExpressionAny | AssetExpressionAll | null; latest_dag_runs: Array; pending_actions: Array; diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/en/components.json b/airflow-core/src/airflow/ui/public/i18n/locales/en/components.json index 08490d89f8d1b..30e0a21b1c456 100644 --- a/airflow-core/src/airflow/ui/public/i18n/locales/en/components.json +++ b/airflow-core/src/airflow/ui/public/i18n/locales/en/components.json @@ -149,6 +149,7 @@ "loadingFailed": "Failed to load Dag information. Please try again.", "manualRunDenied": "Manual runs are not allowed for this Dag", "partitionKeyHelp": "Optional - only applies to partitioned Dags", + "partitionKeySuggestedHelp": "Suggested value - this partition may already have run. Confirm it is the one you want, or edit it, before submitting.", "runIdHelp": "Optional - will be generated if not provided", "selectDescription": "Trigger a single run of this Dag", "selectLabel": "Single Run", diff --git a/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/components.json b/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/components.json index 54488c9372762..6cd96cd3aa512 100644 --- a/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/components.json +++ b/airflow-core/src/airflow/ui/public/i18n/locales/zh-TW/components.json @@ -149,6 +149,7 @@ "loadingFailed": "載入 Dag 資訊失敗，請重試。", "manualRunDenied": "此 Dag 不允許手動執行", "partitionKeyHelp": "選填 - 僅適用於分區的 Dag", + "partitionKeySuggestedHelp": "此為建議值，該分區可能已經執行過 - 送出前請確認是否為您要的分區，或自行修改。", "runIdHelp": "選填 - 若未提供將會自動產生", "selectDescription": "觸發此 Dag 單次執行", "selectLabel": "單次執行", diff --git a/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGAdvancedOptions.tsx b/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGAdvancedOptions.tsx index d5512579093f0..09f1db9b7607e 100644 --- a/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGAdvancedOptions.tsx +++ b/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGAdvancedOptions.tsx @@ -26,12 +26,29 @@ import type { DagRunTriggerParams } from "./types"; type TriggerDAGAdvancedOptionsProps = { readonly control: Control; readonly isPartitioned: boolean; + readonly isPartitionedAtRuntime: boolean; + readonly suggestedPartitionKey?: string | null; }; -const TriggerDAGAdvancedOptions = ({ control, isPartitioned }: TriggerDAGAdvancedOptionsProps) => { +const TriggerDAGAdvancedOptions = ({ + control, + isPartitioned, + isPartitionedAtRuntime, + suggestedPartitionKey, +}: TriggerDAGAdvancedOptionsProps) => { const { t: translate } = useTranslation(["common", "components"]); const { t: rootTranslate } = useTranslation(); + // The suggestion is a guess from whichever source answered first — a pending + // partition, a time-based inference, or the last successful run — and the API + // does not say which. One warning that holds for all three is honest; picking + // wording from isPartitionedAtRuntime would mislabel an asset-driven Dag that + // fell through to its last successful run. + const partitionKeyHelpText = + suggestedPartitionKey === undefined || suggestedPartitionKey === null + ? translate("components:triggerDag.partitionKeyHelp") + : translate("components:triggerDag.partitionKeySuggestedHelp"); + return ( <> - {isPartitioned ? ( + {isPartitioned || isPartitionedAtRuntime ? ( - {translate("components:triggerDag.partitionKeyHelp")} + {partitionKeyHelpText} )} 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..fa58095209927 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 @@ -99,6 +99,7 @@ describe("TriggerDAGForm", () => { error={undefined} hasSchedule={false} isPartitioned={false} + isPartitionedAtRuntime={false} isPaused={false} isPending={false} onSubmitTrigger={vi.fn()} @@ -133,6 +134,7 @@ describe("TriggerDAGForm", () => { error={undefined} hasSchedule={false} isPartitioned={false} + isPartitionedAtRuntime={false} isPaused={false} isPending={false} onSubmitTrigger={vi.fn()} @@ -179,6 +181,7 @@ describe("TriggerDAGForm", () => { error={undefined} hasSchedule={false} isPartitioned={false} + isPartitionedAtRuntime={false} isPaused={false} isPending={false} onSubmitTrigger={vi.fn()} @@ -201,6 +204,7 @@ describe("TriggerDAGForm", () => { error={undefined} hasSchedule={false} isPartitioned + isPartitionedAtRuntime={false} isPaused={false} isPending={false} onSubmitTrigger={vi.fn()} @@ -214,4 +218,105 @@ describe("TriggerDAGForm", () => { await waitFor(() => expect(screen.getByText("dagRun.partitionKey")).toBeInTheDocument()); expect(screen.getByText("components:triggerDag.partitionKeyHelp")).toBeInTheDocument(); }); + + it("shows the partition key field for partitioned_at_runtime Dags even though isPartitioned is false", async () => { + render( + , + { wrapper: Wrapper }, + ); + + fireEvent.click(screen.getByText("Advanced Options")); + + await waitFor(() => expect(screen.getByText("dagRun.partitionKey")).toBeInTheDocument()); + expect(screen.getByText("components:triggerDag.partitionKeyHelp")).toBeInTheDocument(); + }); + + it("pre-fills the suggested partition key for an asset-driven Dag and warns the value is a suggestion", async () => { + render( + , + { wrapper: Wrapper }, + ); + + fireEvent.click(screen.getByText("Advanced Options")); + + await waitFor(() => expect(screen.getByDisplayValue("2024-01-01")).toBeInTheDocument()); + expect(screen.getByText("components:triggerDag.partitionKeySuggestedHelp")).toBeInTheDocument(); + }); + + it("pre-fills the suggested partition key for a partitioned_at_runtime Dag and warns the value is a suggestion", async () => { + render( + , + { wrapper: Wrapper }, + ); + + fireEvent.click(screen.getByText("Advanced Options")); + + await waitFor(() => expect(screen.getByDisplayValue("2024-01-01")).toBeInTheDocument()); + expect(screen.getByText("components:triggerDag.partitionKeySuggestedHelp")).toBeInTheDocument(); + }); + + it("applies a suggested partition key that arrives after the form first rendered", async () => { + const form = (suggestedPartitionKey?: string) => ( + + ); + + // The Dag query resolves after mount, so defaultValues captured no suggestion; + // only the prefill effect can still apply it. + const { rerender } = render(form(undefined), { wrapper: Wrapper }); + + fireEvent.click(screen.getByText("Advanced Options")); + await waitFor(() => expect(screen.getByText("dagRun.partitionKey")).toBeInTheDocument()); + + rerender(form("2024-01-01")); + + await waitFor(() => expect(screen.getByDisplayValue("2024-01-01")).toBeInTheDocument()); + }); }); 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..aa209e0b8874c 100644 --- a/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGForm.tsx +++ b/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGForm.tsx @@ -42,6 +42,7 @@ type TriggerDAGFormProps = { readonly error?: unknown; readonly hasSchedule: boolean; readonly isPartitioned: boolean; + readonly isPartitionedAtRuntime: boolean; readonly isPaused: boolean; readonly isPending?: boolean; readonly onSubmitTrigger?: (params: DagRunTriggerParams) => void; @@ -53,6 +54,7 @@ type TriggerDAGFormProps = { runId: string; } | undefined; + readonly suggestedPartitionKey?: string | null; }; const TriggerDAGForm = ({ @@ -61,11 +63,13 @@ const TriggerDAGForm = ({ error, hasSchedule, isPartitioned, + isPartitionedAtRuntime, isPaused, isPending = false, onSubmitTrigger, open, prefillConfig, + suggestedPartitionKey, }: TriggerDAGFormProps) => { const { t: translate } = useTranslation(["common", "components"]); const [errors, setErrors] = useState<{ conf?: string; date?: unknown }>({}); @@ -87,12 +91,13 @@ const TriggerDAGForm = ({ // For partitioned Dags, logical date is not applicable. logicalDate: isPartitioned ? "" : dayjs().format(DEFAULT_DATETIME_FORMAT), note: "", - partitionKey: undefined, + partitionKey: suggestedPartitionKey ?? undefined, }, }); // 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), not logicalDate or runId, to avoid 409 conflicts. + // partitionKey still gets the suggested value (a guess, not a restore of the prior attempt). useEffect(() => { if (prefillConfig && open) { const confString = prefillConfig.conf ? JSON.stringify(prefillConfig.conf, undefined, 2) : ""; @@ -105,7 +110,7 @@ const TriggerDAGForm = ({ dataIntervalStart: "", logicalDate: isPartitioned ? "" : dayjs().format(DEFAULT_DATETIME_FORMAT), note: "", - partitionKey: undefined, + partitionKey: suggestedPartitionKey ?? undefined, }); // 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,6 +137,7 @@ const TriggerDAGForm = ({ initialParamDict, setInitialParamDict, isPartitioned, + suggestedPartitionKey, ]); // Automatically reset form when conf is fetched (only if no prefillConfig) @@ -248,7 +254,12 @@ const TriggerDAGForm = ({ setErrors={setErrors} setFormError={setFormError} > - + 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..4e65a6d590155 100644 --- a/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGModal.tsx +++ b/airflow-core/src/airflow/ui/src/components/TriggerDag/TriggerDAGModal.tsx @@ -75,6 +75,8 @@ const TriggerDAGModal: React.FC = ({ const isBackfillable = dag?.is_backfillable ?? false; const hasSchedule = dag?.timetable_summary !== null; const isPartitioned = dag ? dag.timetable_partitioned : false; + const isPartitionedAtRuntime = dag?.timetable_partitioned_at_runtime ?? false; + const suggestedPartitionKey = dag?.suggested_partition_key; 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; @@ -144,11 +146,13 @@ const TriggerDAGModal: React.FC = ({ error={error} hasSchedule={hasSchedule} isPartitioned={isPartitioned} + isPartitionedAtRuntime={isPartitionedAtRuntime} isPaused={isPaused} isPending={isPending} onSubmitTrigger={triggerDagRun} open={open} prefillConfig={prefillConfig} + suggestedPartitionKey={suggestedPartitionKey} /> ) : ( isBackfillable && dag && diff --git a/airflow-core/src/airflow/ui/src/pages/Asset/CreateAssetEventModal.tsx b/airflow-core/src/airflow/ui/src/pages/Asset/CreateAssetEventModal.tsx index 76a291d4667f5..268cdaa9bb195 100644 --- a/airflow-core/src/airflow/ui/src/pages/Asset/CreateAssetEventModal.tsx +++ b/airflow-core/src/airflow/ui/src/pages/Asset/CreateAssetEventModal.tsx @@ -230,10 +230,12 @@ export const CreateAssetEventModal = ({ asset, onClose, open }: Props) => { error={materializeError} hasSchedule={dag.timetable_summary !== null} isPartitioned={dag.timetable_partitioned} + isPartitionedAtRuntime={dag.timetable_partitioned_at_runtime} isPaused={dag.is_paused} isPending={isMaterializePending} onSubmitTrigger={handleMaterializeSubmit} open={open} + suggestedPartitionKey={dag.suggested_partition_key} /> ) : undefined} diff --git a/airflow-core/src/airflow/ui/src/pages/DagsList/DagCard.test.tsx b/airflow-core/src/airflow/ui/src/pages/DagsList/DagCard.test.tsx index 04064406576e8..be87b30c2f94c 100644 --- a/airflow-core/src/airflow/ui/src/pages/DagsList/DagCard.test.tsx +++ b/airflow-core/src/airflow/ui/src/pages/DagsList/DagCard.test.tsx @@ -189,6 +189,7 @@ const mockDag = { tags: [], timetable_description: "Every minute", timetable_partitioned: false, + timetable_partitioned_at_runtime: false, timetable_periodic: true, timetable_summary: "* * * * *", } satisfies DAGWithLatestDagRunsResponse; diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 4b873b4214eba..2a10b3895fb02 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "76c46545c91e", + "3.4.0": "b8f1c0a4d276", } # Prefix used to identify tables holding data moved during migration. diff --git a/airflow-core/tests/unit/api_fastapi/common/test_partition_helpers.py b/airflow-core/tests/unit/api_fastapi/common/test_partition_helpers.py index cdadeaf62752f..ba6e6b552f260 100644 --- a/airflow-core/tests/unit/api_fastapi/common/test_partition_helpers.py +++ b/airflow-core/tests/unit/api_fastapi/common/test_partition_helpers.py @@ -18,10 +18,23 @@ from unittest import mock +import pendulum import pytest -from airflow.api_fastapi.common.partition_helpers import _extract_partitioned_timetable +from airflow.api_fastapi.common.partition_helpers import ( + _extract_partitioned_timetable, + suggest_partition_key_for_dag, +) from airflow.exceptions import DeserializationError +from airflow.models.asset import AssetPartitionDagRun +from airflow.partition_mappers.temporal import StartOfDayMapper +from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.sdk import Asset +from airflow.serialization.encoders import ensure_serialized_asset +from airflow.timetables.simple import NullTimetable, PartitionedAssetTimetable, PartitionedAtRuntime +from airflow.utils.state import DagRunState + +NOW = pendulum.datetime(2025, 6, 1, tz="UTC") def _make_serdag(exc: Exception): @@ -95,3 +108,186 @@ def test_extract_partitioned_timetable_refactor_signal_exceptions_propagate(exc) _extract_partitioned_timetable(serdag) mock_log.warning.assert_not_called() + + +def _asset_driven_timetable_with_default_mapper(asset_uri: str) -> PartitionedAssetTimetable: + """ + Asset-driven timetable using the real default mapper (``IdentityMapper``). + + ``IdentityMapper`` has no ``normalize``/``format``, so + ``suggest_partition_key`` always returns ``None`` for it — exercising the + resolver's fallback to the recent-run source (step 3) without a temporal + mapper (step 2) short-circuiting first. + """ + return PartitionedAssetTimetable(assets=ensure_serialized_asset(Asset(name=asset_uri, uri=asset_uri))) + + +@pytest.mark.db_test +class TestSuggestPartitionKeyForDag: + @pytest.mark.parametrize( + "timetable_kind", + ["partitioned_at_runtime", "asset_driven"], + ) + def test_returns_partition_key_of_most_recent_successful_run(self, dag_maker, session, timetable_kind): + dag_id = f"suggest_pk_recent_{timetable_kind}" + timetable = ( + PartitionedAtRuntime() + if timetable_kind == "partitioned_at_runtime" + else _asset_driven_timetable_with_default_mapper(f"s3://bucket/{dag_id}") + ) + with dag_maker(dag_id=dag_id, schedule=timetable, serialized=True, session=session): + EmptyOperator(task_id="t") + dag_maker.create_dagrun( + run_id="older", + state=DagRunState.SUCCESS, + partition_key="2024-01-01", + logical_date=pendulum.datetime(2024, 1, 1, tz="UTC"), + ) + dag_maker.create_dagrun( + run_id="newer", + state=DagRunState.SUCCESS, + partition_key="2024-01-02", + logical_date=pendulum.datetime(2024, 1, 2, tz="UTC"), + ) + session.commit() + + result = suggest_partition_key_for_dag(dag_id=dag_id, timetable=timetable, now=NOW, session=session) + assert result == "2024-01-02" + + @pytest.mark.parametrize( + "timetable_kind", + ["partitioned_at_runtime", "asset_driven"], + ) + def test_returns_none_when_dag_never_ran(self, dag_maker, session, timetable_kind): + dag_id = f"suggest_pk_never_ran_{timetable_kind}" + timetable = ( + PartitionedAtRuntime() + if timetable_kind == "partitioned_at_runtime" + else _asset_driven_timetable_with_default_mapper(f"s3://bucket/{dag_id}") + ) + with dag_maker(dag_id=dag_id, schedule=timetable, serialized=True, session=session): + EmptyOperator(task_id="t") + session.commit() + + result = suggest_partition_key_for_dag(dag_id=dag_id, timetable=timetable, now=NOW, session=session) + assert result is None + + @pytest.mark.parametrize( + "timetable_kind", + ["partitioned_at_runtime", "asset_driven"], + ) + def test_returns_none_when_most_recent_successful_run_is_unpartitioned( + self, dag_maker, session, timetable_kind + ): + dag_id = f"suggest_pk_unpartitioned_{timetable_kind}" + timetable = ( + PartitionedAtRuntime() + if timetable_kind == "partitioned_at_runtime" + else _asset_driven_timetable_with_default_mapper(f"s3://bucket/{dag_id}") + ) + with dag_maker(dag_id=dag_id, schedule=timetable, serialized=True, session=session): + EmptyOperator(task_id="t") + dag_maker.create_dagrun( + run_id="older", + state=DagRunState.SUCCESS, + partition_key="2024-01-01", + logical_date=pendulum.datetime(2024, 1, 1, tz="UTC"), + ) + dag_maker.create_dagrun( + run_id="newer", + state=DagRunState.SUCCESS, + partition_key=None, + logical_date=pendulum.datetime(2024, 1, 2, tz="UTC"), + ) + session.commit() + + result = suggest_partition_key_for_dag(dag_id=dag_id, timetable=timetable, now=NOW, session=session) + assert result is None + + def test_pending_apdr_takes_priority_over_recent_run(self, dag_maker, session): + dag_id = "suggest_pk_pending_apdr_dag" + timetable = _asset_driven_timetable_with_default_mapper(f"s3://bucket/{dag_id}") + with dag_maker(dag_id=dag_id, schedule=timetable, serialized=True, session=session): + EmptyOperator(task_id="t") + dag_maker.create_dagrun(run_id="older", state=DagRunState.SUCCESS, partition_key="2024-01-01") + session.add(AssetPartitionDagRun(target_dag_id=dag_id, partition_key="pending-key")) + session.commit() + + result = suggest_partition_key_for_dag(dag_id=dag_id, timetable=timetable, now=NOW, session=session) + assert result == "pending-key" + + def test_falls_back_to_mapper_derived_key_when_no_apdr_or_run(self, dag_maker, session): + dag_id = "suggest_pk_mapper_derived_dag" + timetable = PartitionedAssetTimetable( + assets=Asset(name=f"s3://bucket/{dag_id}", uri=f"s3://bucket/{dag_id}"), + default_partition_mapper=StartOfDayMapper(), + ) + with dag_maker(dag_id=dag_id, schedule=timetable, serialized=True, session=session): + EmptyOperator(task_id="t") + session.commit() + + result = suggest_partition_key_for_dag(dag_id=dag_id, timetable=timetable, now=NOW, session=session) + assert result == "2025-06-01" + + def test_pending_apdr_takes_priority_over_mapper_derived_key(self, dag_maker, session): + dag_id = "suggest_pk_pending_beats_mapper_dag" + timetable = PartitionedAssetTimetable( + assets=Asset(name=f"s3://bucket/{dag_id}", uri=f"s3://bucket/{dag_id}"), + default_partition_mapper=StartOfDayMapper(), + ) + with dag_maker(dag_id=dag_id, schedule=timetable, serialized=True, session=session): + EmptyOperator(task_id="t") + session.add(AssetPartitionDagRun(target_dag_id=dag_id, partition_key="pending-key")) + session.commit() + + result = suggest_partition_key_for_dag(dag_id=dag_id, timetable=timetable, now=NOW, session=session) + assert result == "pending-key" + + def test_mapper_derived_key_takes_priority_over_recent_run(self, dag_maker, session): + dag_id = "suggest_pk_mapper_beats_run_dag" + timetable = PartitionedAssetTimetable( + assets=Asset(name=f"s3://bucket/{dag_id}", uri=f"s3://bucket/{dag_id}"), + default_partition_mapper=StartOfDayMapper(), + ) + with dag_maker(dag_id=dag_id, schedule=timetable, serialized=True, session=session): + EmptyOperator(task_id="t") + dag_maker.create_dagrun(run_id="older", state=DagRunState.SUCCESS, partition_key="2024-01-01") + session.commit() + + result = suggest_partition_key_for_dag(dag_id=dag_id, timetable=timetable, now=NOW, session=session) + assert result == "2025-06-01" + + def test_pending_apdr_is_picked_in_scheduler_fifo_order(self, dag_maker, session): + """The oldest pending APDR wins, matching the scheduler's FIFO claim order.""" + dag_id = "suggest_pk_apdr_fifo_dag" + timetable = _asset_driven_timetable_with_default_mapper(f"s3://bucket/{dag_id}") + with dag_maker(dag_id=dag_id, schedule=timetable, serialized=True, session=session): + EmptyOperator(task_id="t") + session.add( + AssetPartitionDagRun( + target_dag_id=dag_id, + partition_key="oldest-pending", + created_at=pendulum.datetime(2025, 5, 1, tz="UTC"), + ) + ) + session.add( + AssetPartitionDagRun( + target_dag_id=dag_id, + partition_key="newest-pending", + created_at=pendulum.datetime(2025, 5, 2, tz="UTC"), + ) + ) + session.commit() + + result = suggest_partition_key_for_dag(dag_id=dag_id, timetable=timetable, now=NOW, session=session) + assert result == "oldest-pending" + + def test_returns_none_and_issues_no_query_for_non_partitioned_timetable(self): + mock_session = mock.MagicMock() + + result = suggest_partition_key_for_dag( + dag_id="not-partitioned-dag", timetable=NullTimetable(), now=NOW, session=mock_session + ) + + assert result is None + mock_session.execute.assert_not_called() diff --git a/airflow-core/tests/unit/api_fastapi/core_api/datamodels/test_dags.py b/airflow-core/tests/unit/api_fastapi/core_api/datamodels/test_dags.py index 9536a1719efac..88d11ce07330b 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/datamodels/test_dags.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/datamodels/test_dags.py @@ -22,9 +22,12 @@ from airflow._shared.module_loading import qualname from airflow.api_fastapi.core_api.datamodels.dags import ( + DAG_ALIAS_MAPPING, + DAG_RESPONSE_ROUTE_SUPPLIED_FIELDS, DAGDetailsResponse, DAGResponse, ) +from airflow.models.dag import DagModel from airflow.utils.types import DagRunType @@ -46,6 +49,7 @@ def _make_dag_response(**overrides) -> DAGResponse: "timetable_summary": "0 * * * *", "timetable_description": "At the start of every hour", "timetable_partitioned": False, + "timetable_partitioned_at_runtime": False, "timetable_periodic": True, "tags": [], "max_active_tasks": 16, @@ -174,3 +178,30 @@ def test_non_periodic_with_backfill_in_allowed_run_types(self): allowed_run_types=[DagRunType.BACKFILL_JOB, DagRunType.MANUAL], ) assert dag.is_backfillable is False + + +class TestDagResponseFieldsResolveAgainstDagModel: + """ + Guard the contract that ``routes/ui/dags.py`` relies on. + + That route builds its response by iterating ``DAGResponse.model_fields`` and + reading each name off a bare ``DagModel`` with no ``getattr`` default, so a + field that is neither a ``DagModel`` attribute nor declared route-supplied + turns the Dags list page into a 500. Failing here points at the new field + instead. + """ + + def test_every_non_route_supplied_field_exists_on_dag_model(self): + missing = [ + field_name + for field_name in DAGResponse.model_fields + if field_name not in DAG_RESPONSE_ROUTE_SUPPLIED_FIELDS + and not hasattr(DagModel, DAG_ALIAS_MAPPING.get(field_name, field_name)) + ] + assert missing == [] + + def test_route_supplied_fields_are_optional(self): + for field_name in DAG_RESPONSE_ROUTE_SUPPLIED_FIELDS: + assert not DAGResponse.model_fields[field_name].is_required(), ( + f"{field_name} is route-supplied, so routes that omit it must still validate" + ) 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..c94c5b37de975 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 @@ -23,13 +23,15 @@ import pytest from sqlalchemy import delete, insert, select, update -from airflow.models.asset import AssetModel, DagScheduleAssetReference +from airflow.models.asset import AssetModel, AssetPartitionDagRun, DagScheduleAssetReference from airflow.models.dag import DagModel, DagTag from airflow.models.dag_favorite import DagFavorite from airflow.models.dagbundle import DagBundleModel from airflow.models.dagrun import DagRun from airflow.models.team import Team from airflow.providers.standard.operators.empty import EmptyOperator +from airflow.sdk import Asset +from airflow.timetables.simple import PartitionedAssetTimetable, PartitionedAtRuntime from airflow.utils.state import DagRunState, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType @@ -1078,10 +1080,12 @@ 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", "timetable_partitioned": False, + "timetable_partitioned_at_runtime": False, "timetable_periodic": False, "timetable_summary": None, "timezone": UTC_JSON_REPR, @@ -1181,17 +1185,33 @@ 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, "timetable_description": "Never, external triggers only", "timetable_partitioned": False, + "timetable_partitioned_at_runtime": False, "timetable_periodic": False, "timezone": UTC_JSON_REPR, "team_name": None, } assert res_json == expected + def test_dag_details_returns_suggested_partition_key(self, session, test_client, dag_maker): + dag_id = "test_dag_details_partition_suggestion" + with dag_maker(dag_id=dag_id, schedule=PartitionedAtRuntime(), serialized=True): + EmptyOperator(task_id="task1") + dag_maker.create_dagrun(state=DagRunState.SUCCESS, partition_key="runtime-key") + dag_maker.sync_dagbag_to_db() + session.commit() + + response = test_client.get(f"/dags/{dag_id}/details") + assert response.status_code == 200 + res_json = response.json() + assert res_json["timetable_partitioned_at_runtime"] is True + assert res_json["suggested_partition_key"] == "runtime-key" + def test_dag_details_should_response_401(self, unauthenticated_test_client): response = unauthenticated_test_client.get(f"/dags/{DAG1_ID}/details") assert response.status_code == 401 @@ -1382,9 +1402,11 @@ def test_get_dag( "next_dagrun_run_after": None, "owners": ["airflow"], "relative_fileloc": "test_dags.py", + "suggested_partition_key": None, "tags": tags, "timetable_description": "Never, external triggers only", "timetable_partitioned": False, + "timetable_partitioned_at_runtime": False, "timetable_periodic": False, "timetable_summary": None, } @@ -1417,6 +1439,55 @@ def test_get_dag_tags_sorted_alphabetically(self, session, test_client, dag_make expected_sorted_tags = sorted(tag_names) assert tag_names_in_response == expected_sorted_tags + def test_get_dag_returns_suggested_partition_key_for_asset_driven_dag( + self, session, test_client, dag_maker + ): + dag_id = "test_dag_asset_driven_suggestion" + asset = Asset(name=f"{dag_id}_asset", uri=f"s3://bucket/{dag_id}") + with dag_maker(dag_id=dag_id, schedule=PartitionedAssetTimetable(assets=asset), serialized=True): + EmptyOperator(task_id="task1") + dag_maker.sync_dagbag_to_db() + session.add(AssetPartitionDagRun(target_dag_id=dag_id, partition_key="pending-key")) + session.commit() + + response = test_client.get(f"/dags/{dag_id}") + assert response.status_code == 200 + res_json = response.json() + assert res_json["timetable_partitioned"] is True + assert res_json["timetable_partitioned_at_runtime"] is False + assert res_json["suggested_partition_key"] == "pending-key" + + def test_get_dag_returns_suggested_partition_key_for_partitioned_at_runtime_dag( + self, session, test_client, dag_maker + ): + dag_id = "test_dag_partitioned_at_runtime_suggestion" + with dag_maker(dag_id=dag_id, schedule=PartitionedAtRuntime(), serialized=True): + EmptyOperator(task_id="task1") + dag_maker.create_dagrun(state=DagRunState.SUCCESS, partition_key="runtime-key") + dag_maker.sync_dagbag_to_db() + session.commit() + + response = test_client.get(f"/dags/{dag_id}") + assert response.status_code == 200 + res_json = response.json() + assert res_json["timetable_partitioned"] is False + assert res_json["timetable_partitioned_at_runtime"] is True + assert res_json["suggested_partition_key"] == "runtime-key" + + def test_get_dags_list_reports_partitioned_at_runtime(self, session, test_client, dag_maker): + """The list endpoint returns bare DagModel rows, so this only holds if the flag is a column.""" + dag_id = "test_dag_list_partitioned_at_runtime" + with dag_maker(dag_id=dag_id, schedule=PartitionedAtRuntime(), serialized=True): + EmptyOperator(task_id="task1") + dag_maker.sync_dagbag_to_db() + session.commit() + + response = test_client.get("/dags", params={"dag_id_pattern": dag_id}) + assert response.status_code == 200 + [dag] = response.json()["dags"] + assert dag["timetable_partitioned_at_runtime"] is True + assert dag["suggested_partition_key"] is None + def test_get_dag_should_response_401(self, unauthenticated_test_client): response = unauthenticated_test_client.get(f"/dags/{DAG1_ID}") assert response.status_code == 401 diff --git a/airflow-core/tests/unit/timetables/test_partitioned_timetable.py b/airflow-core/tests/unit/timetables/test_partitioned_timetable.py index 72dd2773bc915..8ea2566c8ad2d 100644 --- a/airflow-core/tests/unit/timetables/test_partitioned_timetable.py +++ b/airflow-core/tests/unit/timetables/test_partitioned_timetable.py @@ -19,7 +19,6 @@ from collections.abc import Callable, Iterable from contextlib import ExitStack -from typing import TYPE_CHECKING from unittest import mock import pendulum @@ -27,19 +26,20 @@ from airflow._shared.module_loading import qualname from airflow.exceptions import InvalidPartitionKeyError -from airflow.partition_mappers.base import RollupMapper +from airflow.partition_mappers.allowed_key import AllowedKeyMapper +from airflow.partition_mappers.base import PartitionMapper, RollupMapper +from airflow.partition_mappers.chain import ChainMapper +from airflow.partition_mappers.fixed_key import FixedKeyMapper from airflow.partition_mappers.identity import IdentityMapper as IdentityMapper -from airflow.partition_mappers.temporal import StartOfDayMapper -from airflow.partition_mappers.window import DayWindow +from airflow.partition_mappers.product import ProductMapper +from airflow.partition_mappers.temporal import FanOutMapper, StartOfDayMapper, StartOfHourMapper +from airflow.partition_mappers.window import DayWindow, WeekWindow from airflow.sdk import Asset, AssetAlias from airflow.serialization.definitions.assets import SerializedAsset from airflow.serialization.encoders import ensure_serialized_asset from airflow.serialization.enums import DagAttributeTypes from airflow.timetables.simple import PartitionedAssetTimetable -if TYPE_CHECKING: - from airflow.partition_mappers.base import PartitionMapper - class Key1Mapper(IdentityMapper): """Partition Mapper that returns only key-1 as downstream key""" @@ -279,3 +279,95 @@ def test_decode_partition_date_returns_period_start_for_valid_key(self): default_partition_mapper=StartOfDayMapper(), ) assert timetable._decode_partition_date("2025-01-01") == pendulum.datetime(2025, 1, 1, tz="UTC") + + +class _RaisingMapper(PartitionMapper): + """Mapper with a callable ``normalize``/``format`` pair where ``format`` always raises.""" + + def to_downstream(self, key: str) -> str: + return key + + def normalize(self, dt): + return dt + + def format(self, dt) -> str: + raise ValueError("boom") + + +class TestPartitionedAssetTimetableSuggestPartitionKey: + """ + ``suggest_partition_key`` enumeration of ``airflow.partition_mappers`` (Step 2): + + - Can answer (has callable ``normalize`` and ``format``): the + ``_BaseTemporalMapper`` family — ``StartOfHourMapper``, ``StartOfDayMapper``, + ``StartOfWeekMapper``, ``StartOfMonthMapper``, ``StartOfQuarterMapper``, + ``StartOfYearMapper``. Only ``StartOfDayMapper``/``StartOfHourMapper`` are + exercised directly here; the others share the same base implementation. + - Cannot answer: ``RollupMapper``, ``FanOutMapper``, ``ChainMapper``, + ``ProductMapper``, ``IdentityMapper``, ``FixedKeyMapper``, ``AllowedKeyMapper``. + """ + + def test_suggest_partition_key_returns_mapper_answer_when_all_agree(self): + timetable = PartitionedAssetTimetable( + assets=Asset(name="daily", uri="s3://bucket/daily"), + default_partition_mapper=StartOfDayMapper(), + ) + now = pendulum.datetime(2025, 1, 1, 10, tz="UTC") + assert timetable.suggest_partition_key(now) == "2025-01-01" + + def test_suggest_partition_key_returns_none_when_mappers_disagree(self): + name_ref = ensure_serialized_asset(Asset.ref(name="daily")) + uri_ref = ensure_serialized_asset(Asset.ref(uri="s3://bucket/hourly")) + timetable = PartitionedAssetTimetable( + assets=[name_ref, uri_ref], + partition_mapper_config={ + name_ref: StartOfDayMapper(), + uri_ref: StartOfHourMapper(), + }, + ) + now = pendulum.datetime(2025, 1, 1, 10, tz="UTC") + assert timetable.suggest_partition_key(now) is None + + @pytest.mark.parametrize( + "mapper", + [ + IdentityMapper(), + FixedKeyMapper("fixed"), + AllowedKeyMapper(["a", "b"]), + ChainMapper(IdentityMapper(), IdentityMapper()), + ProductMapper(IdentityMapper(), IdentityMapper()), + RollupMapper(upstream_mapper=StartOfDayMapper(), window=DayWindow()), + FanOutMapper(upstream_mapper=StartOfDayMapper(), window=WeekWindow()), + ], + ids=[ + "IdentityMapper", + "FixedKeyMapper", + "AllowedKeyMapper", + "ChainMapper", + "ProductMapper", + "RollupMapper", + "FanOutMapper", + ], + ) + def test_suggest_partition_key_returns_none_for_mappers_without_normalize_and_format(self, mapper): + timetable = PartitionedAssetTimetable( + assets=Asset(name="daily", uri="s3://bucket/daily"), + default_partition_mapper=mapper, + ) + assert timetable.suggest_partition_key(pendulum.datetime(2025, 1, 1, tz="UTC")) is None + + def test_suggest_partition_key_returns_none_when_no_asset(self): + timetable = PartitionedAssetTimetable(assets=AssetAlias(name="alias_only")) + assert timetable.suggest_partition_key(pendulum.datetime(2025, 1, 1, tz="UTC")) is None + + @mock.patch("airflow.timetables.simple.log") + def test_suggest_partition_key_swallows_mapper_exceptions(self, mock_log): + timetable = PartitionedAssetTimetable( + assets=Asset(name="daily", uri="s3://bucket/daily"), + default_partition_mapper=_RaisingMapper(), + ) + assert timetable.suggest_partition_key(pendulum.datetime(2025, 1, 1, tz="UTC")) is None + assert any( + call.args[0] == "Failed to suggest partition key from mapper; ignoring" + for call in mock_log.warning.mock_calls + ) diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index ca3c63d6168e2..9de3273b2126e 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -1795,6 +1795,7 @@ class DAGResponse(BaseModel): timetable_summary: Annotated[str | None, Field(title="Timetable Summary")] timetable_description: Annotated[str | None, Field(title="Timetable Description")] timetable_partitioned: Annotated[bool, Field(title="Timetable Partitioned")] + timetable_partitioned_at_runtime: Annotated[bool, Field(title="Timetable Partitioned At Runtime")] timetable_periodic: Annotated[bool, Field(title="Timetable Periodic")] tags: Annotated[list[DagTagResponse], Field(title="Tags")] max_active_tasks: Annotated[int, Field(title="Max Active Tasks")] @@ -1810,6 +1811,7 @@ class DAGResponse(BaseModel): next_dagrun_run_after: Annotated[datetime | None, Field(title="Next Dagrun Run After")] allowed_run_types: Annotated[list[DagRunType] | None, Field(title="Allowed Run Types")] owners: Annotated[list[str], Field(title="Owners")] + 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") ] @@ -2600,6 +2602,7 @@ class DAGDetailsResponse(BaseModel): timetable_summary: Annotated[str | None, Field(title="Timetable Summary")] timetable_description: Annotated[str | None, Field(title="Timetable Description")] timetable_partitioned: Annotated[bool, Field(title="Timetable Partitioned")] + timetable_partitioned_at_runtime: Annotated[bool, Field(title="Timetable Partitioned At Runtime")] timetable_periodic: Annotated[bool, Field(title="Timetable Periodic")] tags: Annotated[list[DagTagResponse], Field(title="Tags")] max_active_tasks: Annotated[int, Field(title="Max Active Tasks")] @@ -2615,6 +2618,7 @@ class DAGDetailsResponse(BaseModel): next_dagrun_run_after: Annotated[datetime | None, Field(title="Next Dagrun Run After")] allowed_run_types: Annotated[list[DagRunType] | None, Field(title="Allowed Run Types")] owners: Annotated[list[str], Field(title="Owners")] + suggested_partition_key: Annotated[str | None, Field(title="Suggested Partition Key")] = None catchup: Annotated[bool, Field(title="Catchup")] dag_run_timeout: Annotated[timedelta | None, Field(title="Dag Run Timeout")] asset_expression: Annotated[ diff --git a/airflow-ctl/tests/airflow_ctl/api/test_operations.py b/airflow-ctl/tests/airflow_ctl/api/test_operations.py index 52fa594afd5ee..0218c7fe62f41 100644 --- a/airflow-ctl/tests/airflow_ctl/api/test_operations.py +++ b/airflow-ctl/tests/airflow_ctl/api/test_operations.py @@ -959,6 +959,7 @@ class TestDagOperations: timetable_summary="timetable_summary", timetable_description="timetable_description", timetable_partitioned=False, + timetable_partitioned_at_runtime=False, timetable_periodic=True, tags=[], max_active_tasks=1, @@ -992,6 +993,7 @@ class TestDagOperations: timetable_summary="timetable_summary", timetable_description="timetable_description", timetable_partitioned=False, + timetable_partitioned_at_runtime=False, timetable_periodic=True, tags=[], max_active_tasks=1, diff --git a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_dag_command.py b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_dag_command.py index 1ee613b85f8f6..a54fbca413571 100644 --- a/airflow-ctl/tests/airflow_ctl/ctl/commands/test_dag_command.py +++ b/airflow-ctl/tests/airflow_ctl/ctl/commands/test_dag_command.py @@ -57,6 +57,7 @@ class TestDagCommands: timetable_summary="timetable_summary", timetable_description="timetable_description", timetable_partitioned=False, + timetable_partitioned_at_runtime=False, timetable_periodic=True, tags=[], max_active_tasks=1, @@ -90,6 +91,7 @@ class TestDagCommands: timetable_summary="timetable_summary", timetable_description="timetable_description", timetable_partitioned=False, + timetable_partitioned_at_runtime=False, timetable_periodic=True, tags=[], max_active_tasks=1, @@ -123,6 +125,7 @@ class TestDagCommands: timetable_summary=None, timetable_description=None, timetable_partitioned=False, + timetable_partitioned_at_runtime=False, timetable_periodic=False, tags=[], max_active_tasks=1,