From 25bd45aabffdda6074c58e6fa0fd6281ea30b9f4 Mon Sep 17 00:00:00 2001 From: Henry Chen Date: Thu, 23 Jul 2026 22:55:38 +0800 Subject: [PATCH] [v3-3-test] Add partition date filters to Dag run API (#68682) * Add partition date support to Dag run API * Reject partition date filters on non-partitioned Dags in Dag run API Runs of a non-partitioned Dag never carry a partition_date, so filtering by partition_date_start/partition_date_end silently returned an empty list, which reads as "no runs in that window" rather than "this Dag has no partition dimension". Return 400 instead, mirroring how supplying a partition_key for a non-partitioned Dag is rejected. * Rename Dag run partition date filters to partition_date_gte/lte The other range filters on the list Dag runs endpoint use the _gte/_lte suffix convention (start_date_gte, logical_date_lte, ...), so the new partition date bounds should follow it too. Also clarify in the query parameter descriptions that the inclusive bounds cover the whole calendar day, and fold the partition-date branches into a single condition so the 400 paths and the filter path read in one place. (cherry picked from commit 8e6f28fbd33b2e8ba55f2f8b0d775847b3244516) Co-authored-by: Henry Chen --- .../openapi/v2-rest-api-generated.yaml | 38 ++++++++-- .../core_api/routes/public/dag_run.py | 58 ++++++++++++++- airflow-core/src/airflow/models/dagrun.py | 14 +++- .../airflow/ui/openapi-gen/queries/common.ts | 6 +- .../ui/openapi-gen/queries/ensureQueryData.ts | 10 ++- .../ui/openapi-gen/queries/prefetch.ts | 10 ++- .../airflow/ui/openapi-gen/queries/queries.ts | 10 ++- .../ui/openapi-gen/queries/suspense.ts | 10 ++- .../ui/openapi-gen/requests/services.gen.ts | 6 +- .../ui/openapi-gen/requests/types.gen.ts | 10 ++- .../core_api/routes/public/test_dag_run.py | 72 ++++++++++++++++++- 11 files changed, 217 insertions(+), 27 deletions(-) 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 75caaf4f033ea..1f73622f36a23 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 @@ -2325,6 +2325,36 @@ paths: schema: type: string title: Dag Id + - name: partition_date_gte + in: query + required: false + schema: + anyOf: + - type: string + format: date + - type: 'null' + description: Inclusive lower bound of the partition_date window, interpreted + as a local calendar day in the Dag's timetable timezone. Runs from the + start of this day onwards match. + title: Partition Date Gte + description: Inclusive lower bound of the partition_date window, interpreted + as a local calendar day in the Dag's timetable timezone. Runs from the start + of this day onwards match. + - name: partition_date_lte + in: query + required: false + schema: + anyOf: + - type: string + format: date + - type: 'null' + description: 'Inclusive upper bound of the partition_date window, interpreted + as a local calendar day in the Dag''s timetable timezone. The whole day + is included: runs up to the end of this day match.' + title: Partition Date Lte + description: 'Inclusive upper bound of the partition_date window, interpreted + as a local calendar day in the Dag''s timetable timezone. The whole day + is included: runs up to the end of this day match.' - name: cursor in: query required: false @@ -2614,15 +2644,15 @@ paths: type: string description: 'Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, state, - dag_id, run_id, logical_date, run_after, start_date, end_date, updated_at, - conf, duration, dag_run_id`' + dag_id, run_id, logical_date, partition_date, run_after, start_date, end_date, + updated_at, conf, duration, dag_run_id`' default: - id title: Order By description: 'Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, state, dag_id, - run_id, logical_date, run_after, start_date, end_date, updated_at, conf, - duration, dag_run_id`' + run_id, logical_date, partition_date, run_after, start_date, end_date, updated_at, + conf, duration, dag_run_id`' - name: run_id_pattern in: query required: false diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py index 955c7aa34ce10..84e2f77f8faa4 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_run.py @@ -17,6 +17,7 @@ from __future__ import annotations +import datetime import textwrap from typing import Annotated, Literal, cast @@ -509,6 +510,7 @@ def get_dag_runs( "dag_id", "run_id", "logical_date", + "partition_date", "run_after", "start_date", "end_date", @@ -547,6 +549,21 @@ def get_dag_runs( partition_key_pattern: QueryDagRunPartitionKeySearch, partition_key_prefix_pattern: QueryDagRunPartitionKeyPrefixSearch, consuming_asset_pattern: QueryConsumingAssetPatternSearch, + partition_date_gte: datetime.date | None = Query( + None, + description=( + "Inclusive lower bound of the partition_date window, interpreted as a local calendar " + "day in the Dag's timetable timezone. Runs from the start of this day onwards match." + ), + ), + partition_date_lte: datetime.date | None = Query( + None, + description=( + "Inclusive upper bound of the partition_date window, interpreted as a local calendar " + "day in the Dag's timetable timezone. The whole day is included: runs up to the end " + "of this day match." + ), + ), cursor: str | None = Query( None, description="Cursor for keyset-based pagination. " @@ -571,9 +588,46 @@ def get_dag_runs( use_cursor = cursor is not None query = select(DagRun).options(*eager_load_dag_run_for_list()) - if dag_id != "~": - get_latest_version_of_dag(dag_bag, dag_id, session) # Check if the Dag exists. + has_partition_date_filter = partition_date_gte is not None or partition_date_lte is not None + + if dag_id == "~": + if has_partition_date_filter: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, + "partition_date_gte and partition_date_lte require a specific dag_id.", + ) + else: + dag = get_latest_version_of_dag(dag_bag, dag_id, session) # Check if the Dag exists. query = query.filter(DagRun.dag_id == dag_id).options() + if has_partition_date_filter: + # Runs of a non-partitioned Dag never carry a partition_date (this includes + # partitioned-at-runtime Dags, whose runs keep it NULL), so the filter would + # silently match nothing; reject it instead. + if not dag.timetable.partitioned: + raise HTTPException( + status.HTTP_400_BAD_REQUEST, + f"Dag with dag_id: '{dag_id}' is not partitioned; " + "partition_date_gte and partition_date_lte are not supported.", + ) + # The bounds are calendar days, so the whole of partition_date_lte belongs to the + # window: widen it to the following local midnight and exclude that edge. + query = DagRun.apply_partition_date_window( + query, + timetable=dag.timetable, + start=( + datetime.datetime.combine(partition_date_gte, datetime.time.min) + if partition_date_gte is not None + else None + ), + end=( + datetime.datetime.combine( + partition_date_lte + datetime.timedelta(days=1), datetime.time.min + ) + if partition_date_lte is not None + else None + ), + end_exclusive=True, + ) # Add join with DagVersion if dag_version filter is active if dag_version.value: diff --git a/airflow-core/src/airflow/models/dagrun.py b/airflow-core/src/airflow/models/dagrun.py index f5ae6a28d5824..1ca2199252b4e 100644 --- a/airflow-core/src/airflow/models/dagrun.py +++ b/airflow-core/src/airflow/models/dagrun.py @@ -2225,14 +2225,24 @@ def apply_partition_date_window( timetable: Timetable, start: datetime | None, end: datetime | None, + end_exclusive: bool = False, ) -> Select: - """Filter stmt to the inclusive interval [lower, upper] on partition_date.""" + """ + Filter stmt to a partition_date window bounded by *start* and *end*. + + The lower bound is always inclusive. The upper bound is inclusive by default; + pass ``end_exclusive=True`` when *end* is the open edge of the window, as with + a date-only filter widened to the following local midnight. + """ if start is not None: lower = timetable.localize_partition_datetime(start) stmt = stmt.where(DagRun.partition_date >= lower) if end is not None: upper = timetable.localize_partition_datetime(end) - stmt = stmt.where(DagRun.partition_date <= upper) + if end_exclusive: + stmt = stmt.where(DagRun.partition_date < upper) + else: + stmt = stmt.where(DagRun.partition_date <= upper) return stmt diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts index 3d923f44e9a3b..d7baf71cf539a 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/common.ts @@ -156,7 +156,7 @@ export const UseDagRunServiceGetDagRunKeyFn = ({ dagId, dagRunId }: { export type DagRunServiceGetDagRunsDefaultResponse = Awaited>; export type DagRunServiceGetDagRunsQueryResult = UseQueryResult; export const useDagRunServiceGetDagRunsKey = "DagRunServiceGetDagRuns"; -export const UseDagRunServiceGetDagRunsKeyFn = ({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }: { +export const UseDagRunServiceGetDagRunsKeyFn = ({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }: { bundleVersion?: string; confContains?: string; consumingAssetPattern?: string; @@ -180,6 +180,8 @@ export const UseDagRunServiceGetDagRunsKeyFn = ({ bundleVersion, confContains, c logicalDateLte?: string; offset?: number; orderBy?: string[]; + partitionDateGte?: string; + partitionDateLte?: string; partitionKeyPattern?: string; partitionKeyPrefixPattern?: string; runAfterGt?: string; @@ -200,7 +202,7 @@ export const UseDagRunServiceGetDagRunsKeyFn = ({ bundleVersion, confContains, c updatedAtGte?: string; updatedAtLt?: string; updatedAtLte?: string; -}, queryKey?: Array) => [useDagRunServiceGetDagRunsKey, ...(queryKey ?? [{ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }])]; +}, queryKey?: Array) => [useDagRunServiceGetDagRunsKey, ...(queryKey ?? [{ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }])]; export type DagRunServiceGetUpstreamAssetEventsDefaultResponse = Awaited>; export type DagRunServiceGetUpstreamAssetEventsQueryResult = UseQueryResult; export const useDagRunServiceGetUpstreamAssetEventsKey = "DagRunServiceGetUpstreamAssetEvents"; diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts index c4a9f18e8d8f3..edcc8e6d4dd08 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/ensureQueryData.ts @@ -307,6 +307,8 @@ export const ensureUseDagRunServiceGetDagRunData = (queryClient: QueryClient, { * on the first page. * @param data The data for the request. * @param data.dagId +* @param data.partitionDateGte Inclusive lower bound of the partition_date window, interpreted as a local calendar day in the Dag's timetable timezone. Runs from the start of this day onwards match. +* @param data.partitionDateLte Inclusive upper bound of the partition_date window, interpreted as a local calendar day in the Dag's timetable timezone. The whole day is included: runs up to the end of this day match. * @param data.cursor Cursor for keyset-based pagination. Pass an empty string for the first page, then use ``next_cursor`` from the response. When ``cursor`` is provided, ``offset`` is ignored. * @param data.limit * @param data.offset @@ -339,7 +341,7 @@ export const ensureUseDagRunServiceGetDagRunData = (queryClient: QueryClient, { * @param data.state * @param data.dagVersion * @param data.bundleVersion -* @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, state, dag_id, run_id, logical_date, run_after, start_date, end_date, updated_at, conf, duration, dag_run_id` +* @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, state, dag_id, run_id, logical_date, partition_date, run_after, start_date, end_date, updated_at, conf, duration, dag_run_id` * @param data.runIdPattern SQL LIKE expression — use `%` / `_` wildcards (e.g. `%customer_%`). Use the pipe `|` operator for OR logic (e.g. `dag1 | dag2`). Regular expressions are **not** supported. * * **Performance note:** this full-match pattern is evaluated as ``ILIKE '%term%'`` and most of the time prevents the database from using B-tree indexes, which can be very slow on large tables. Prefer the equivalent ``run_id_prefix_pattern`` parameter when possible. @@ -360,7 +362,7 @@ export const ensureUseDagRunServiceGetDagRunData = (queryClient: QueryClient, { * @returns DAGRunCollectionResponse Successful Response * @throws ApiError */ -export const ensureUseDagRunServiceGetDagRunsData = (queryClient: QueryClient, { bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }: { +export const ensureUseDagRunServiceGetDagRunsData = (queryClient: QueryClient, { bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }: { bundleVersion?: string; confContains?: string; consumingAssetPattern?: string; @@ -384,6 +386,8 @@ export const ensureUseDagRunServiceGetDagRunsData = (queryClient: QueryClient, { logicalDateLte?: string; offset?: number; orderBy?: string[]; + partitionDateGte?: string; + partitionDateLte?: string; partitionKeyPattern?: string; partitionKeyPrefixPattern?: string; runAfterGt?: string; @@ -404,7 +408,7 @@ export const ensureUseDagRunServiceGetDagRunsData = (queryClient: QueryClient, { updatedAtGte?: string; updatedAtLt?: string; updatedAtLte?: string; -}) => queryClient.ensureQueryData({ queryKey: Common.UseDagRunServiceGetDagRunsKeyFn({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }), queryFn: () => DagRunService.getDagRuns({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }) }); +}) => queryClient.ensureQueryData({ queryKey: Common.UseDagRunServiceGetDagRunsKeyFn({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }), queryFn: () => DagRunService.getDagRuns({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }) }); /** * Get Upstream Asset Events * If dag run is asset-triggered, return the asset events that triggered it. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts index ce348567d74a7..df3d75071d492 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/prefetch.ts @@ -307,6 +307,8 @@ export const prefetchUseDagRunServiceGetDagRun = (queryClient: QueryClient, { da * on the first page. * @param data The data for the request. * @param data.dagId +* @param data.partitionDateGte Inclusive lower bound of the partition_date window, interpreted as a local calendar day in the Dag's timetable timezone. Runs from the start of this day onwards match. +* @param data.partitionDateLte Inclusive upper bound of the partition_date window, interpreted as a local calendar day in the Dag's timetable timezone. The whole day is included: runs up to the end of this day match. * @param data.cursor Cursor for keyset-based pagination. Pass an empty string for the first page, then use ``next_cursor`` from the response. When ``cursor`` is provided, ``offset`` is ignored. * @param data.limit * @param data.offset @@ -339,7 +341,7 @@ export const prefetchUseDagRunServiceGetDagRun = (queryClient: QueryClient, { da * @param data.state * @param data.dagVersion * @param data.bundleVersion -* @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, state, dag_id, run_id, logical_date, run_after, start_date, end_date, updated_at, conf, duration, dag_run_id` +* @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, state, dag_id, run_id, logical_date, partition_date, run_after, start_date, end_date, updated_at, conf, duration, dag_run_id` * @param data.runIdPattern SQL LIKE expression — use `%` / `_` wildcards (e.g. `%customer_%`). Use the pipe `|` operator for OR logic (e.g. `dag1 | dag2`). Regular expressions are **not** supported. * * **Performance note:** this full-match pattern is evaluated as ``ILIKE '%term%'`` and most of the time prevents the database from using B-tree indexes, which can be very slow on large tables. Prefer the equivalent ``run_id_prefix_pattern`` parameter when possible. @@ -360,7 +362,7 @@ export const prefetchUseDagRunServiceGetDagRun = (queryClient: QueryClient, { da * @returns DAGRunCollectionResponse Successful Response * @throws ApiError */ -export const prefetchUseDagRunServiceGetDagRuns = (queryClient: QueryClient, { bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }: { +export const prefetchUseDagRunServiceGetDagRuns = (queryClient: QueryClient, { bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }: { bundleVersion?: string; confContains?: string; consumingAssetPattern?: string; @@ -384,6 +386,8 @@ export const prefetchUseDagRunServiceGetDagRuns = (queryClient: QueryClient, { b logicalDateLte?: string; offset?: number; orderBy?: string[]; + partitionDateGte?: string; + partitionDateLte?: string; partitionKeyPattern?: string; partitionKeyPrefixPattern?: string; runAfterGt?: string; @@ -404,7 +408,7 @@ export const prefetchUseDagRunServiceGetDagRuns = (queryClient: QueryClient, { b updatedAtGte?: string; updatedAtLt?: string; updatedAtLte?: string; -}) => queryClient.prefetchQuery({ queryKey: Common.UseDagRunServiceGetDagRunsKeyFn({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }), queryFn: () => DagRunService.getDagRuns({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }) }); +}) => queryClient.prefetchQuery({ queryKey: Common.UseDagRunServiceGetDagRunsKeyFn({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }), queryFn: () => DagRunService.getDagRuns({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }) }); /** * Get Upstream Asset Events * If dag run is asset-triggered, return the asset events that triggered it. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts index 8bc74b305aa17..ab04de50cb199 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/queries.ts @@ -307,6 +307,8 @@ export const useDagRunServiceGetDagRun = = unknown[]>({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }: { +export const useDagRunServiceGetDagRuns = = unknown[]>({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }: { bundleVersion?: string; confContains?: string; consumingAssetPattern?: string; @@ -384,6 +386,8 @@ export const useDagRunServiceGetDagRuns = , "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseDagRunServiceGetDagRunsKeyFn({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }, queryKey), queryFn: () => DagRunService.getDagRuns({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }) as TData, ...options }); +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useQuery({ queryKey: Common.UseDagRunServiceGetDagRunsKeyFn({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }, queryKey), queryFn: () => DagRunService.getDagRuns({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }) as TData, ...options }); /** * Get Upstream Asset Events * If dag run is asset-triggered, return the asset events that triggered it. diff --git a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts index 5c22b3bf33873..4fc85f3bc864d 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/queries/suspense.ts @@ -307,6 +307,8 @@ export const useDagRunServiceGetDagRunSuspense = = unknown[]>({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }: { +export const useDagRunServiceGetDagRunsSuspense = = unknown[]>({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }: { bundleVersion?: string; confContains?: string; consumingAssetPattern?: string; @@ -384,6 +386,8 @@ export const useDagRunServiceGetDagRunsSuspense = , "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseDagRunServiceGetDagRunsKeyFn({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }, queryKey), queryFn: () => DagRunService.getDagRuns({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }) as TData, ...options }); +}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">) => useSuspenseQuery({ queryKey: Common.UseDagRunServiceGetDagRunsKeyFn({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }, queryKey), queryFn: () => DagRunService.getDagRuns({ bundleVersion, confContains, consumingAssetPattern, cursor, dagId, dagIdPattern, dagIdPrefixPattern, dagVersion, durationGt, durationGte, durationLt, durationLte, endDateGt, endDateGte, endDateLt, endDateLte, limit, logicalDateGt, logicalDateGte, logicalDateLt, logicalDateLte, offset, orderBy, partitionDateGte, partitionDateLte, partitionKeyPattern, partitionKeyPrefixPattern, runAfterGt, runAfterGte, runAfterLt, runAfterLte, runIdPattern, runIdPrefixPattern, runType, startDateGt, startDateGte, startDateLt, startDateLte, state, triggeringUserNamePattern, triggeringUserNamePrefixPattern, updatedAtGt, updatedAtGte, updatedAtLt, updatedAtLte }) as TData, ...options }); /** * Get Upstream Asset Events * If dag run is asset-triggered, return the asset events that triggered it. diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts index 69e12c4514c33..8bbaed47da9ce 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/services.gen.ts @@ -1059,6 +1059,8 @@ export class DagRunService { * on the first page. * @param data The data for the request. * @param data.dagId + * @param data.partitionDateGte Inclusive lower bound of the partition_date window, interpreted as a local calendar day in the Dag's timetable timezone. Runs from the start of this day onwards match. + * @param data.partitionDateLte Inclusive upper bound of the partition_date window, interpreted as a local calendar day in the Dag's timetable timezone. The whole day is included: runs up to the end of this day match. * @param data.cursor Cursor for keyset-based pagination. Pass an empty string for the first page, then use ``next_cursor`` from the response. When ``cursor`` is provided, ``offset`` is ignored. * @param data.limit * @param data.offset @@ -1091,7 +1093,7 @@ export class DagRunService { * @param data.state * @param data.dagVersion * @param data.bundleVersion - * @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, state, dag_id, run_id, logical_date, run_after, start_date, end_date, updated_at, conf, duration, dag_run_id` + * @param data.orderBy Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, state, dag_id, run_id, logical_date, partition_date, run_after, start_date, end_date, updated_at, conf, duration, dag_run_id` * @param data.runIdPattern SQL LIKE expression — use `%` / `_` wildcards (e.g. `%customer_%`). Use the pipe `|` operator for OR logic (e.g. `dag1 | dag2`). Regular expressions are **not** supported. * * **Performance note:** this full-match pattern is evaluated as ``ILIKE '%term%'`` and most of the time prevents the database from using B-tree indexes, which can be very slow on large tables. Prefer the equivalent ``run_id_prefix_pattern`` parameter when possible. @@ -1120,6 +1122,8 @@ export class DagRunService { dag_id: data.dagId }, query: { + partition_date_gte: data.partitionDateGte, + partition_date_lte: data.partitionDateLte, cursor: data.cursor, limit: data.limit, offset: data.offset, 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 ce406b5fcb671..1c6e5e073b73b 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 @@ -3103,9 +3103,17 @@ export type GetDagRunsData = { logicalDateLte?: string | null; offset?: number; /** - * Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, state, dag_id, run_id, logical_date, run_after, start_date, end_date, updated_at, conf, duration, dag_run_id` + * Attributes to order by, multi criteria sort is supported. Prefix with `-` for descending order. Supported attributes: `id, state, dag_id, run_id, logical_date, partition_date, run_after, start_date, end_date, updated_at, conf, duration, dag_run_id` */ orderBy?: Array<(string)>; + /** + * Inclusive lower bound of the partition_date window, interpreted as a local calendar day in the Dag's timetable timezone. Runs from the start of this day onwards match. + */ + partitionDateGte?: string | null; + /** + * Inclusive upper bound of the partition_date window, interpreted as a local calendar day in the Dag's timetable timezone. The whole day is included: runs up to the end of this day match. + */ + partitionDateLte?: string | null; /** * SQL LIKE expression — use `%` / `_` wildcards (e.g. `%customer_%`). The pipe `|` is matched literally, not as an OR separator. Regular expressions are **not** supported. * diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py index 52ffc947edaec..930ba9786836b 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_dag_run.py @@ -128,6 +128,10 @@ def custom_timetable_plugin(monkeypatch): START_DATE2 = datetime(2024, 4, 15, 0, 0, tzinfo=timezone.utc) LOGICAL_DATE3 = datetime(2024, 5, 16, 0, 0, tzinfo=timezone.utc) LOGICAL_DATE4 = datetime(2024, 5, 25, 0, 0, tzinfo=timezone.utc) +PARTITION_DATE1 = datetime(2024, 6, 1, 0, 0, tzinfo=timezone.utc) +PARTITION_DATE2 = datetime(2024, 6, 2, 0, 0, tzinfo=timezone.utc) +PARTITION_DATE3 = datetime(2024, 6, 3, 0, 0, tzinfo=timezone.utc) +PARTITION_DATE4 = datetime(2024, 6, 4, 0, 0, tzinfo=timezone.utc) DAG1_RUN1_NOTE = "test_note" DAG2_PARAM = {"validated_number": Param(1, minimum=1, maximum=10)} @@ -173,6 +177,7 @@ def setup(request, dag_maker, *, session=None): dag_run1.end_date = dag_run1.start_date + timedelta(seconds=101) # Set conf for testing conf_contains filter (values ordered for predictable sorting) dag_run1.conf = {"env": "development", "version": "1.0"} + dag_run1.partition_date = PARTITION_DATE1 for i, t in enumerate([task1, task2], start=1): ti = dag_run1.get_task_instance(task_id=t.task_id) @@ -203,6 +208,7 @@ def setup(request, dag_maker, *, session=None): dag_run2.end_date = dag_run2.start_date + timedelta(seconds=201) # Set conf for testing conf_contains filter dag_run2.conf = {"env": "production", "debug": True} + dag_run2.partition_date = PARTITION_DATE2 ti1 = dag_run2.get_task_instance(task_id=task1.task_id) ti1.task = task1 @@ -230,6 +236,7 @@ def setup(request, dag_maker, *, session=None): dag_run3.end_date = dag_run3.start_date + timedelta(seconds=51) # Set conf for testing conf_contains filter dag_run3.conf = {"env": "staging", "test_mode": True} + dag_run3.partition_date = PARTITION_DATE3 dag_run4 = dag_maker.create_dagrun( run_id=DAG2_RUN2_ID, @@ -244,6 +251,7 @@ def setup(request, dag_maker, *, session=None): dag_run4.end_date = dag_run4.start_date + timedelta(seconds=150) # Set conf for testing conf_contains filter dag_run4.conf = {"env": "testing", "mode": "ci"} + dag_run4.partition_date = PARTITION_DATE4 dag_maker.sync_dagbag_to_db() dag_maker.dag_model.has_task_concurrency_limits = True @@ -318,9 +326,9 @@ def get_dag_run_dict(run: DagRun): "note": run.note, "dag_versions": get_dag_versions_dict(run.dag_versions), "partition_key": run.partition_key, - "partition_date": from_datetime_to_zulu_without_ms(run.partition_date) - if run.partition_date - else None, + "partition_date": ( + from_datetime_to_zulu_without_ms(run.partition_date) if run.partition_date else None + ), } @@ -413,6 +421,21 @@ def test_get_dag_runs_not_found(self, test_client): body = response.json() assert body["detail"] == "The Dag with ID: `invalid` was not found" + def test_partition_date_day_filters_reject_all_dags_selector(self, test_client): + response = test_client.get("/dags/~/dagRuns", params={"partition_date_gte": "2025-01-01"}) + assert response.status_code == 400 + assert response.json()["detail"] == ( + "partition_date_gte and partition_date_lte require a specific dag_id." + ) + + def test_partition_date_day_filters_reject_non_partitioned_dag(self, test_client): + response = test_client.get(f"/dags/{DAG1_ID}/dagRuns", params={"partition_date_gte": "2025-01-01"}) + assert response.status_code == 400 + assert response.json()["detail"] == ( + f"Dag with dag_id: '{DAG1_ID}' is not partitioned; " + "partition_date_gte and partition_date_lte are not supported." + ) + def test_invalid_order_by_raises_400(self, test_client): response = test_client.get("/dags/test_dag1/dagRuns?order_by=invalid") assert response.status_code == 400 @@ -437,6 +460,7 @@ def test_should_respond_403(self, unauthorized_test_client): pytest.param("state", [DAG1_RUN2_ID, DAG1_RUN1_ID], id="order_by_state"), pytest.param("dag_id", [DAG1_RUN1_ID, DAG1_RUN2_ID], id="order_by_dag_id"), pytest.param("logical_date", [DAG1_RUN1_ID, DAG1_RUN2_ID], id="order_by_logical_date"), + pytest.param("partition_date", [DAG1_RUN1_ID, DAG1_RUN2_ID], id="order_by_partition_date"), pytest.param("dag_run_id", [DAG1_RUN1_ID, DAG1_RUN2_ID], id="order_by_dag_run_id"), pytest.param("start_date", [DAG1_RUN1_ID, DAG1_RUN2_ID], id="order_by_start_date"), pytest.param("end_date", [DAG1_RUN1_ID, DAG1_RUN2_ID], id="order_by_end_date"), @@ -970,6 +994,48 @@ def test_filters(self, test_client, dag_id, query_params, expected_dag_id_list): body = response.json() assert [each["dag_run_id"] for each in body["dag_runs"]] == expected_dag_id_list + @pytest.mark.usefixtures("configure_git_connection_for_dag_bundle") + def test_partition_date_day_filters_use_timetable_timezone(self, test_client, dag_maker, session): + dag_id = "test_partition_date_local_day" + with dag_maker( + dag_id=dag_id, + schedule=CronPartitionTimetable("0 0 * * *", timezone="Asia/Taipei"), + start_date=START_DATE1, + session=session, + serialized=True, + ): + EmptyOperator(task_id="task") + + for run_id, partition_date in [ + ("before_local_day", datetime(2025, 1, 1, 15, 59, 59, tzinfo=timezone.utc)), + ("start_local_day", datetime(2025, 1, 1, 16, 0, 0, tzinfo=timezone.utc)), + ("end_local_day", datetime(2025, 1, 2, 15, 59, 59, tzinfo=timezone.utc)), + ("after_local_day", datetime(2025, 1, 2, 16, 0, 0, tzinfo=timezone.utc)), + ]: + dag_maker.create_dagrun( + run_id=run_id, + state=DagRunState.SUCCESS, + logical_date=None, + partition_date=partition_date, + partition_key=run_id, + ) + dag_maker.sync_dagbag_to_db() + session.commit() + + response = test_client.get( + f"/dags/{dag_id}/dagRuns", + params={ + "partition_date_gte": "2025-01-02", + "partition_date_lte": "2025-01-02", + "order_by": "partition_date", + }, + ) + assert response.status_code == 200 + assert [each["dag_run_id"] for each in response.json()["dag_runs"]] == [ + "start_local_day", + "end_local_day", + ] + def test_bad_filters(self, test_client): query_params = { "logical_date_gte": "invalid",