Refactor Get Event Logs with filter_param_factory

apache · Nov 13, 2024 · cabb68d · cabb68d
1 parent 48e7385
commit cabb68d
Show file tree

Hide file tree

Showing 6 changed files with 79 additions and 72 deletions.
diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml
@@ -1554,6 +1554,27 @@ paths:
       description: Get all Event Logs.
       operationId: get_event_logs
       parameters:
+      - name: limit
+        in: query
+        required: false
+        schema:
+          type: integer
+          default: 100
+          title: Limit
+      - name: offset
+        in: query
+        required: false
+        schema:
+          type: integer
+          default: 0
+          title: Offset
+      - name: order_by
+        in: query
+        required: false
+        schema:
+          type: string
+          default: id
+          title: Order By
       - name: dag_id
         in: query
         required: false
@@ -1648,27 +1669,6 @@ paths:
             format: date-time
           - type: 'null'
           title: After
-      - name: limit
-        in: query
-        required: false
-        schema:
-          type: integer
-          default: 100
-          title: Limit
-      - name: offset
-        in: query
-        required: false
-        schema:
-          type: integer
-          default: 0
-          title: Offset
-      - name: order_by
-        in: query
-        required: false
-        schema:
-          type: string
-          default: id
-          title: Order By
       responses:
         '200':
           description: Successful Response

diff --git a/airflow/api_fastapi/core_api/routes/public/event_logs.py b/airflow/api_fastapi/core_api/routes/public/event_logs.py
@@ -19,7 +19,7 @@
 from datetime import datetime
 from typing import Annotated
 
-from fastapi import Depends, HTTPException, Query, status
+from fastapi import Depends, HTTPException, status
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
@@ -28,9 +28,15 @@
     paginated_select,
 )
 from airflow.api_fastapi.common.parameters import (
+    FilterOptionEnum,
+    FilterParam,
     QueryLimit,
     QueryOffset,
     SortParam,
+    datetime_filter_param_factory,
+    int_filter_param_factory,
+    str_filter_param_factory,
+    str_list_filter_param_factory,
 )
 from airflow.api_fastapi.common.router import AirflowRouter
 from airflow.api_fastapi.core_api.datamodels.event_logs import (
@@ -90,46 +96,47 @@ def get_event_logs(
             ).dynamic_depends()
         ),
     ],
-    dag_id: str | None = None,
-    task_id: str | None = None,
-    run_id: str | None = None,
-    map_index: int | None = None,
-    try_number: int | None = None,
-    owner: str | None = None,
-    event: str | None = None,
-    excluded_events: list[str] | None = Query(None),
-    included_events: list[str] | None = Query(None),
-    before: datetime | None = None,
-    after: datetime | None = None,
+    dag_id: Annotated[FilterParam[str | None], Depends(str_filter_param_factory(Log.dag_id))],
+    task_id: Annotated[FilterParam[str | None], Depends(str_filter_param_factory(Log.task_id))],
+    run_id: Annotated[FilterParam[str | None], Depends(str_filter_param_factory(Log.run_id))],
+    map_index: Annotated[FilterParam[int | None], Depends(int_filter_param_factory(Log.map_index))],
+    try_number: Annotated[FilterParam[int | None], Depends(int_filter_param_factory(Log.try_number))],
+    owner: Annotated[FilterParam[str | None], Depends(str_filter_param_factory(Log.owner))],
+    event: Annotated[FilterParam[str | None], Depends(str_filter_param_factory(Log.event))],
+    excluded_events: Annotated[
+        FilterParam[list[str] | None],
+        Depends(str_list_filter_param_factory(Log.event, FilterOptionEnum.NOT_IN, "excluded_events")),
+    ],
+    included_events: Annotated[
+        FilterParam[list[str] | None],
+        Depends(str_list_filter_param_factory(Log.event, FilterOptionEnum.IN, "included_events")),
+    ],
+    before: Annotated[
+        FilterParam[datetime | None],
+        Depends(datetime_filter_param_factory(Log.dttm, FilterOptionEnum.LESS_THAN, "before")),
+    ],
+    after: Annotated[
+        FilterParam[datetime | None],
+        Depends(datetime_filter_param_factory(Log.dttm, FilterOptionEnum.GREATER_THAN, "after")),
+    ],
 ) -> EventLogCollectionResponse:
     """Get all Event Logs."""
     base_select = select(Log).group_by(Log.id)
-    # TODO: Refactor using the `FilterParam` class in commit `574b72e41cc5ed175a2bbf4356522589b836bb11`
-    if dag_id is not None:
-        base_select = base_select.where(Log.dag_id == dag_id)
-    if task_id is not None:
-        base_select = base_select.where(Log.task_id == task_id)
-    if run_id is not None:
-        base_select = base_select.where(Log.run_id == run_id)
-    if map_index is not None:
-        base_select = base_select.where(Log.map_index == map_index)
-    if try_number is not None:
-        base_select = base_select.where(Log.try_number == try_number)
-    if owner is not None:
-        base_select = base_select.where(Log.owner == owner)
-    if event is not None:
-        base_select = base_select.where(Log.event == event)
-    if excluded_events is not None:
-        base_select = base_select.where(Log.event.notin_(excluded_events))
-    if included_events is not None:
-        base_select = base_select.where(Log.event.in_(included_events))
-    if before is not None:
-        base_select = base_select.where(Log.dttm < before)
-    if after is not None:
-        base_select = base_select.where(Log.dttm > after)
     event_logs_select, total_entries = paginated_select(
         base_select,
-        [],
+        [
+            dag_id,
+            task_id,
+            run_id,
+            map_index,
+            try_number,
+            owner,
+            event,
+            excluded_events,
+            included_events,
+            before,
+            after,
+        ],
         order_by,
         offset,
         limit,

diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts
@@ -499,6 +499,9 @@ export const prefetchUseEventLogServiceGetEventLog = (
  * Get Event Logs
  * Get all Event Logs.
  * @param data The data for the request.
+ * @param data.limit
+ * @param data.offset
+ * @param data.orderBy
  * @param data.dagId
  * @param data.taskId
  * @param data.runId
@@ -510,9 +513,6 @@ export const prefetchUseEventLogServiceGetEventLog = (
  * @param data.includedEvents
  * @param data.before
  * @param data.after
- * @param data.limit
- * @param data.offset
- * @param data.orderBy
  * @returns EventLogCollectionResponse Successful Response
  * @throws ApiError
  */

diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts
@@ -628,6 +628,9 @@ export const useEventLogServiceGetEventLog = <
  * Get Event Logs
  * Get all Event Logs.
  * @param data The data for the request.
+ * @param data.limit
+ * @param data.offset
+ * @param data.orderBy
  * @param data.dagId
  * @param data.taskId
  * @param data.runId
@@ -639,9 +642,6 @@ export const useEventLogServiceGetEventLog = <
  * @param data.includedEvents
  * @param data.before
  * @param data.after
- * @param data.limit
- * @param data.offset
- * @param data.orderBy
  * @returns EventLogCollectionResponse Successful Response
  * @throws ApiError
  */

diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts
@@ -613,6 +613,9 @@ export const useEventLogServiceGetEventLogSuspense = <
  * Get Event Logs
  * Get all Event Logs.
  * @param data The data for the request.
+ * @param data.limit
+ * @param data.offset
+ * @param data.orderBy
  * @param data.dagId
  * @param data.taskId
  * @param data.runId
@@ -624,9 +627,6 @@ export const useEventLogServiceGetEventLogSuspense = <
  * @param data.includedEvents
  * @param data.before
  * @param data.after
- * @param data.limit
- * @param data.offset
- * @param data.orderBy
  * @returns EventLogCollectionResponse Successful Response
  * @throws ApiError
  */

diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts
@@ -962,6 +962,9 @@ export class EventLogService {
    * Get Event Logs
    * Get all Event Logs.
    * @param data The data for the request.
+   * @param data.limit
+   * @param data.offset
+   * @param data.orderBy
    * @param data.dagId
    * @param data.taskId
    * @param data.runId
@@ -973,9 +976,6 @@ export class EventLogService {
    * @param data.includedEvents
    * @param data.before
    * @param data.after
-   * @param data.limit
-   * @param data.offset
-   * @param data.orderBy
    * @returns EventLogCollectionResponse Successful Response
    * @throws ApiError
    */
@@ -986,6 +986,9 @@ export class EventLogService {
       method: "GET",
       url: "/public/eventLogs/",
       query: {
+        limit: data.limit,
+        offset: data.offset,
+        order_by: data.orderBy,
         dag_id: data.dagId,
         task_id: data.taskId,
         run_id: data.runId,
@@ -997,9 +1000,6 @@ export class EventLogService {
         included_events: data.includedEvents,
         before: data.before,
         after: data.after,
-        limit: data.limit,
-        offset: data.offset,
-        order_by: data.orderBy,
       },
       errors: {
         401: "Unauthorized",