AIP-84 Add filter by dag_id in recent_dag_runs (apache#44737)

astronomer · Dec 11, 2024 · f0ec10e · f0ec10e
1 parent a77fca2
commit f0ec10e
Show file tree

Hide file tree

Showing 9 changed files with 44 additions and 0 deletions.
diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
@@ -101,6 +101,16 @@ paths:
           items:
             type: string
           title: Owners
+      - name: dag_ids
+        in: query
+        required: false
+        schema:
+          anyOf:
+          - type: array
+            items:
+              type: string
+          - type: 'null'
+          title: Dag Ids
       - name: dag_id_pattern
         in: query
         required: false

diff --git a/airflow/api_fastapi/core_api/routes/ui/dags.py b/airflow/api_fastapi/core_api/routes/ui/dags.py
@@ -17,13 +17,18 @@
 
 from __future__ import annotations
 
+from typing import Annotated
+
+from fastapi import Depends
 from sqlalchemy import and_, func, select
 
 from airflow.api_fastapi.common.db.common import (
     SessionDep,
     paginated_select,
 )
 from airflow.api_fastapi.common.parameters import (
+    FilterOptionEnum,
+    FilterParam,
     QueryDagDisplayNamePatternSearch,
     QueryDagIdPatternSearch,
     QueryLastDagRunStateFilter,
@@ -33,6 +38,7 @@
     QueryOwnersFilter,
     QueryPausedFilter,
     QueryTagsFilter,
+    filter_param_factory,
 )
 from airflow.api_fastapi.common.router import AirflowRouter
 from airflow.api_fastapi.core_api.datamodels.dag_run import DAGRunResponse
@@ -52,6 +58,10 @@ def recent_dag_runs(
     offset: QueryOffset,
     tags: QueryTagsFilter,
     owners: QueryOwnersFilter,
+    dag_ids: Annotated[
+        FilterParam[list[str] | None],
+        Depends(filter_param_factory(DagRun.dag_id, list[str] | None, FilterOptionEnum.IN, "dag_ids")),
+    ],
     dag_id_pattern: QueryDagIdPatternSearch,
     dag_display_name_pattern: QueryDagDisplayNamePatternSearch,
     only_active: QueryOnlyActiveFilter,
@@ -104,6 +114,7 @@ def recent_dag_runs(
             only_active,
             paused,
             dag_id_pattern,
+            dag_ids,
             dag_display_name_pattern,
             tags,
             owners,

diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts
@@ -312,6 +312,7 @@ export const UseDagsServiceRecentDagRunsKeyFn = (
   {
     dagDisplayNamePattern,
     dagIdPattern,
+    dagIds,
     dagRunsLimit,
     lastDagRunState,
     limit,
@@ -323,6 +324,7 @@ export const UseDagsServiceRecentDagRunsKeyFn = (
   }: {
     dagDisplayNamePattern?: string;
     dagIdPattern?: string;
+    dagIds?: string[];
     dagRunsLimit?: number;
     lastDagRunState?: DagRunState;
     limit?: number;
@@ -339,6 +341,7 @@ export const UseDagsServiceRecentDagRunsKeyFn = (
     {
       dagDisplayNamePattern,
       dagIdPattern,
+      dagIds,
       dagRunsLimit,
       lastDagRunState,
       limit,

diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -389,6 +389,7 @@ export const prefetchUseConfigServiceGetConfigValue = (
  * @param data.offset
  * @param data.tags
  * @param data.owners
+ * @param data.dagIds
  * @param data.dagIdPattern
  * @param data.dagDisplayNamePattern
  * @param data.onlyActive
@@ -402,6 +403,7 @@ export const prefetchUseDagsServiceRecentDagRuns = (
   {
     dagDisplayNamePattern,
     dagIdPattern,
+    dagIds,
     dagRunsLimit,
     lastDagRunState,
     limit,
@@ -413,6 +415,7 @@ export const prefetchUseDagsServiceRecentDagRuns = (
   }: {
     dagDisplayNamePattern?: string;
     dagIdPattern?: string;
+    dagIds?: string[];
     dagRunsLimit?: number;
     lastDagRunState?: DagRunState;
     limit?: number;
@@ -427,6 +430,7 @@ export const prefetchUseDagsServiceRecentDagRuns = (
     queryKey: Common.UseDagsServiceRecentDagRunsKeyFn({
       dagDisplayNamePattern,
       dagIdPattern,
+      dagIds,
       dagRunsLimit,
       lastDagRunState,
       limit,
@@ -440,6 +444,7 @@ export const prefetchUseDagsServiceRecentDagRuns = (
       DagsService.recentDagRuns({
         dagDisplayNamePattern,
         dagIdPattern,
+        dagIds,
         dagRunsLimit,
         lastDagRunState,
         limit,

diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts
@@ -497,6 +497,7 @@ export const useConfigServiceGetConfigValue = <
  * @param data.offset
  * @param data.tags
  * @param data.owners
+ * @param data.dagIds
  * @param data.dagIdPattern
  * @param data.dagDisplayNamePattern
  * @param data.onlyActive
@@ -513,6 +514,7 @@ export const useDagsServiceRecentDagRuns = <
   {
     dagDisplayNamePattern,
     dagIdPattern,
+    dagIds,
     dagRunsLimit,
     lastDagRunState,
     limit,
@@ -524,6 +526,7 @@ export const useDagsServiceRecentDagRuns = <
   }: {
     dagDisplayNamePattern?: string;
     dagIdPattern?: string;
+    dagIds?: string[];
     dagRunsLimit?: number;
     lastDagRunState?: DagRunState;
     limit?: number;
@@ -541,6 +544,7 @@ export const useDagsServiceRecentDagRuns = <
       {
         dagDisplayNamePattern,
         dagIdPattern,
+        dagIds,
         dagRunsLimit,
         lastDagRunState,
         limit,
@@ -556,6 +560,7 @@ export const useDagsServiceRecentDagRuns = <
       DagsService.recentDagRuns({
         dagDisplayNamePattern,
         dagIdPattern,
+        dagIds,
         dagRunsLimit,
         lastDagRunState,
         limit,

diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts
@@ -472,6 +472,7 @@ export const useConfigServiceGetConfigValueSuspense = <
  * @param data.offset
  * @param data.tags
  * @param data.owners
+ * @param data.dagIds
  * @param data.dagIdPattern
  * @param data.dagDisplayNamePattern
  * @param data.onlyActive
@@ -488,6 +489,7 @@ export const useDagsServiceRecentDagRunsSuspense = <
   {
     dagDisplayNamePattern,
     dagIdPattern,
+    dagIds,
     dagRunsLimit,
     lastDagRunState,
     limit,
@@ -499,6 +501,7 @@ export const useDagsServiceRecentDagRunsSuspense = <
   }: {
     dagDisplayNamePattern?: string;
     dagIdPattern?: string;
+    dagIds?: string[];
     dagRunsLimit?: number;
     lastDagRunState?: DagRunState;
     limit?: number;
@@ -516,6 +519,7 @@ export const useDagsServiceRecentDagRunsSuspense = <
       {
         dagDisplayNamePattern,
         dagIdPattern,
+        dagIds,
         dagRunsLimit,
         lastDagRunState,
         limit,
@@ -531,6 +535,7 @@ export const useDagsServiceRecentDagRunsSuspense = <
       DagsService.recentDagRuns({
         dagDisplayNamePattern,
         dagIdPattern,
+        dagIds,
         dagRunsLimit,
         lastDagRunState,
         limit,

diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -668,6 +668,7 @@ export class DagsService {
    * @param data.offset
    * @param data.tags
    * @param data.owners
+   * @param data.dagIds
    * @param data.dagIdPattern
    * @param data.dagDisplayNamePattern
    * @param data.onlyActive
@@ -688,6 +689,7 @@ export class DagsService {
         offset: data.offset,
         tags: data.tags,
         owners: data.owners,
+        dag_ids: data.dagIds,
         dag_id_pattern: data.dagIdPattern,
         dag_display_name_pattern: data.dagDisplayNamePattern,
         only_active: data.onlyActive,

diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -1421,6 +1421,7 @@ export type GetConfigValueResponse = Config;
 export type RecentDagRunsData = {
   dagDisplayNamePattern?: string | null;
   dagIdPattern?: string | null;
+  dagIds?: Array<string> | null;
   dagRunsLimit?: number;
   lastDagRunState?: DagRunState | null;
   limit?: number;

diff --git a/tests/api_fastapi/core_api/routes/ui/test_dags.py b/tests/api_fastapi/core_api/routes/ui/test_dags.py
@@ -73,6 +73,8 @@ def setup_dag_runs(self, session=None) -> None:
             ({"paused": False}, [DAG1_ID, DAG2_ID], 11),
             ({"owners": ["airflow"]}, [DAG1_ID, DAG2_ID], 11),
             ({"owners": ["test_owner"], "only_active": False}, [DAG3_ID], 4),
+            ({"dag_ids": [DAG1_ID]}, [DAG1_ID], 6),
+            ({"dag_ids": [DAG1_ID, DAG2_ID]}, [DAG1_ID, DAG2_ID], 11),
             ({"last_dag_run_state": "success", "only_active": False}, [DAG1_ID, DAG2_ID, DAG3_ID], 6),
             ({"last_dag_run_state": "failed", "only_active": False}, [DAG1_ID, DAG2_ID, DAG3_ID], 9),
             # Search