refactor: fix SortParm for connections, dag api

apache · Oct 18, 2024 · 2f1cb76 · 2f1cb76
1 parent 11fa7bc
commit 2f1cb76
Show file tree

Hide file tree

Showing 9 changed files with 63 additions and 29 deletions.
diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
@@ -299,12 +299,6 @@ paths:
       description: Get all DAG tags.
       operationId: get_dag_tags
       parameters:
-      - name: order_by
-        in: query
-        required: true
-        schema:
-          type: 'null'
-          title: Order By
       - name: limit
         in: query
         required: false
@@ -319,6 +313,13 @@ paths:
           type: integer
           default: 0
           title: Offset
+      - name: order_by
+        in: query
+        required: false
+        schema:
+          type: string
+          default: name
+          title: Order By
       - name: tag_name_pattern
         in: query
         required: false
@@ -334,6 +335,12 @@ paths:
             application/json:
               schema:
                 $ref: '#/components/schemas/DAGTagCollectionResponse'
+        '400':
+          content:
+            application/json:
+              schema:
+                $ref: '#/components/schemas/HTTPExceptionResponse'
+          description: Bad Request
         '401':
           content:
             application/json:

diff --git a/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow/api_fastapi/core_api/routes/public/connections.py
@@ -80,7 +80,15 @@ async def get_connections(
         SortParam,
         Depends(
             SortParam(
-                ["connection_id", "conn_type", "description", "host", "port", "id"], Connection
+                {
+                    "connection_id": Connection.conn_id,
+                    "conn_type": Connection.conn_type,
+                    "description": Connection.description,
+                    "host": Connection.host,
+                    "port": Connection.port,
+                    "id": Connection.id,
+                },
+                Connection,
             ).dynamic_depends()
         ),
     ],

diff --git a/airflow/api_fastapi/core_api/routes/public/dags.py b/airflow/api_fastapi/core_api/routes/public/dags.py
@@ -18,7 +18,7 @@
 from __future__ import annotations
 
 from fastapi import Depends, HTTPException, Query, Request, Response
-from sqlalchemy import distinct, select, update
+from sqlalchemy import select, update
 from sqlalchemy.orm import Session
 from typing_extensions import Annotated
 
@@ -32,7 +32,6 @@
     QueryDagDisplayNamePatternSearch,
     QueryDagIdPatternSearch,
     QueryDagIdPatternSearchWithNone,
-    QueryDagTagOrderBy,
     QueryDagTagPatternSearch,
     QueryLastDagRunStateFilter,
     QueryLimit,
@@ -53,7 +52,7 @@
     DAGTagCollectionResponse,
 )
 from airflow.exceptions import AirflowException, DagNotFound
-from airflow.models import DAG, DagModel, DagTag
+from airflow.models import DAG, DagModel, DagRun, DagTag
 
 dags_router = AirflowRouter(tags=["DAG"], prefix="/dags")
 
@@ -73,7 +72,13 @@ async def get_dags(
         SortParam,
         Depends(
             SortParam(
-                ["dag_id", "dag_display_name", "next_dagrun", "last_run_state", "last_run_start_date"],
+                {
+                    "dag_id": DagModel.dag_id,
+                    "dag_display_name": DagModel.dag_display_name,
+                    "next_dagrun": DagModel.next_dagrun,
+                    "last_run_state": DagRun.state,
+                    "last_run_start_date": DagRun.start_date,
+                },
                 DagModel,
             ).dynamic_depends()
         ),
@@ -100,17 +105,26 @@ async def get_dags(
 
 @dags_router.get(
     "/tags",
-    responses=create_openapi_http_exception_doc([401, 403]),
+    responses=create_openapi_http_exception_doc([400, 401, 403]),
 )
 async def get_dag_tags(
     limit: QueryLimit,
     offset: QueryOffset,
-    order_by: QueryDagTagOrderBy,
+    order_by: Annotated[
+        SortParam,
+        Depends(
+            SortParam(
+                {"name": DagTag.name},
+                DagTag,
+                "name",
+            ).dynamic_depends()
+        ),
+    ],
     tag_name_pattern: QueryDagTagPatternSearch,
     session: Annotated[Session, Depends(get_session)],
 ) -> DAGTagCollectionResponse:
     """Get all DAG tags."""
-    base_select = select(distinct(DagTag.name))
+    base_select = select(DagTag.name).group_by(DagTag.name)
     dag_tags_select, total_entries = paginated_select(
         base_select=base_select,
         filters=[tag_name_pattern],

diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts
@@ -117,9 +117,9 @@ export const UseDagServiceGetDagTagsKeyFn = (
   }: {
     limit?: number;
     offset?: number;
-    orderBy: null;
+    orderBy?: string;
     tagNamePattern?: string;
-  },
+  } = {},
   queryKey?: Array<unknown>,
 ) => [
   useDagServiceGetDagTagsKey,

diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -132,9 +132,9 @@ export const prefetchUseDagServiceGetDags = (
  * Get Dag Tags
  * Get all DAG tags.
  * @param data The data for the request.
- * @param data.orderBy
  * @param data.limit
  * @param data.offset
+ * @param data.orderBy
  * @param data.tagNamePattern
  * @returns DAGTagCollectionResponse Successful Response
  * @throws ApiError
@@ -149,9 +149,9 @@ export const prefetchUseDagServiceGetDagTags = (
   }: {
     limit?: number;
     offset?: number;
-    orderBy: null;
+    orderBy?: string;
     tagNamePattern?: string;
-  },
+  } = {},
 ) =>
   queryClient.prefetchQuery({
     queryKey: Common.UseDagServiceGetDagTagsKeyFn({

diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts
@@ -159,9 +159,9 @@ export const useDagServiceGetDags = <
  * Get Dag Tags
  * Get all DAG tags.
  * @param data The data for the request.
- * @param data.orderBy
  * @param data.limit
  * @param data.offset
+ * @param data.orderBy
  * @param data.tagNamePattern
  * @returns DAGTagCollectionResponse Successful Response
  * @throws ApiError
@@ -179,9 +179,9 @@ export const useDagServiceGetDagTags = <
   }: {
     limit?: number;
     offset?: number;
-    orderBy: null;
+    orderBy?: string;
     tagNamePattern?: string;
-  },
+  } = {},
   queryKey?: TQueryKey,
   options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">,
 ) =>

diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts
@@ -154,9 +154,9 @@ export const useDagServiceGetDagsSuspense = <
  * Get Dag Tags
  * Get all DAG tags.
  * @param data The data for the request.
- * @param data.orderBy
  * @param data.limit
  * @param data.offset
+ * @param data.orderBy
  * @param data.tagNamePattern
  * @returns DAGTagCollectionResponse Successful Response
  * @throws ApiError
@@ -174,9 +174,9 @@ export const useDagServiceGetDagTagsSuspense = <
   }: {
     limit?: number;
     offset?: number;
-    orderBy: null;
+    orderBy?: string;
     tagNamePattern?: string;
-  },
+  } = {},
   queryKey?: TQueryKey,
   options?: Omit<UseQueryOptions<TData, TError>, "queryKey" | "queryFn">,
 ) =>

diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -188,26 +188,27 @@ export class DagService {
    * Get Dag Tags
    * Get all DAG tags.
    * @param data The data for the request.
-   * @param data.orderBy
    * @param data.limit
    * @param data.offset
+   * @param data.orderBy
    * @param data.tagNamePattern
    * @returns DAGTagCollectionResponse Successful Response
    * @throws ApiError
    */
   public static getDagTags(
-    data: GetDagTagsData,
+    data: GetDagTagsData = {},
   ): CancelablePromise<GetDagTagsResponse> {
     return __request(OpenAPI, {
       method: "GET",
       url: "/public/dags/tags",
       query: {
-        order_by: data.orderBy,
         limit: data.limit,
         offset: data.offset,
+        order_by: data.orderBy,
         tag_name_pattern: data.tagNamePattern,
       },
       errors: {
+        400: "Bad Request",
         401: "Unauthorized",
         403: "Forbidden",
         422: "Validation Error",

diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts
@@ -381,7 +381,7 @@ export type PatchDagsResponse = DAGCollectionResponse;
 export type GetDagTagsData = {
   limit?: number;
   offset?: number;
-  orderBy: null;
+  orderBy?: string;
   tagNamePattern?: string | null;
 };
 
@@ -572,6 +572,10 @@ export type $OpenApiTs = {
          * Successful Response
          */
         200: DAGTagCollectionResponse;
+        /**
+         * Bad Request
+         */
+        400: HTTPExceptionResponse;
         /**
          * Unauthorized
          */