From cf5218ccf98d694101ba1d37e9fe88ad83155e9e Mon Sep 17 00:00:00 2001 From: Amogh Date: Wed, 6 Nov 2024 15:22:45 +0530 Subject: [PATCH 01/19] AIP-84: Migrating GET Assets to fastAPI --- .../api_connexion/endpoints/asset_endpoint.py | 2 + airflow/api_fastapi/common/parameters.py | 15 ++++ .../core_api/openapi/v1-generated.yaml | 27 +++++++ .../core_api/routes/public/__init__.py | 2 + .../core_api/routes/public/assets.py | 72 +++++++++++++++++++ .../core_api/serializers/assets.py | 64 +++++++++++++++++ airflow/ui/openapi-gen/queries/common.ts | 16 +++++ airflow/ui/openapi-gen/queries/prefetch.ts | 19 +++++ airflow/ui/openapi-gen/queries/queries.ts | 25 +++++++ airflow/ui/openapi-gen/queries/suspense.ts | 25 +++++++ .../ui/openapi-gen/requests/services.gen.ts | 24 +++++++ airflow/ui/openapi-gen/requests/types.gen.ts | 25 +++++++ 12 files changed, 316 insertions(+) create mode 100644 airflow/api_fastapi/core_api/routes/public/assets.py create mode 100644 airflow/api_fastapi/core_api/serializers/assets.py diff --git a/airflow/api_connexion/endpoints/asset_endpoint.py b/airflow/api_connexion/endpoints/asset_endpoint.py index 1ea1db2b3bbb8..085817213d0a4 100644 --- a/airflow/api_connexion/endpoints/asset_endpoint.py +++ b/airflow/api_connexion/endpoints/asset_endpoint.py @@ -47,6 +47,7 @@ from airflow.assets.manager import asset_manager from airflow.models.asset import AssetDagRunQueue, AssetEvent, AssetModel from airflow.utils import timezone +from airflow.utils.api_migration import mark_fastapi_migration_done from airflow.utils.db import get_query_count from airflow.utils.session import NEW_SESSION, provide_session from airflow.www.decorators import action_logging @@ -77,6 +78,7 @@ def get_asset(*, uri: str, session: Session = NEW_SESSION) -> APIResponse: return asset_schema.dump(asset) +@mark_fastapi_migration_done @security.requires_access_asset("GET") @format_parameters({"limit": check_limit}) @provide_session diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index 64ae9406f08c3..224f6e0b7e775 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -29,6 +29,7 @@ from typing_extensions import Annotated, Self from airflow.models import Base, Connection +from airflow.models.asset import AssetModel from airflow.models.dag import DagModel, DagTag from airflow.models.dagrun import DagRun from airflow.models.dagwarning import DagWarning, DagWarningType @@ -335,6 +336,17 @@ def depends(self, dag_id: str | None = None) -> _DagIdFilter: return self.set_value(dag_id) +class _UriPatternSearch(_SearchParam): + """Search on dag_id.""" + + def __init__(self, skip_none: bool = True) -> None: + super().__init__(AssetModel.uri, skip_none) + + def depends(self, uri_pattern: str | None = None) -> _UriPatternSearch: + uri_pattern = super().transform_aliases(uri_pattern) + return self.set_value(uri_pattern) + + # Common Safe DateTime DateTimeQuery = Annotated[str, AfterValidator(_safe_parse_datetime)] @@ -363,3 +375,6 @@ def depends(self, dag_id: str | None = None) -> _DagIdFilter: # DAGTags QueryDagTagPatternSearch = Annotated[_DagTagNamePatternSearch, Depends(_DagTagNamePatternSearch().depends)] + +# Assets +QueryUriPatternSearch = Annotated[_UriPatternSearch, Depends(_UriPatternSearch().depends)] diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 46a2be61d64a0..4974d9540bc7d 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2515,6 +2515,33 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' + /public/next_run_assets/{dag_id}: + get: + tags: + - Asset + summary: Next Run Assets + operationId: next_run_assets + parameters: + - name: dag_id + in: path + required: true + schema: + type: string + title: Dag Id + responses: + '200': + description: Successful Response + content: + application/json: + schema: + type: object + title: Response Next Run Assets + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' components: schemas: AppBuilderMenuItemResponse: diff --git a/airflow/api_fastapi/core_api/routes/public/__init__.py b/airflow/api_fastapi/core_api/routes/public/__init__.py index b7c8affe4a9cb..8359a8b679298 100644 --- a/airflow/api_fastapi/core_api/routes/public/__init__.py +++ b/airflow/api_fastapi/core_api/routes/public/__init__.py @@ -34,6 +34,7 @@ from airflow.api_fastapi.core_api.routes.public.task_instances import task_instances_router from airflow.api_fastapi.core_api.routes.public.variables import variables_router from airflow.api_fastapi.core_api.routes.public.version import version_router +from airflow.api_fastapi.core_api.routes.ui import assets_router public_router = AirflowRouter(prefix="/public") @@ -56,3 +57,4 @@ public_router.include_router(variables_router) public_router.include_router(version_router) public_router.include_router(dag_stats_router) +public_router.include_router(assets_router) diff --git a/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow/api_fastapi/core_api/routes/public/assets.py new file mode 100644 index 0000000000000..a90741dfd4479 --- /dev/null +++ b/airflow/api_fastapi/core_api/routes/public/assets.py @@ -0,0 +1,72 @@ +# 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. + +from __future__ import annotations + +from typing import Annotated + +from fastapi import Depends +from sqlalchemy import select +from sqlalchemy.orm import Session + +from airflow.api_fastapi.common.db.common import get_session, paginated_select +from airflow.api_fastapi.common.parameters import ( + QueryDagIdsFilter, + QueryLimit, + QueryOffset, + QueryUriPatternSearch, + SortParam, +) +from airflow.api_fastapi.common.router import AirflowRouter +from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.core_api.serializers.assets import AssetCollectionResponse, AssetResponse +from airflow.assets import Asset +from airflow.models.asset import AssetModel + +assets_router = AirflowRouter(tags=["Asset"], prefix="/assets") + + +@assets_router.get( + "/", + responses=create_openapi_http_exception_doc([401, 403, 404]), +) +async def get_assets( + limit: QueryLimit, + offset: QueryOffset, + uri_pattern: QueryUriPatternSearch, + dag_ids: QueryDagIdsFilter, + order_by: Annotated[ + SortParam, + Depends(SortParam(["id", "uri", "created_at", "updated_at"], Asset).dynamic_depends()), + ], + session: Annotated[Session, Depends(get_session)], +) -> AssetCollectionResponse: + """Get assets.""" + assets_select, total_entries = paginated_select( + select(AssetModel), + filters=[dag_ids, uri_pattern], + order_by=order_by, + offset=offset, + limit=limit, + session=session, + ) + + assets = session.scalars(assets_select).all() + return AssetCollectionResponse( + assets=[AssetResponse.model_validate(x, from_attributes=True) for x in assets], + total_entries=total_entries, + ) diff --git a/airflow/api_fastapi/core_api/serializers/assets.py b/airflow/api_fastapi/core_api/serializers/assets.py new file mode 100644 index 0000000000000..7acb8cc6357dc --- /dev/null +++ b/airflow/api_fastapi/core_api/serializers/assets.py @@ -0,0 +1,64 @@ +# 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. + +from __future__ import annotations + +from datetime import datetime + +from pydantic import BaseModel, ConfigDict, Field + + +class DagScheduleAssetReference(BaseModel): + """Serializable version of the DagScheduleAssetReference ORM SqlAlchemyModel.""" + + asset_id: int + dag_id: str + created_at: datetime + updated_at: datetime + + model_config = ConfigDict(from_attributes=True) + + +class TaskOutletAssetReference(BaseModel): + """Serializable version of the TaskOutletAssetReference ORM SqlAlchemyModel.""" + + asset_id: int + dag_id: str + task_id: str + created_at: datetime + updated_at: datetime + + model_config = ConfigDict(from_attributes=True) + + +class AssetResponse(BaseModel): + """Asset serializer for responses.""" + + id: int + uri: str + extra: str | None = Field(default=None) + created_at: str + updated_at: str + consuming_dags: DagScheduleAssetReference | None = None + producing_tasks: TaskOutletAssetReference | None = None + + +class AssetCollectionResponse(BaseModel): + """Asset collection response.""" + + assets: list[AssetResponse] + total_entries: int diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index cec1f0f314dc7..27b9d5c59a629 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -40,6 +40,22 @@ export const UseAssetServiceNextRunAssetsKeyFn = ( }, queryKey?: Array, ) => [useAssetServiceNextRunAssetsKey, ...(queryKey ?? [{ dagId }])]; +export type AssetServiceNextRunAssets1DefaultResponse = Awaited< + ReturnType +>; +export type AssetServiceNextRunAssets1QueryResult< + TData = AssetServiceNextRunAssets1DefaultResponse, + TError = unknown, +> = UseQueryResult; +export const useAssetServiceNextRunAssets1Key = "AssetServiceNextRunAssets1"; +export const UseAssetServiceNextRunAssets1KeyFn = ( + { + dagId, + }: { + dagId: string; + }, + queryKey?: Array, +) => [useAssetServiceNextRunAssets1Key, ...(queryKey ?? [{ dagId }])]; export type DashboardServiceHistoricalMetricsDefaultResponse = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index 04443427aca70..c1b00fa352a20 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -44,6 +44,25 @@ export const prefetchUseAssetServiceNextRunAssets = ( queryKey: Common.UseAssetServiceNextRunAssetsKeyFn({ dagId }), queryFn: () => AssetService.nextRunAssets({ dagId }), }); +/** + * Next Run Assets + * @param data The data for the request. + * @param data.dagId + * @returns unknown Successful Response + * @throws ApiError + */ +export const prefetchUseAssetServiceNextRunAssets1 = ( + queryClient: QueryClient, + { + dagId, + }: { + dagId: string; + }, +) => + queryClient.prefetchQuery({ + queryKey: Common.UseAssetServiceNextRunAssets1KeyFn({ dagId }), + queryFn: () => AssetService.nextRunAssets1({ dagId }), + }); /** * Historical Metrics * Return cluster activity historical metrics. diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 11dea6f3df589..c3057b19e5997 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -65,6 +65,31 @@ export const useAssetServiceNextRunAssets = < queryFn: () => AssetService.nextRunAssets({ dagId }) as TData, ...options, }); +/** + * Next Run Assets + * @param data The data for the request. + * @param data.dagId + * @returns unknown Successful Response + * @throws ApiError + */ +export const useAssetServiceNextRunAssets1 = < + TData = Common.AssetServiceNextRunAssets1DefaultResponse, + TError = unknown, + TQueryKey extends Array = unknown[], +>( + { + dagId, + }: { + dagId: string; + }, + queryKey?: TQueryKey, + options?: Omit, "queryKey" | "queryFn">, +) => + useQuery({ + queryKey: Common.UseAssetServiceNextRunAssets1KeyFn({ dagId }, queryKey), + queryFn: () => AssetService.nextRunAssets1({ dagId }) as TData, + ...options, + }); /** * Historical Metrics * Return cluster activity historical metrics. diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index eed1a0afe8057..033b4cdb513c8 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -50,6 +50,31 @@ export const useAssetServiceNextRunAssetsSuspense = < queryFn: () => AssetService.nextRunAssets({ dagId }) as TData, ...options, }); +/** + * Next Run Assets + * @param data The data for the request. + * @param data.dagId + * @returns unknown Successful Response + * @throws ApiError + */ +export const useAssetServiceNextRunAssets1Suspense = < + TData = Common.AssetServiceNextRunAssets1DefaultResponse, + TError = unknown, + TQueryKey extends Array = unknown[], +>( + { + dagId, + }: { + dagId: string; + }, + queryKey?: TQueryKey, + options?: Omit, "queryKey" | "queryFn">, +) => + useSuspenseQuery({ + queryKey: Common.UseAssetServiceNextRunAssets1KeyFn({ dagId }, queryKey), + queryFn: () => AssetService.nextRunAssets1({ dagId }) as TData, + ...options, + }); /** * Historical Metrics * Return cluster activity historical metrics. diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 4bfc986b56ba1..41676411a89c2 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -5,6 +5,8 @@ import { request as __request } from "./core/request"; import type { NextRunAssetsData, NextRunAssetsResponse, + NextRunAssets1Data, + NextRunAssets1Response, HistoricalMetricsData, HistoricalMetricsResponse, RecentDagRunsData, @@ -117,6 +119,28 @@ export class AssetService { }, }); } + + /** + * Next Run Assets + * @param data The data for the request. + * @param data.dagId + * @returns unknown Successful Response + * @throws ApiError + */ + public static nextRunAssets1( + data: NextRunAssets1Data, + ): CancelablePromise { + return __request(OpenAPI, { + method: "GET", + url: "/public/next_run_assets/{dag_id}", + path: { + dag_id: data.dagId, + }, + errors: { + 422: "Validation Error", + }, + }); + } } export class DashboardService { diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 06799783653bd..3cb914e09d381 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -735,6 +735,14 @@ export type NextRunAssetsResponse = { [key: string]: unknown; }; +export type NextRunAssets1Data = { + dagId: string; +}; + +export type NextRunAssets1Response = { + [key: string]: unknown; +}; + export type HistoricalMetricsData = { endDate: string; startDate: string; @@ -1093,6 +1101,23 @@ export type $OpenApiTs = { }; }; }; + "/public/next_run_assets/{dag_id}": { + get: { + req: NextRunAssets1Data; + res: { + /** + * Successful Response + */ + 200: { + [key: string]: unknown; + }; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; + }; "/ui/dashboard/historical_metrics_data": { get: { req: HistoricalMetricsData; From 5a49280ca9e5463219f5db2a3a29ef00162a9749 Mon Sep 17 00:00:00 2001 From: Amogh Date: Wed, 6 Nov 2024 19:00:47 +0530 Subject: [PATCH 02/19] matching response to legacy --- .../core_api/openapi/v1-generated.yaml | 188 +++++++++++++++++- .../core_api/routes/public/__init__.py | 2 +- .../core_api/routes/public/assets.py | 3 +- .../core_api/serializers/assets.py | 24 ++- airflow/ui/openapi-gen/queries/common.ts | 31 ++- airflow/ui/openapi-gen/queries/prefetch.ts | 38 +++- airflow/ui/openapi-gen/queries/queries.ts | 43 +++- airflow/ui/openapi-gen/queries/suspense.ts | 43 +++- .../ui/openapi-gen/requests/schemas.gen.ts | 157 +++++++++++++++ .../ui/openapi-gen/requests/services.gen.ts | 34 +++- airflow/ui/openapi-gen/requests/types.gen.ts | 83 +++++++- 11 files changed, 563 insertions(+), 83 deletions(-) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 4974d9540bc7d..473bc55b122db 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2515,27 +2515,76 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' - /public/next_run_assets/{dag_id}: + /public/assets/: get: tags: - Asset - summary: Next Run Assets - operationId: next_run_assets + summary: Get Assets + description: Get assets. + operationId: get_assets parameters: - - name: dag_id - in: path - required: true + - 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: uri_pattern + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + title: Uri Pattern + - name: dag_ids + in: query + required: false + schema: + type: array + items: + type: string + title: Dag Ids + - name: order_by + in: query + required: false schema: type: string - title: Dag Id + default: id + title: Order By responses: '200': description: Successful Response content: application/json: schema: - type: object - title: Response Next Run Assets + $ref: '#/components/schemas/AssetCollectionResponse' + '401': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Unauthorized + '403': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Forbidden + '404': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Not Found '422': description: Validation Error content: @@ -2591,6 +2640,83 @@ components: type: object title: AppBuilderViewResponse description: Serializer for AppBuilder View responses. + AssetAliasSchema: + properties: + id: + type: integer + title: Id + name: + type: string + title: Name + type: object + required: + - id + - name + title: AssetAliasSchema + description: Serializable version of the AssetAliasSchema ORM SqlAlchemyModel. + AssetCollectionResponse: + properties: + assets: + items: + $ref: '#/components/schemas/AssetResponse' + type: array + title: Assets + total_entries: + type: integer + title: Total Entries + type: object + required: + - assets + - total_entries + title: AssetCollectionResponse + description: Asset collection response. + AssetResponse: + properties: + id: + type: integer + title: Id + uri: + type: string + title: Uri + extra: + anyOf: + - type: object + - type: 'null' + title: Extra + created_at: + type: string + format: date-time + title: Created At + updated_at: + type: string + format: date-time + title: Updated At + consuming_dags: + items: + $ref: '#/components/schemas/DagScheduleAssetReference' + type: array + title: Consuming Dags + producing_tasks: + items: + $ref: '#/components/schemas/TaskOutletAssetReference' + type: array + title: Producing Tasks + aliases: + items: + $ref: '#/components/schemas/AssetAliasSchema' + type: array + title: Aliases + type: object + required: + - id + - uri + - created_at + - updated_at + - consuming_dags + - producing_tasks + - aliases + title: AssetResponse + description: Asset serializer for responses. BackfillPostBody: properties: dag_id: @@ -3547,6 +3673,26 @@ components: - asset_triggered title: DagRunType description: Class with DagRun types. + DagScheduleAssetReference: + properties: + dag_id: + type: string + title: Dag Id + created_at: + type: string + format: date-time + title: Created At + updated_at: + type: string + format: date-time + title: Updated At + type: object + required: + - dag_id + - created_at + - updated_at + title: DagScheduleAssetReference + description: Serializable version of the DagScheduleAssetReference ORM SqlAlchemyModel. DagStatsCollectionResponse: properties: dags: @@ -4303,6 +4449,30 @@ components: - triggerer_job title: TaskInstanceResponse description: TaskInstance serializer for responses. + TaskOutletAssetReference: + properties: + dag_id: + type: string + title: Dag Id + task_id: + type: string + title: Task Id + created_at: + type: string + format: date-time + title: Created At + updated_at: + type: string + format: date-time + title: Updated At + type: object + required: + - dag_id + - task_id + - created_at + - updated_at + title: TaskOutletAssetReference + description: Serializable version of the TaskOutletAssetReference ORM SqlAlchemyModel. TriggerResponse: properties: id: diff --git a/airflow/api_fastapi/core_api/routes/public/__init__.py b/airflow/api_fastapi/core_api/routes/public/__init__.py index 8359a8b679298..134bab4fc5190 100644 --- a/airflow/api_fastapi/core_api/routes/public/__init__.py +++ b/airflow/api_fastapi/core_api/routes/public/__init__.py @@ -18,6 +18,7 @@ from __future__ import annotations from airflow.api_fastapi.common.router import AirflowRouter +from airflow.api_fastapi.core_api.routes.public.assets import assets_router from airflow.api_fastapi.core_api.routes.public.backfills import backfills_router from airflow.api_fastapi.core_api.routes.public.connections import connections_router from airflow.api_fastapi.core_api.routes.public.dag_run import dag_run_router @@ -34,7 +35,6 @@ from airflow.api_fastapi.core_api.routes.public.task_instances import task_instances_router from airflow.api_fastapi.core_api.routes.public.variables import variables_router from airflow.api_fastapi.core_api.routes.public.version import version_router -from airflow.api_fastapi.core_api.routes.ui import assets_router public_router = AirflowRouter(prefix="/public") diff --git a/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow/api_fastapi/core_api/routes/public/assets.py index a90741dfd4479..584fb39b6b2bd 100644 --- a/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/airflow/api_fastapi/core_api/routes/public/assets.py @@ -34,7 +34,6 @@ from airflow.api_fastapi.common.router import AirflowRouter from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc from airflow.api_fastapi.core_api.serializers.assets import AssetCollectionResponse, AssetResponse -from airflow.assets import Asset from airflow.models.asset import AssetModel assets_router = AirflowRouter(tags=["Asset"], prefix="/assets") @@ -51,7 +50,7 @@ async def get_assets( dag_ids: QueryDagIdsFilter, order_by: Annotated[ SortParam, - Depends(SortParam(["id", "uri", "created_at", "updated_at"], Asset).dynamic_depends()), + Depends(SortParam(["id", "uri", "created_at", "updated_at"], AssetModel).dynamic_depends()), ], session: Annotated[Session, Depends(get_session)], ) -> AssetCollectionResponse: diff --git a/airflow/api_fastapi/core_api/serializers/assets.py b/airflow/api_fastapi/core_api/serializers/assets.py index 7acb8cc6357dc..498dc4edae6fd 100644 --- a/airflow/api_fastapi/core_api/serializers/assets.py +++ b/airflow/api_fastapi/core_api/serializers/assets.py @@ -19,30 +19,31 @@ from datetime import datetime -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel class DagScheduleAssetReference(BaseModel): """Serializable version of the DagScheduleAssetReference ORM SqlAlchemyModel.""" - asset_id: int dag_id: str created_at: datetime updated_at: datetime - model_config = ConfigDict(from_attributes=True) - class TaskOutletAssetReference(BaseModel): """Serializable version of the TaskOutletAssetReference ORM SqlAlchemyModel.""" - asset_id: int dag_id: str task_id: str created_at: datetime updated_at: datetime - model_config = ConfigDict(from_attributes=True) + +class AssetAliasSchema(BaseModel): + """Serializable version of the AssetAliasSchema ORM SqlAlchemyModel.""" + + id: int + name: str class AssetResponse(BaseModel): @@ -50,11 +51,12 @@ class AssetResponse(BaseModel): id: int uri: str - extra: str | None = Field(default=None) - created_at: str - updated_at: str - consuming_dags: DagScheduleAssetReference | None = None - producing_tasks: TaskOutletAssetReference | None = None + extra: dict | None = None + created_at: datetime + updated_at: datetime + consuming_dags: list[DagScheduleAssetReference] + producing_tasks: list[TaskOutletAssetReference] + aliases: list[AssetAliasSchema] class AssetCollectionResponse(BaseModel): diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 27b9d5c59a629..9100c69a29380 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -40,22 +40,33 @@ export const UseAssetServiceNextRunAssetsKeyFn = ( }, queryKey?: Array, ) => [useAssetServiceNextRunAssetsKey, ...(queryKey ?? [{ dagId }])]; -export type AssetServiceNextRunAssets1DefaultResponse = Awaited< - ReturnType +export type AssetServiceGetAssetsDefaultResponse = Awaited< + ReturnType >; -export type AssetServiceNextRunAssets1QueryResult< - TData = AssetServiceNextRunAssets1DefaultResponse, +export type AssetServiceGetAssetsQueryResult< + TData = AssetServiceGetAssetsDefaultResponse, TError = unknown, > = UseQueryResult; -export const useAssetServiceNextRunAssets1Key = "AssetServiceNextRunAssets1"; -export const UseAssetServiceNextRunAssets1KeyFn = ( +export const useAssetServiceGetAssetsKey = "AssetServiceGetAssets"; +export const UseAssetServiceGetAssetsKeyFn = ( { - dagId, + dagIds, + limit, + offset, + orderBy, + uriPattern, }: { - dagId: string; - }, + dagIds?: string[]; + limit?: number; + offset?: number; + orderBy?: string; + uriPattern?: string; + } = {}, queryKey?: Array, -) => [useAssetServiceNextRunAssets1Key, ...(queryKey ?? [{ dagId }])]; +) => [ + useAssetServiceGetAssetsKey, + ...(queryKey ?? [{ dagIds, limit, offset, orderBy, uriPattern }]), +]; export type DashboardServiceHistoricalMetricsDefaultResponse = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index c1b00fa352a20..7e5e9a12b77fc 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -45,23 +45,43 @@ export const prefetchUseAssetServiceNextRunAssets = ( queryFn: () => AssetService.nextRunAssets({ dagId }), }); /** - * Next Run Assets + * Get Assets + * Get assets. * @param data The data for the request. - * @param data.dagId - * @returns unknown Successful Response + * @param data.limit + * @param data.offset + * @param data.uriPattern + * @param data.dagIds + * @param data.orderBy + * @returns AssetCollectionResponse Successful Response * @throws ApiError */ -export const prefetchUseAssetServiceNextRunAssets1 = ( +export const prefetchUseAssetServiceGetAssets = ( queryClient: QueryClient, { - dagId, + dagIds, + limit, + offset, + orderBy, + uriPattern, }: { - dagId: string; - }, + dagIds?: string[]; + limit?: number; + offset?: number; + orderBy?: string; + uriPattern?: string; + } = {}, ) => queryClient.prefetchQuery({ - queryKey: Common.UseAssetServiceNextRunAssets1KeyFn({ dagId }), - queryFn: () => AssetService.nextRunAssets1({ dagId }), + queryKey: Common.UseAssetServiceGetAssetsKeyFn({ + dagIds, + limit, + offset, + orderBy, + uriPattern, + }), + queryFn: () => + AssetService.getAssets({ dagIds, limit, offset, orderBy, uriPattern }), }); /** * Historical Metrics diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index c3057b19e5997..12a61571dee48 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -66,28 +66,51 @@ export const useAssetServiceNextRunAssets = < ...options, }); /** - * Next Run Assets + * Get Assets + * Get assets. * @param data The data for the request. - * @param data.dagId - * @returns unknown Successful Response + * @param data.limit + * @param data.offset + * @param data.uriPattern + * @param data.dagIds + * @param data.orderBy + * @returns AssetCollectionResponse Successful Response * @throws ApiError */ -export const useAssetServiceNextRunAssets1 = < - TData = Common.AssetServiceNextRunAssets1DefaultResponse, +export const useAssetServiceGetAssets = < + TData = Common.AssetServiceGetAssetsDefaultResponse, TError = unknown, TQueryKey extends Array = unknown[], >( { - dagId, + dagIds, + limit, + offset, + orderBy, + uriPattern, }: { - dagId: string; - }, + dagIds?: string[]; + limit?: number; + offset?: number; + orderBy?: string; + uriPattern?: string; + } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => useQuery({ - queryKey: Common.UseAssetServiceNextRunAssets1KeyFn({ dagId }, queryKey), - queryFn: () => AssetService.nextRunAssets1({ dagId }) as TData, + queryKey: Common.UseAssetServiceGetAssetsKeyFn( + { dagIds, limit, offset, orderBy, uriPattern }, + queryKey, + ), + queryFn: () => + AssetService.getAssets({ + dagIds, + limit, + offset, + orderBy, + uriPattern, + }) as TData, ...options, }); /** diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index 033b4cdb513c8..6cce83916431b 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -51,28 +51,51 @@ export const useAssetServiceNextRunAssetsSuspense = < ...options, }); /** - * Next Run Assets + * Get Assets + * Get assets. * @param data The data for the request. - * @param data.dagId - * @returns unknown Successful Response + * @param data.limit + * @param data.offset + * @param data.uriPattern + * @param data.dagIds + * @param data.orderBy + * @returns AssetCollectionResponse Successful Response * @throws ApiError */ -export const useAssetServiceNextRunAssets1Suspense = < - TData = Common.AssetServiceNextRunAssets1DefaultResponse, +export const useAssetServiceGetAssetsSuspense = < + TData = Common.AssetServiceGetAssetsDefaultResponse, TError = unknown, TQueryKey extends Array = unknown[], >( { - dagId, + dagIds, + limit, + offset, + orderBy, + uriPattern, }: { - dagId: string; - }, + dagIds?: string[]; + limit?: number; + offset?: number; + orderBy?: string; + uriPattern?: string; + } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => useSuspenseQuery({ - queryKey: Common.UseAssetServiceNextRunAssets1KeyFn({ dagId }, queryKey), - queryFn: () => AssetService.nextRunAssets1({ dagId }) as TData, + queryKey: Common.UseAssetServiceGetAssetsKeyFn( + { dagIds, limit, offset, orderBy, uriPattern }, + queryKey, + ), + queryFn: () => + AssetService.getAssets({ + dagIds, + limit, + offset, + orderBy, + uriPattern, + }) as TData, ...options, }); /** diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index d64abb3853b49..f30ea53d66084 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -89,6 +89,111 @@ export const $AppBuilderViewResponse = { description: "Serializer for AppBuilder View responses.", } as const; +export const $AssetAliasSchema = { + properties: { + id: { + type: "integer", + title: "Id", + }, + name: { + type: "string", + title: "Name", + }, + }, + type: "object", + required: ["id", "name"], + title: "AssetAliasSchema", + description: + "Serializable version of the AssetAliasSchema ORM SqlAlchemyModel.", +} as const; + +export const $AssetCollectionResponse = { + properties: { + assets: { + items: { + $ref: "#/components/schemas/AssetResponse", + }, + type: "array", + title: "Assets", + }, + total_entries: { + type: "integer", + title: "Total Entries", + }, + }, + type: "object", + required: ["assets", "total_entries"], + title: "AssetCollectionResponse", + description: "Asset collection response.", +} as const; + +export const $AssetResponse = { + properties: { + id: { + type: "integer", + title: "Id", + }, + uri: { + type: "string", + title: "Uri", + }, + extra: { + anyOf: [ + { + type: "object", + }, + { + type: "null", + }, + ], + title: "Extra", + }, + created_at: { + type: "string", + format: "date-time", + title: "Created At", + }, + updated_at: { + type: "string", + format: "date-time", + title: "Updated At", + }, + consuming_dags: { + items: { + $ref: "#/components/schemas/DagScheduleAssetReference", + }, + type: "array", + title: "Consuming Dags", + }, + producing_tasks: { + items: { + $ref: "#/components/schemas/TaskOutletAssetReference", + }, + type: "array", + title: "Producing Tasks", + }, + aliases: { + items: { + $ref: "#/components/schemas/AssetAliasSchema", + }, + type: "array", + title: "Aliases", + }, + }, + type: "object", + required: [ + "id", + "uri", + "created_at", + "updated_at", + "consuming_dags", + "producing_tasks", + "aliases", + ], + title: "AssetResponse", + description: "Asset serializer for responses.", +} as const; + export const $BackfillPostBody = { properties: { dag_id: { @@ -1580,6 +1685,30 @@ export const $DagRunType = { description: "Class with DagRun types.", } as const; +export const $DagScheduleAssetReference = { + properties: { + dag_id: { + type: "string", + title: "Dag Id", + }, + created_at: { + type: "string", + format: "date-time", + title: "Created At", + }, + updated_at: { + type: "string", + format: "date-time", + title: "Updated At", + }, + }, + type: "object", + required: ["dag_id", "created_at", "updated_at"], + title: "DagScheduleAssetReference", + description: + "Serializable version of the DagScheduleAssetReference ORM SqlAlchemyModel.", +} as const; + export const $DagStatsCollectionResponse = { properties: { dags: { @@ -2716,6 +2845,34 @@ export const $TaskInstanceResponse = { description: "TaskInstance serializer for responses.", } as const; +export const $TaskOutletAssetReference = { + properties: { + dag_id: { + type: "string", + title: "Dag Id", + }, + task_id: { + type: "string", + title: "Task Id", + }, + created_at: { + type: "string", + format: "date-time", + title: "Created At", + }, + updated_at: { + type: "string", + format: "date-time", + title: "Updated At", + }, + }, + type: "object", + required: ["dag_id", "task_id", "created_at", "updated_at"], + title: "TaskOutletAssetReference", + description: + "Serializable version of the TaskOutletAssetReference ORM SqlAlchemyModel.", +} as const; + export const $TriggerResponse = { properties: { id: { diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 41676411a89c2..f5f43d82060a3 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -5,8 +5,8 @@ import { request as __request } from "./core/request"; import type { NextRunAssetsData, NextRunAssetsResponse, - NextRunAssets1Data, - NextRunAssets1Response, + GetAssetsData, + GetAssetsResponse, HistoricalMetricsData, HistoricalMetricsResponse, RecentDagRunsData, @@ -121,22 +121,34 @@ export class AssetService { } /** - * Next Run Assets + * Get Assets + * Get assets. * @param data The data for the request. - * @param data.dagId - * @returns unknown Successful Response + * @param data.limit + * @param data.offset + * @param data.uriPattern + * @param data.dagIds + * @param data.orderBy + * @returns AssetCollectionResponse Successful Response * @throws ApiError */ - public static nextRunAssets1( - data: NextRunAssets1Data, - ): CancelablePromise { + public static getAssets( + data: GetAssetsData = {}, + ): CancelablePromise { return __request(OpenAPI, { method: "GET", - url: "/public/next_run_assets/{dag_id}", - path: { - dag_id: data.dagId, + url: "/public/assets/", + query: { + limit: data.limit, + offset: data.offset, + uri_pattern: data.uriPattern, + dag_ids: data.dagIds, + order_by: data.orderBy, }, errors: { + 401: "Unauthorized", + 403: "Forbidden", + 404: "Not Found", 422: "Validation Error", }, }); diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 3cb914e09d381..71e313043f14a 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -21,6 +21,38 @@ export type AppBuilderViewResponse = { [key: string]: unknown; }; +/** + * Serializable version of the AssetAliasSchema ORM SqlAlchemyModel. + */ +export type AssetAliasSchema = { + id: number; + name: string; +}; + +/** + * Asset collection response. + */ +export type AssetCollectionResponse = { + assets: Array; + total_entries: number; +}; + +/** + * Asset serializer for responses. + */ +export type AssetResponse = { + id: number; + uri: string; + extra?: { + [key: string]: unknown; + } | null; + created_at: string; + updated_at: string; + consuming_dags: Array; + producing_tasks: Array; + aliases: Array; +}; + /** * Object used for create backfill request. */ @@ -347,6 +379,15 @@ export type DagRunType = | "manual" | "asset_triggered"; +/** + * Serializable version of the DagScheduleAssetReference ORM SqlAlchemyModel. + */ +export type DagScheduleAssetReference = { + dag_id: string; + created_at: string; + updated_at: string; +}; + /** * DAG Stats Collection serializer for responses. */ @@ -629,6 +670,16 @@ export type TaskInstanceResponse = { triggerer_job: JobResponse | null; }; +/** + * Serializable version of the TaskOutletAssetReference ORM SqlAlchemyModel. + */ +export type TaskOutletAssetReference = { + dag_id: string; + task_id: string; + created_at: string; + updated_at: string; +}; + /** * Trigger serializer for responses. */ @@ -735,13 +786,15 @@ export type NextRunAssetsResponse = { [key: string]: unknown; }; -export type NextRunAssets1Data = { - dagId: string; +export type GetAssetsData = { + dagIds?: Array; + limit?: number; + offset?: number; + orderBy?: string; + uriPattern?: string | null; }; -export type NextRunAssets1Response = { - [key: string]: unknown; -}; +export type GetAssetsResponse = AssetCollectionResponse; export type HistoricalMetricsData = { endDate: string; @@ -1101,16 +1154,26 @@ export type $OpenApiTs = { }; }; }; - "/public/next_run_assets/{dag_id}": { + "/public/assets/": { get: { - req: NextRunAssets1Data; + req: GetAssetsData; res: { /** * Successful Response */ - 200: { - [key: string]: unknown; - }; + 200: AssetCollectionResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Not Found + */ + 404: HTTPExceptionResponse; /** * Validation Error */ From 962572b5a08c17b8160f9a865895b4a2dbcd33c0 Mon Sep 17 00:00:00 2001 From: Amogh Date: Fri, 8 Nov 2024 10:59:04 +0530 Subject: [PATCH 03/19] Adding unit tests - part 1 --- airflow/api_fastapi/common/parameters.py | 3 +- .../core_api/routes/public/assets.py | 2 +- .../core_api/routes/public/test_assets.py | 190 ++++++++++++++++++ 3 files changed, 192 insertions(+), 3 deletions(-) create mode 100644 tests/api_fastapi/core_api/routes/public/test_assets.py diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index 224f6e0b7e775..c98b8540e1b7c 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -135,7 +135,7 @@ def __init__(self, attribute: ColumnElement, skip_none: bool = True) -> None: def to_orm(self, select: Select) -> Select: if self.value is None and self.skip_none: return select - return select.where(self.attribute.ilike(f"%{self.value}")) + return select.where(self.attribute.ilike(f"%{self.value}%")) def transform_aliases(self, value: str | None) -> str | None: if value == "~": @@ -343,7 +343,6 @@ def __init__(self, skip_none: bool = True) -> None: super().__init__(AssetModel.uri, skip_none) def depends(self, uri_pattern: str | None = None) -> _UriPatternSearch: - uri_pattern = super().transform_aliases(uri_pattern) return self.set_value(uri_pattern) diff --git a/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow/api_fastapi/core_api/routes/public/assets.py index 584fb39b6b2bd..961672a333844 100644 --- a/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/airflow/api_fastapi/core_api/routes/public/assets.py @@ -57,7 +57,7 @@ async def get_assets( """Get assets.""" assets_select, total_entries = paginated_select( select(AssetModel), - filters=[dag_ids, uri_pattern], + filters=[uri_pattern, dag_ids], order_by=order_by, offset=offset, limit=limit, diff --git a/tests/api_fastapi/core_api/routes/public/test_assets.py b/tests/api_fastapi/core_api/routes/public/test_assets.py new file mode 100644 index 0000000000000..04e2755412d33 --- /dev/null +++ b/tests/api_fastapi/core_api/routes/public/test_assets.py @@ -0,0 +1,190 @@ +# 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. +from __future__ import annotations + +import pytest + +from airflow.models.asset import AssetModel +from airflow.utils import timezone +from airflow.utils.session import provide_session + +from tests_common.test_utils.config import conf_vars +from tests_common.test_utils.db import clear_db_assets + +pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode] + + +def _create_assets(session, num: int = 2) -> None: + default_time = "2020-06-11T18:00:00+00:00" + assets = [ + AssetModel( + id=i, + uri=f"s3://bucket/key/{i}", + extra={"foo": "bar"}, + created_at=timezone.parse(default_time), + updated_at=timezone.parse(default_time), + ) + for i in range(1, 1 + num) + ] + session.add_all(assets) + session.commit() + + +def _create_provided_asset(asset: AssetModel, session) -> None: + session.add(asset) + session.commit() + + +class TestAssets: + default_time = "2020-06-11T18:00:00+00:00" + + @pytest.fixture(autouse=True) + def setup(self) -> None: + clear_db_assets() + + def teardown_method(self) -> None: + clear_db_assets() + + @provide_session + def create_assets(self, session, num: int = 2): + _create_assets(session, num=num) + + @provide_session + def create_provided_asset(self, session, asset: AssetModel): + _create_provided_asset(session, asset) + + +class TestGetAssets(TestAssets): + def test_should_respond_200(self, test_client, session): + self.create_assets() + assets = session.query(AssetModel).all() + assert len(assets) == 2 + + response = test_client.get("/public/assets") + assert response.status_code == 200 + response_data = response.json() + tz_datetime_format = self.default_time.replace("+00:00", "Z") + assert response_data == { + "assets": [ + { + "id": 1, + "uri": "s3://bucket/key/1", + "extra": {"foo": "bar"}, + "created_at": tz_datetime_format, + "updated_at": tz_datetime_format, + "consuming_dags": [], + "producing_tasks": [], + "aliases": [], + }, + { + "id": 2, + "uri": "s3://bucket/key/2", + "extra": {"foo": "bar"}, + "created_at": tz_datetime_format, + "updated_at": tz_datetime_format, + "consuming_dags": [], + "producing_tasks": [], + "aliases": [], + }, + ], + "total_entries": 2, + } + + def test_order_by_raises_400_for_invalid_attr(self, test_client, session): + response = test_client.get("/public/assets?order_by=fake") + + assert response.status_code == 400 + msg = "Ordering with 'fake' is disallowed or the attribute does not exist on the model" + assert response.json()["detail"] == msg + + @pytest.mark.parametrize( + "url, expected_assets", + [ + ("/public/assets?uri_pattern=s3", {"s3://folder/key"}), + ("/public/assets?uri_pattern=bucket", {"gcp://bucket/key", "wasb://some_asset_bucket_/key"}), + ( + "/public/assets?uri_pattern=asset", + {"somescheme://asset/key", "wasb://some_asset_bucket_/key"}, + ), + ( + "/public/assets?uri_pattern=", + { + "gcp://bucket/key", + "s3://folder/key", + "somescheme://asset/key", + "wasb://some_asset_bucket_/key", + }, + ), + ], + ) + @provide_session + def test_filter_assets_by_uri_pattern_works(self, test_client, url, expected_assets, session): + asset1 = AssetModel("s3://folder/key") + asset2 = AssetModel("gcp://bucket/key") + asset3 = AssetModel("somescheme://asset/key") + asset4 = AssetModel("wasb://some_asset_bucket_/key") + + assets = [asset1, asset2, asset3, asset4] + for a in assets: + self.create_provided_asset(a) + + response = test_client.get(url) + assert response.status_code == 200 + asset_urls = {asset["uri"] for asset in response.json()["assets"]} + assert expected_assets == asset_urls + + +class TestGetAssetsEndpointPagination(TestAssets): + @pytest.mark.parametrize( + "url, expected_asset_uris", + [ + # Limit test data + ("/public/assets?limit=1", ["s3://bucket/key/1"]), + ("/public/assets?limit=100", [f"s3://bucket/key/{i}" for i in range(1, 101)]), + # Offset test data + ("/public/assets?offset=1", [f"s3://bucket/key/{i}" for i in range(2, 102)]), + ("/public/assets?offset=3", [f"s3://bucket/key/{i}" for i in range(4, 104)]), + # Limit and offset test data + ("/public/assets?offset=3&limit=3", [f"s3://bucket/key/{i}" for i in [4, 5, 6]]), + ], + ) + def test_limit_and_offset(self, test_client, url, expected_asset_uris): + self.create_assets(110) + + response = test_client.get(url) + + assert response.status_code == 200 + asset_uris = [asset["uri"] for asset in response.json()["assets"]] + assert asset_uris == expected_asset_uris + + def test_should_respect_page_size_limit_default(self, test_client): + self.create_assets(110) + + response = test_client.get("/public/assets") + + assert response.status_code == 200 + assert len(response.json()["assets"]) == 100 + + @conf_vars({("api", "maximum_page_limit"): "150"}) + def test_should_return_conf_max_if_req_max_above_conf(self, test_client): + self.create_assets(200) + + # change to 180 once format_parameters is integrated + response = test_client.get("/public/assets?limit=150") + + assert response.status_code == 200 + assert len(response.json()["assets"]) == 150 From 428cb6c6bd9084370e1e244092c1de326d291d8c Mon Sep 17 00:00:00 2001 From: Amogh Desai Date: Fri, 8 Nov 2024 11:03:12 +0530 Subject: [PATCH 04/19] Update airflow/api_fastapi/common/parameters.py Co-authored-by: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com> --- airflow/api_fastapi/common/parameters.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index c98b8540e1b7c..8c32f114172fe 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -337,7 +337,7 @@ def depends(self, dag_id: str | None = None) -> _DagIdFilter: class _UriPatternSearch(_SearchParam): - """Search on dag_id.""" + """Search on uri.""" def __init__(self, skip_none: bool = True) -> None: super().__init__(AssetModel.uri, skip_none) From a78d3cb9209d160f53b304e989c412b0a324d28a Mon Sep 17 00:00:00 2001 From: Amogh Date: Fri, 8 Nov 2024 11:55:18 +0530 Subject: [PATCH 05/19] fixing the dag_ids filter --- airflow/api_fastapi/common/parameters.py | 27 ++++++++++++++++++- .../core_api/openapi/v1-generated.yaml | 6 ++--- .../core_api/routes/public/assets.py | 4 +-- airflow/ui/openapi-gen/queries/common.ts | 4 +-- airflow/ui/openapi-gen/queries/prefetch.ts | 6 ++--- airflow/ui/openapi-gen/queries/queries.ts | 6 ++--- airflow/ui/openapi-gen/queries/suspense.ts | 6 ++--- .../ui/openapi-gen/requests/services.gen.ts | 4 +-- airflow/ui/openapi-gen/requests/types.gen.ts | 2 +- 9 files changed, 44 insertions(+), 21 deletions(-) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index 8c32f114172fe..8ac1803f91737 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -29,7 +29,7 @@ from typing_extensions import Annotated, Self from airflow.models import Base, Connection -from airflow.models.asset import AssetModel +from airflow.models.asset import AssetModel, DagScheduleAssetReference, TaskOutletAssetReference from airflow.models.dag import DagModel, DagTag from airflow.models.dagrun import DagRun from airflow.models.dagwarning import DagWarning, DagWarningType @@ -346,6 +346,28 @@ def depends(self, uri_pattern: str | None = None) -> _UriPatternSearch: return self.set_value(uri_pattern) +class _DagIdAssetReferencePatternSearch(_SearchParam): + """Search on dag_id.""" + + def __init__(self, skip_none: bool = True) -> None: + super().__init__(AssetModel.consuming_dags, skip_none) + self.task_attribute = AssetModel.consuming_dags + + def depends(self, dag_ids: str) -> _DagIdAssetReferencePatternSearch: + return self.set_value(dag_ids) + + def to_orm(self, select: Select) -> Select: + if self.value is None and self.skip_none: + return select + if self.value is not None: + dags_list = self.value.split(",") + print("dags list", dags_list) + return select.filter( + (self.attribute.any(DagScheduleAssetReference.dag_id.in_(dags_list))) + | (self.task_attribute.any(TaskOutletAssetReference.dag_id.in_(dags_list))) + ) + + # Common Safe DateTime DateTimeQuery = Annotated[str, AfterValidator(_safe_parse_datetime)] @@ -377,3 +399,6 @@ def depends(self, uri_pattern: str | None = None) -> _UriPatternSearch: # Assets QueryUriPatternSearch = Annotated[_UriPatternSearch, Depends(_UriPatternSearch().depends)] +QueryAssetDagIdPatternSearch = Annotated[ + _DagIdAssetReferencePatternSearch, Depends(_DagIdAssetReferencePatternSearch().depends) +] diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 473bc55b122db..ab1369a4373cf 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -2547,11 +2547,9 @@ paths: title: Uri Pattern - name: dag_ids in: query - required: false + required: true schema: - type: array - items: - type: string + type: string title: Dag Ids - name: order_by in: query diff --git a/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow/api_fastapi/core_api/routes/public/assets.py index 961672a333844..8597dc40d66a3 100644 --- a/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/airflow/api_fastapi/core_api/routes/public/assets.py @@ -25,7 +25,7 @@ from airflow.api_fastapi.common.db.common import get_session, paginated_select from airflow.api_fastapi.common.parameters import ( - QueryDagIdsFilter, + QueryAssetDagIdPatternSearch, QueryLimit, QueryOffset, QueryUriPatternSearch, @@ -47,7 +47,7 @@ async def get_assets( limit: QueryLimit, offset: QueryOffset, uri_pattern: QueryUriPatternSearch, - dag_ids: QueryDagIdsFilter, + dag_ids: QueryAssetDagIdPatternSearch, order_by: Annotated[ SortParam, Depends(SortParam(["id", "uri", "created_at", "updated_at"], AssetModel).dynamic_depends()), diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 9100c69a29380..3cf15c9aa42c5 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -56,12 +56,12 @@ export const UseAssetServiceGetAssetsKeyFn = ( orderBy, uriPattern, }: { - dagIds?: string[]; + dagIds: string; limit?: number; offset?: number; orderBy?: string; uriPattern?: string; - } = {}, + }, queryKey?: Array, ) => [ useAssetServiceGetAssetsKey, diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index 7e5e9a12b77fc..acef4e764b06f 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -48,10 +48,10 @@ export const prefetchUseAssetServiceNextRunAssets = ( * Get Assets * Get assets. * @param data The data for the request. + * @param data.dagIds * @param data.limit * @param data.offset * @param data.uriPattern - * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response * @throws ApiError @@ -65,12 +65,12 @@ export const prefetchUseAssetServiceGetAssets = ( orderBy, uriPattern, }: { - dagIds?: string[]; + dagIds: string; limit?: number; offset?: number; orderBy?: string; uriPattern?: string; - } = {}, + }, ) => queryClient.prefetchQuery({ queryKey: Common.UseAssetServiceGetAssetsKeyFn({ diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 12a61571dee48..6924213faab4c 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -69,10 +69,10 @@ export const useAssetServiceNextRunAssets = < * Get Assets * Get assets. * @param data The data for the request. + * @param data.dagIds * @param data.limit * @param data.offset * @param data.uriPattern - * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response * @throws ApiError @@ -89,12 +89,12 @@ export const useAssetServiceGetAssets = < orderBy, uriPattern, }: { - dagIds?: string[]; + dagIds: string; limit?: number; offset?: number; orderBy?: string; uriPattern?: string; - } = {}, + }, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index 6cce83916431b..5672b2dcb2947 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -54,10 +54,10 @@ export const useAssetServiceNextRunAssetsSuspense = < * Get Assets * Get assets. * @param data The data for the request. + * @param data.dagIds * @param data.limit * @param data.offset * @param data.uriPattern - * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response * @throws ApiError @@ -74,12 +74,12 @@ export const useAssetServiceGetAssetsSuspense = < orderBy, uriPattern, }: { - dagIds?: string[]; + dagIds: string; limit?: number; offset?: number; orderBy?: string; uriPattern?: string; - } = {}, + }, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index f5f43d82060a3..bd0c3466fc78f 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -124,16 +124,16 @@ export class AssetService { * Get Assets * Get assets. * @param data The data for the request. + * @param data.dagIds * @param data.limit * @param data.offset * @param data.uriPattern - * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response * @throws ApiError */ public static getAssets( - data: GetAssetsData = {}, + data: GetAssetsData, ): CancelablePromise { return __request(OpenAPI, { method: "GET", diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 71e313043f14a..a3e35c8082889 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -787,7 +787,7 @@ export type NextRunAssetsResponse = { }; export type GetAssetsData = { - dagIds?: Array; + dagIds: string; limit?: number; offset?: number; orderBy?: string; From 882d20cc7e7d488e736ca5159026cbd3e5860d73 Mon Sep 17 00:00:00 2001 From: Amogh Date: Fri, 8 Nov 2024 12:16:21 +0530 Subject: [PATCH 06/19] fixing the dag_ids filter --- airflow/api_fastapi/common/parameters.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index 8ac1803f91737..03581f689f889 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -351,7 +351,7 @@ class _DagIdAssetReferencePatternSearch(_SearchParam): def __init__(self, skip_none: bool = True) -> None: super().__init__(AssetModel.consuming_dags, skip_none) - self.task_attribute = AssetModel.consuming_dags + self.task_attribute = AssetModel.producing_tasks def depends(self, dag_ids: str) -> _DagIdAssetReferencePatternSearch: return self.set_value(dag_ids) From 25bb08e637aa4fe0f76912a11abe07709addd36c Mon Sep 17 00:00:00 2001 From: Amogh Date: Fri, 8 Nov 2024 12:19:40 +0530 Subject: [PATCH 07/19] Adding unit tests - part 2 --- .../core_api/routes/public/test_assets.py | 53 ++++++++++++++++++- 1 file changed, 52 insertions(+), 1 deletion(-) diff --git a/tests/api_fastapi/core_api/routes/public/test_assets.py b/tests/api_fastapi/core_api/routes/public/test_assets.py index 04e2755412d33..d3eee703d01d6 100644 --- a/tests/api_fastapi/core_api/routes/public/test_assets.py +++ b/tests/api_fastapi/core_api/routes/public/test_assets.py @@ -18,7 +18,8 @@ import pytest -from airflow.models.asset import AssetModel +from airflow.models import DagModel +from airflow.models.asset import AssetModel, DagScheduleAssetReference, TaskOutletAssetReference from airflow.utils import timezone from airflow.utils.session import provide_session @@ -147,6 +148,56 @@ def test_filter_assets_by_uri_pattern_works(self, test_client, url, expected_ass asset_urls = {asset["uri"] for asset in response.json()["assets"]} assert expected_assets == asset_urls + @pytest.mark.parametrize("dag_ids, expected_num", [("dag1,dag2", 2), ("dag3", 1), ("dag2,dag3", 2)]) + @provide_session + def test_filter_assets_by_dag_ids_works(self, test_client, dag_ids, expected_num, session): + session.query(DagModel).delete() + session.commit() + dag1 = DagModel(dag_id="dag1") + dag2 = DagModel(dag_id="dag2") + dag3 = DagModel(dag_id="dag3") + asset1 = AssetModel("s3://folder/key") + asset2 = AssetModel("gcp://bucket/key") + asset3 = AssetModel("somescheme://asset/key") + dag_ref1 = DagScheduleAssetReference(dag_id="dag1", asset=asset1) + dag_ref2 = DagScheduleAssetReference(dag_id="dag2", asset=asset2) + task_ref1 = TaskOutletAssetReference(dag_id="dag3", task_id="task1", asset=asset3) + session.add_all([asset1, asset2, asset3, dag1, dag2, dag3, dag_ref1, dag_ref2, task_ref1]) + session.commit() + response = test_client.get( + f"/public/assets?dag_ids={dag_ids}", + ) + assert response.status_code == 200 + response_data = response.json() + assert len(response_data["assets"]) == expected_num + + @pytest.mark.parametrize( + "dag_ids, uri_pattern,expected_num", + [("dag1,dag2", "folder", 1), ("dag3", "nothing", 0), ("dag2,dag3", "key", 2)], + ) + def test_filter_assets_by_dag_ids_and_uri_pattern_works( + self, test_client, dag_ids, uri_pattern, expected_num, session + ): + session.query(DagModel).delete() + session.commit() + dag1 = DagModel(dag_id="dag1") + dag2 = DagModel(dag_id="dag2") + dag3 = DagModel(dag_id="dag3") + asset1 = AssetModel("s3://folder/key") + asset2 = AssetModel("gcp://bucket/key") + asset3 = AssetModel("somescheme://asset/key") + dag_ref1 = DagScheduleAssetReference(dag_id="dag1", asset=asset1) + dag_ref2 = DagScheduleAssetReference(dag_id="dag2", asset=asset2) + task_ref1 = TaskOutletAssetReference(dag_id="dag3", task_id="task1", asset=asset3) + session.add_all([asset1, asset2, asset3, dag1, dag2, dag3, dag_ref1, dag_ref2, task_ref1]) + session.commit() + response = test_client.get( + f"/public/assets?dag_ids={dag_ids}&uri_pattern={uri_pattern}", + ) + assert response.status_code == 200 + response_data = response.json() + assert len(response_data["assets"]) == expected_num + class TestGetAssetsEndpointPagination(TestAssets): @pytest.mark.parametrize( From fa0cd23cc9cadee09f1a53abdcd6a0a7164a316b Mon Sep 17 00:00:00 2001 From: Amogh Date: Fri, 8 Nov 2024 13:18:20 +0530 Subject: [PATCH 08/19] fixing unit tests & updating parameter type --- airflow/api_fastapi/common/parameters.py | 2 +- .../api_fastapi/core_api/openapi/v1-generated.yaml | 6 ++++-- airflow/ui/openapi-gen/queries/common.ts | 4 ++-- airflow/ui/openapi-gen/queries/prefetch.ts | 6 +++--- airflow/ui/openapi-gen/queries/queries.ts | 6 +++--- airflow/ui/openapi-gen/queries/suspense.ts | 6 +++--- airflow/ui/openapi-gen/requests/services.gen.ts | 4 ++-- airflow/ui/openapi-gen/requests/types.gen.ts | 2 +- .../core_api/routes/public/test_assets.py | 14 +++++++------- 9 files changed, 26 insertions(+), 24 deletions(-) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index 74f0b7da0c35d..b334fe7a4a987 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -424,7 +424,7 @@ def __init__(self, skip_none: bool = True) -> None: super().__init__(AssetModel.consuming_dags, skip_none) self.task_attribute = AssetModel.producing_tasks - def depends(self, dag_ids: str) -> _DagIdAssetReferencePatternSearch: + def depends(self, dag_ids: str | None = None) -> _DagIdAssetReferencePatternSearch: return self.set_value(dag_ids) def to_orm(self, select: Select) -> Select: diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 307525a835c2a..d12d7f6deb8e8 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -3136,9 +3136,11 @@ paths: title: Uri Pattern - name: dag_ids in: query - required: true + required: false schema: - type: string + anyOf: + - type: string + - type: 'null' title: Dag Ids - name: order_by in: query diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 048fe3e156e1a..161e7627d41aa 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -56,12 +56,12 @@ export const UseAssetServiceGetAssetsKeyFn = ( orderBy, uriPattern, }: { - dagIds: string; + dagIds?: string; limit?: number; offset?: number; orderBy?: string; uriPattern?: string; - }, + } = {}, queryKey?: Array, ) => [ useAssetServiceGetAssetsKey, diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index 2f6250ff20ad5..9052d8b71989d 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -48,10 +48,10 @@ export const prefetchUseAssetServiceNextRunAssets = ( * Get Assets * Get assets. * @param data The data for the request. - * @param data.dagIds * @param data.limit * @param data.offset * @param data.uriPattern + * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response * @throws ApiError @@ -65,12 +65,12 @@ export const prefetchUseAssetServiceGetAssets = ( orderBy, uriPattern, }: { - dagIds: string; + dagIds?: string; limit?: number; offset?: number; orderBy?: string; uriPattern?: string; - }, + } = {}, ) => queryClient.prefetchQuery({ queryKey: Common.UseAssetServiceGetAssetsKeyFn({ diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index cf706ddf17e57..73df2f3f5fc23 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -69,10 +69,10 @@ export const useAssetServiceNextRunAssets = < * Get Assets * Get assets. * @param data The data for the request. - * @param data.dagIds * @param data.limit * @param data.offset * @param data.uriPattern + * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response * @throws ApiError @@ -89,12 +89,12 @@ export const useAssetServiceGetAssets = < orderBy, uriPattern, }: { - dagIds: string; + dagIds?: string; limit?: number; offset?: number; orderBy?: string; uriPattern?: string; - }, + } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index e71412a1e5df9..46545c6138e28 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -54,10 +54,10 @@ export const useAssetServiceNextRunAssetsSuspense = < * Get Assets * Get assets. * @param data The data for the request. - * @param data.dagIds * @param data.limit * @param data.offset * @param data.uriPattern + * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response * @throws ApiError @@ -74,12 +74,12 @@ export const useAssetServiceGetAssetsSuspense = < orderBy, uriPattern, }: { - dagIds: string; + dagIds?: string; limit?: number; offset?: number; orderBy?: string; uriPattern?: string; - }, + } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 34da5fbff304c..921af0e471286 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -134,16 +134,16 @@ export class AssetService { * Get Assets * Get assets. * @param data The data for the request. - * @param data.dagIds * @param data.limit * @param data.offset * @param data.uriPattern + * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response * @throws ApiError */ public static getAssets( - data: GetAssetsData, + data: GetAssetsData = {}, ): CancelablePromise { return __request(OpenAPI, { method: "GET", diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 809edebd13de5..730a0a53ea23b 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -810,7 +810,7 @@ export type NextRunAssetsResponse = { }; export type GetAssetsData = { - dagIds: string; + dagIds?: string | null; limit?: number; offset?: number; orderBy?: string; diff --git a/tests/api_fastapi/core_api/routes/public/test_assets.py b/tests/api_fastapi/core_api/routes/public/test_assets.py index d3eee703d01d6..722573339c710 100644 --- a/tests/api_fastapi/core_api/routes/public/test_assets.py +++ b/tests/api_fastapi/core_api/routes/public/test_assets.py @@ -45,7 +45,7 @@ def _create_assets(session, num: int = 2) -> None: session.commit() -def _create_provided_asset(asset: AssetModel, session) -> None: +def _create_provided_asset(session, asset: AssetModel) -> None: session.add(asset) session.commit() @@ -62,11 +62,11 @@ def teardown_method(self) -> None: @provide_session def create_assets(self, session, num: int = 2): - _create_assets(session, num=num) + _create_assets(session=session, num=num) @provide_session def create_provided_asset(self, session, asset: AssetModel): - _create_provided_asset(session, asset) + _create_provided_asset(session=session, asset=asset) class TestGetAssets(TestAssets): @@ -141,7 +141,7 @@ def test_filter_assets_by_uri_pattern_works(self, test_client, url, expected_ass assets = [asset1, asset2, asset3, asset4] for a in assets: - self.create_provided_asset(a) + self.create_provided_asset(asset=a) response = test_client.get(url) assert response.status_code == 200 @@ -214,7 +214,7 @@ class TestGetAssetsEndpointPagination(TestAssets): ], ) def test_limit_and_offset(self, test_client, url, expected_asset_uris): - self.create_assets(110) + self.create_assets(num=110) response = test_client.get(url) @@ -223,7 +223,7 @@ def test_limit_and_offset(self, test_client, url, expected_asset_uris): assert asset_uris == expected_asset_uris def test_should_respect_page_size_limit_default(self, test_client): - self.create_assets(110) + self.create_assets(num=110) response = test_client.get("/public/assets") @@ -232,7 +232,7 @@ def test_should_respect_page_size_limit_default(self, test_client): @conf_vars({("api", "maximum_page_limit"): "150"}) def test_should_return_conf_max_if_req_max_above_conf(self, test_client): - self.create_assets(200) + self.create_assets(num=200) # change to 180 once format_parameters is integrated response = test_client.get("/public/assets?limit=150") From dd791c214af58c452ead6e8836c255ada3a5d175 Mon Sep 17 00:00:00 2001 From: Amogh Date: Fri, 8 Nov 2024 20:42:12 +0530 Subject: [PATCH 09/19] review comments pierre --- airflow/api_fastapi/common/parameters.py | 12 ++++++------ .../api_fastapi/core_api/openapi/v1-generated.yaml | 6 ++---- .../api_fastapi/core_api/routes/public/assets.py | 4 ++-- airflow/ui/openapi-gen/queries/common.ts | 4 ++-- airflow/ui/openapi-gen/queries/prefetch.ts | 6 +++--- airflow/ui/openapi-gen/queries/queries.ts | 6 +++--- airflow/ui/openapi-gen/queries/suspense.ts | 6 +++--- airflow/ui/openapi-gen/requests/services.gen.ts | 4 ++-- airflow/ui/openapi-gen/requests/types.gen.ts | 2 +- .../core_api/routes/public/test_assets.py | 14 +++++++------- 10 files changed, 31 insertions(+), 33 deletions(-) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index b334fe7a4a987..7cc7350e91902 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -413,18 +413,18 @@ class _UriPatternSearch(_SearchParam): def __init__(self, skip_none: bool = True) -> None: super().__init__(AssetModel.uri, skip_none) - def depends(self, uri_pattern: str | None = None) -> _UriPatternSearch: + def depends(self, uri_pattern: str) -> _UriPatternSearch: return self.set_value(uri_pattern) -class _DagIdAssetReferencePatternSearch(_SearchParam): +class _DagIdAssetReferenceFilter(BaseParam[str]): """Search on dag_id.""" def __init__(self, skip_none: bool = True) -> None: super().__init__(AssetModel.consuming_dags, skip_none) self.task_attribute = AssetModel.producing_tasks - def depends(self, dag_ids: str | None = None) -> _DagIdAssetReferencePatternSearch: + def depends(self, dag_ids: str | None = None) -> _DagIdAssetReferenceFilter: return self.set_value(dag_ids) def to_orm(self, select: Select) -> Select: @@ -433,8 +433,8 @@ def to_orm(self, select: Select) -> Select: if self.value is not None: dags_list = self.value.split(",") return select.filter( - (self.attribute.any(DagScheduleAssetReference.dag_id.in_(dags_list))) - | (self.task_attribute.any(TaskOutletAssetReference.dag_id.in_(dags_list))) + (AssetModel.consuming_dags.any(DagScheduleAssetReference.dag_id.in_(dags_list))) + | (AssetModel.producing_tasks.any(TaskOutletAssetReference.dag_id.in_(dags_list))) ) @@ -533,5 +533,5 @@ def depends_float( # Assets QueryUriPatternSearch = Annotated[_UriPatternSearch, Depends(_UriPatternSearch().depends)] QueryAssetDagIdPatternSearch = Annotated[ - _DagIdAssetReferencePatternSearch, Depends(_DagIdAssetReferencePatternSearch().depends) + _DagIdAssetReferenceFilter, Depends(_DagIdAssetReferenceFilter().depends) ] diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index d12d7f6deb8e8..dd3bedabf0a04 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -3128,11 +3128,9 @@ paths: title: Offset - name: uri_pattern in: query - required: false + required: true schema: - anyOf: - - type: string - - type: 'null' + type: string title: Uri Pattern - name: dag_ids in: query diff --git a/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow/api_fastapi/core_api/routes/public/assets.py index 8597dc40d66a3..26d09e7ff38ab 100644 --- a/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/airflow/api_fastapi/core_api/routes/public/assets.py @@ -43,7 +43,7 @@ "/", responses=create_openapi_http_exception_doc([401, 403, 404]), ) -async def get_assets( +def get_assets( limit: QueryLimit, offset: QueryOffset, uri_pattern: QueryUriPatternSearch, @@ -66,6 +66,6 @@ async def get_assets( assets = session.scalars(assets_select).all() return AssetCollectionResponse( - assets=[AssetResponse.model_validate(x, from_attributes=True) for x in assets], + assets=[AssetResponse.model_validate(asset, from_attributes=True) for asset in assets], total_entries=total_entries, ) diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 161e7627d41aa..01d92f7c0576d 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -60,8 +60,8 @@ export const UseAssetServiceGetAssetsKeyFn = ( limit?: number; offset?: number; orderBy?: string; - uriPattern?: string; - } = {}, + uriPattern: string; + }, queryKey?: Array, ) => [ useAssetServiceGetAssetsKey, diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index 9052d8b71989d..274e6062120d1 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -48,9 +48,9 @@ export const prefetchUseAssetServiceNextRunAssets = ( * Get Assets * Get assets. * @param data The data for the request. + * @param data.uriPattern * @param data.limit * @param data.offset - * @param data.uriPattern * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response @@ -69,8 +69,8 @@ export const prefetchUseAssetServiceGetAssets = ( limit?: number; offset?: number; orderBy?: string; - uriPattern?: string; - } = {}, + uriPattern: string; + }, ) => queryClient.prefetchQuery({ queryKey: Common.UseAssetServiceGetAssetsKeyFn({ diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 73df2f3f5fc23..e472edba8f4f1 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -69,9 +69,9 @@ export const useAssetServiceNextRunAssets = < * Get Assets * Get assets. * @param data The data for the request. + * @param data.uriPattern * @param data.limit * @param data.offset - * @param data.uriPattern * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response @@ -93,8 +93,8 @@ export const useAssetServiceGetAssets = < limit?: number; offset?: number; orderBy?: string; - uriPattern?: string; - } = {}, + uriPattern: string; + }, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index 46545c6138e28..54a3632e95b1b 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -54,9 +54,9 @@ export const useAssetServiceNextRunAssetsSuspense = < * Get Assets * Get assets. * @param data The data for the request. + * @param data.uriPattern * @param data.limit * @param data.offset - * @param data.uriPattern * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response @@ -78,8 +78,8 @@ export const useAssetServiceGetAssetsSuspense = < limit?: number; offset?: number; orderBy?: string; - uriPattern?: string; - } = {}, + uriPattern: string; + }, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 921af0e471286..f06a4417f0491 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -134,16 +134,16 @@ export class AssetService { * Get Assets * Get assets. * @param data The data for the request. + * @param data.uriPattern * @param data.limit * @param data.offset - * @param data.uriPattern * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response * @throws ApiError */ public static getAssets( - data: GetAssetsData = {}, + data: GetAssetsData, ): CancelablePromise { return __request(OpenAPI, { method: "GET", diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 730a0a53ea23b..c84223f95e35a 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -814,7 +814,7 @@ export type GetAssetsData = { limit?: number; offset?: number; orderBy?: string; - uriPattern?: string | null; + uriPattern: string; }; export type GetAssetsResponse = AssetCollectionResponse; diff --git a/tests/api_fastapi/core_api/routes/public/test_assets.py b/tests/api_fastapi/core_api/routes/public/test_assets.py index 722573339c710..c6d085d060399 100644 --- a/tests/api_fastapi/core_api/routes/public/test_assets.py +++ b/tests/api_fastapi/core_api/routes/public/test_assets.py @@ -113,16 +113,16 @@ def test_order_by_raises_400_for_invalid_attr(self, test_client, session): assert response.json()["detail"] == msg @pytest.mark.parametrize( - "url, expected_assets", + "params, expected_assets", [ - ("/public/assets?uri_pattern=s3", {"s3://folder/key"}), - ("/public/assets?uri_pattern=bucket", {"gcp://bucket/key", "wasb://some_asset_bucket_/key"}), + ({"uri_pattern": "s3"}, {"s3://folder/key"}), + ({"uri_pattern": "bucket"}, {"gcp://bucket/key", "wasb://some_asset_bucket_/key"}), ( - "/public/assets?uri_pattern=asset", + {"uri_pattern": "asset"}, {"somescheme://asset/key", "wasb://some_asset_bucket_/key"}, ), ( - "/public/assets?uri_pattern=", + {"uri_pattern": ""}, { "gcp://bucket/key", "s3://folder/key", @@ -133,7 +133,7 @@ def test_order_by_raises_400_for_invalid_attr(self, test_client, session): ], ) @provide_session - def test_filter_assets_by_uri_pattern_works(self, test_client, url, expected_assets, session): + def test_filter_assets_by_uri_pattern_works(self, test_client, params, expected_assets, session): asset1 = AssetModel("s3://folder/key") asset2 = AssetModel("gcp://bucket/key") asset3 = AssetModel("somescheme://asset/key") @@ -143,7 +143,7 @@ def test_filter_assets_by_uri_pattern_works(self, test_client, url, expected_ass for a in assets: self.create_provided_asset(asset=a) - response = test_client.get(url) + response = test_client.get("/public/assets", params=params) assert response.status_code == 200 asset_urls = {asset["uri"] for asset in response.json()["assets"]} assert expected_assets == asset_urls From 06fa0a701a9ac7ceae045b1b9d99984febc9fd17 Mon Sep 17 00:00:00 2001 From: Amogh Date: Fri, 8 Nov 2024 21:01:34 +0530 Subject: [PATCH 10/19] fixing last commit --- airflow/api_fastapi/common/parameters.py | 15 ++++++--------- .../core_api/openapi/v1-generated.yaml | 12 +++++++----- airflow/ui/openapi-gen/queries/common.ts | 6 +++--- airflow/ui/openapi-gen/queries/prefetch.ts | 8 ++++---- airflow/ui/openapi-gen/queries/queries.ts | 8 ++++---- airflow/ui/openapi-gen/queries/suspense.ts | 8 ++++---- airflow/ui/openapi-gen/requests/services.gen.ts | 4 ++-- airflow/ui/openapi-gen/requests/types.gen.ts | 4 ++-- 8 files changed, 32 insertions(+), 33 deletions(-) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index 7cc7350e91902..6a5cf4f3a1628 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -413,28 +413,25 @@ class _UriPatternSearch(_SearchParam): def __init__(self, skip_none: bool = True) -> None: super().__init__(AssetModel.uri, skip_none) - def depends(self, uri_pattern: str) -> _UriPatternSearch: + def depends(self, uri_pattern: str | None = None) -> _UriPatternSearch: return self.set_value(uri_pattern) -class _DagIdAssetReferenceFilter(BaseParam[str]): +class _DagIdAssetReferenceFilter(BaseParam[list[str]]): """Search on dag_id.""" def __init__(self, skip_none: bool = True) -> None: super().__init__(AssetModel.consuming_dags, skip_none) - self.task_attribute = AssetModel.producing_tasks - def depends(self, dag_ids: str | None = None) -> _DagIdAssetReferenceFilter: + def depends(self, dag_ids: list[str] = Query(None)) -> _DagIdAssetReferenceFilter: return self.set_value(dag_ids) def to_orm(self, select: Select) -> Select: if self.value is None and self.skip_none: return select - if self.value is not None: - dags_list = self.value.split(",") - return select.filter( - (AssetModel.consuming_dags.any(DagScheduleAssetReference.dag_id.in_(dags_list))) - | (AssetModel.producing_tasks.any(TaskOutletAssetReference.dag_id.in_(dags_list))) + return select.where( + (AssetModel.consuming_dags.any(DagScheduleAssetReference.dag_id.in_(self.value))) + | (AssetModel.producing_tasks.any(TaskOutletAssetReference.dag_id.in_(self.value))) ) diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index dd3bedabf0a04..fa3c3085f9cc2 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -3128,17 +3128,19 @@ paths: title: Offset - name: uri_pattern in: query - required: true + required: false schema: - type: string + anyOf: + - type: string + - type: 'null' title: Uri Pattern - name: dag_ids in: query required: false schema: - anyOf: - - type: string - - type: 'null' + type: array + items: + type: string title: Dag Ids - name: order_by in: query diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 01d92f7c0576d..9f5657391bb3f 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -56,12 +56,12 @@ export const UseAssetServiceGetAssetsKeyFn = ( orderBy, uriPattern, }: { - dagIds?: string; + dagIds?: string[]; limit?: number; offset?: number; orderBy?: string; - uriPattern: string; - }, + uriPattern?: string; + } = {}, queryKey?: Array, ) => [ useAssetServiceGetAssetsKey, diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index 274e6062120d1..9354ff89bed55 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -48,9 +48,9 @@ export const prefetchUseAssetServiceNextRunAssets = ( * Get Assets * Get assets. * @param data The data for the request. - * @param data.uriPattern * @param data.limit * @param data.offset + * @param data.uriPattern * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response @@ -65,12 +65,12 @@ export const prefetchUseAssetServiceGetAssets = ( orderBy, uriPattern, }: { - dagIds?: string; + dagIds?: string[]; limit?: number; offset?: number; orderBy?: string; - uriPattern: string; - }, + uriPattern?: string; + } = {}, ) => queryClient.prefetchQuery({ queryKey: Common.UseAssetServiceGetAssetsKeyFn({ diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index e472edba8f4f1..a70f6c5dbc6f0 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -69,9 +69,9 @@ export const useAssetServiceNextRunAssets = < * Get Assets * Get assets. * @param data The data for the request. - * @param data.uriPattern * @param data.limit * @param data.offset + * @param data.uriPattern * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response @@ -89,12 +89,12 @@ export const useAssetServiceGetAssets = < orderBy, uriPattern, }: { - dagIds?: string; + dagIds?: string[]; limit?: number; offset?: number; orderBy?: string; - uriPattern: string; - }, + uriPattern?: string; + } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index 54a3632e95b1b..90d14c8eb6ca9 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -54,9 +54,9 @@ export const useAssetServiceNextRunAssetsSuspense = < * Get Assets * Get assets. * @param data The data for the request. - * @param data.uriPattern * @param data.limit * @param data.offset + * @param data.uriPattern * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response @@ -74,12 +74,12 @@ export const useAssetServiceGetAssetsSuspense = < orderBy, uriPattern, }: { - dagIds?: string; + dagIds?: string[]; limit?: number; offset?: number; orderBy?: string; - uriPattern: string; - }, + uriPattern?: string; + } = {}, queryKey?: TQueryKey, options?: Omit, "queryKey" | "queryFn">, ) => diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index f06a4417f0491..921af0e471286 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -134,16 +134,16 @@ export class AssetService { * Get Assets * Get assets. * @param data The data for the request. - * @param data.uriPattern * @param data.limit * @param data.offset + * @param data.uriPattern * @param data.dagIds * @param data.orderBy * @returns AssetCollectionResponse Successful Response * @throws ApiError */ public static getAssets( - data: GetAssetsData, + data: GetAssetsData = {}, ): CancelablePromise { return __request(OpenAPI, { method: "GET", diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index c84223f95e35a..f750cc258ae87 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -810,11 +810,11 @@ export type NextRunAssetsResponse = { }; export type GetAssetsData = { - dagIds?: string | null; + dagIds?: Array; limit?: number; offset?: number; orderBy?: string; - uriPattern: string; + uriPattern?: string | null; }; export type GetAssetsResponse = AssetCollectionResponse; From 7a97220995f922bb94c82bf38059ca6924b27215 Mon Sep 17 00:00:00 2001 From: Amogh Date: Sat, 9 Nov 2024 08:44:03 +0530 Subject: [PATCH 11/19] fixing unit tests --- airflow/api_fastapi/common/parameters.py | 3 +++ .../core_api/routes/public/test_assets.py | 12 +----------- 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index 6a5cf4f3a1628..b60f5751a828d 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -424,6 +424,9 @@ def __init__(self, skip_none: bool = True) -> None: super().__init__(AssetModel.consuming_dags, skip_none) def depends(self, dag_ids: list[str] = Query(None)) -> _DagIdAssetReferenceFilter: + # needed to handle cases where dag_ids=a1,b1 + if dag_ids and len(dag_ids) == 1 and "," in dag_ids[0]: + dag_ids = dag_ids[0].split(",") return self.set_value(dag_ids) def to_orm(self, select: Select) -> Select: diff --git a/tests/api_fastapi/core_api/routes/public/test_assets.py b/tests/api_fastapi/core_api/routes/public/test_assets.py index c6d085d060399..e243c3c2930de 100644 --- a/tests/api_fastapi/core_api/routes/public/test_assets.py +++ b/tests/api_fastapi/core_api/routes/public/test_assets.py @@ -23,7 +23,6 @@ from airflow.utils import timezone from airflow.utils.session import provide_session -from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import clear_db_assets pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode] @@ -175,6 +174,7 @@ def test_filter_assets_by_dag_ids_works(self, test_client, dag_ids, expected_num "dag_ids, uri_pattern,expected_num", [("dag1,dag2", "folder", 1), ("dag3", "nothing", 0), ("dag2,dag3", "key", 2)], ) + @provide_session def test_filter_assets_by_dag_ids_and_uri_pattern_works( self, test_client, dag_ids, uri_pattern, expected_num, session ): @@ -229,13 +229,3 @@ def test_should_respect_page_size_limit_default(self, test_client): assert response.status_code == 200 assert len(response.json()["assets"]) == 100 - - @conf_vars({("api", "maximum_page_limit"): "150"}) - def test_should_return_conf_max_if_req_max_above_conf(self, test_client): - self.create_assets(num=200) - - # change to 180 once format_parameters is integrated - response = test_client.get("/public/assets?limit=150") - - assert response.status_code == 200 - assert len(response.json()["assets"]) == 150 From 8b6b09ea58b384f96a6cdd07dd6a1ad43fd41856 Mon Sep 17 00:00:00 2001 From: vatsrahul1001 Date: Mon, 11 Nov 2024 17:54:25 +0530 Subject: [PATCH 12/19] migrating get assets events endpoint to fastapi --- .../api_connexion/endpoints/asset_endpoint.py | 1 + airflow/api_fastapi/common/parameters.py | 97 +++++++- .../api_fastapi/core_api/datamodels/assets.py | 44 +++- .../core_api/openapi/v1-generated.yaml | 212 ++++++++++++++++++ .../core_api/routes/public/assets.py | 66 +++++- airflow/ui/openapi-gen/queries/common.ts | 44 ++++ airflow/ui/openapi-gen/queries/prefetch.ts | 60 +++++ airflow/ui/openapi-gen/queries/queries.ts | 69 ++++++ airflow/ui/openapi-gen/queries/suspense.ts | 69 ++++++ .../ui/openapi-gen/requests/schemas.gen.ts | 164 ++++++++++++++ .../ui/openapi-gen/requests/services.gen.ts | 42 ++++ airflow/ui/openapi-gen/requests/types.gen.ts | 80 +++++++ .../core_api/routes/public/test_assets.py | 145 +++++++++++- 13 files changed, 1086 insertions(+), 7 deletions(-) diff --git a/airflow/api_connexion/endpoints/asset_endpoint.py b/airflow/api_connexion/endpoints/asset_endpoint.py index 085817213d0a4..84461de029c8a 100644 --- a/airflow/api_connexion/endpoints/asset_endpoint.py +++ b/airflow/api_connexion/endpoints/asset_endpoint.py @@ -114,6 +114,7 @@ def get_assets( return asset_collection_schema.dump(AssetCollection(assets=assets, total_entries=total_entries)) +@mark_fastapi_migration_done @security.requires_access_asset("GET") @provide_session @format_parameters({"limit": check_limit}) diff --git a/airflow/api_fastapi/common/parameters.py b/airflow/api_fastapi/common/parameters.py index 6a5cf4f3a1628..7361677ebaacf 100644 --- a/airflow/api_fastapi/common/parameters.py +++ b/airflow/api_fastapi/common/parameters.py @@ -30,7 +30,7 @@ from airflow.api_connexion.endpoints.task_instance_endpoint import _convert_ti_states from airflow.models import Base, Connection -from airflow.models.asset import AssetModel, DagScheduleAssetReference, TaskOutletAssetReference +from airflow.models.asset import AssetEvent, AssetModel, DagScheduleAssetReference, TaskOutletAssetReference from airflow.models.dag import DagModel, DagTag from airflow.models.dagrun import DagRun from airflow.models.dagwarning import DagWarning, DagWarningType @@ -435,6 +435,86 @@ def to_orm(self, select: Select) -> Select: ) +class _AssetIdFilter(BaseParam[int]): + """Filter on asset_id.""" + + def __init__(self, attribute: ColumnElement, skip_none: bool = True) -> None: + super().__init__(skip_none=skip_none) + self.attribute = attribute + + def to_orm(self, select: Select) -> Select: + if self.value is None and self.skip_none: + return select + return select.where(self.attribute == self.value) + + def depends(self, asset_id: int | None = None) -> _AssetIdFilter: + return self.set_value(asset_id) + + +class _SourceDagIdFilter(BaseParam[str]): + """Filter on source_dag_id.""" + + def __init__(self, attribute: ColumnElement, skip_none: bool = True) -> None: + super().__init__(skip_none=skip_none) + self.attribute = attribute + + def to_orm(self, select: Select) -> Select: + if self.value is None and self.skip_none: + return select + return select.where(self.attribute == self.value) + + def depends(self, source_dag_id: str | None = None) -> _SourceDagIdFilter: + return self.set_value(source_dag_id) + + +class _SourceTaskIdFilter(BaseParam[str]): + """Filter on source_task_id.""" + + def __init__(self, attribute: ColumnElement, skip_none: bool = True) -> None: + super().__init__(skip_none=skip_none) + self.attribute = attribute + + def to_orm(self, select: Select) -> Select: + if self.value is None and self.skip_none: + return select + return select.where(self.attribute == self.value) + + def depends(self, source_task_id: str | None = None) -> _SourceTaskIdFilter: + return self.set_value(source_task_id) + + +class _SourceRunIdFilter(BaseParam[str]): + """filter on source_run_id.""" + + def __init__(self, attribute: ColumnElement, skip_none: bool = True) -> None: + super().__init__(skip_none=skip_none) + self.attribute = attribute + + def to_orm(self, select: Select) -> Select: + if self.value is None and self.skip_none: + return select + return select.where(self.attribute == self.value) + + def depends(self, source_run_id: str | None = None) -> _SourceRunIdFilter: + return self.set_value(source_run_id) + + +class _SourceMapIndexFilter(BaseParam[int]): + """Filter on source_map_index.""" + + def __init__(self, attribute: ColumnElement, skip_none: bool = True) -> None: + super().__init__(skip_none=skip_none) + self.attribute = attribute + + def to_orm(self, select: Select) -> Select: + if self.value is None and self.skip_none: + return select + return select.where(self.attribute == self.value) + + def depends(self, source_map_index: int | None = None) -> _SourceMapIndexFilter: + return self.set_value(source_map_index) + + class Range(BaseModel, Generic[T]): """Range with a lower and upper bound.""" @@ -532,3 +612,18 @@ def depends_float( QueryAssetDagIdPatternSearch = Annotated[ _DagIdAssetReferenceFilter, Depends(_DagIdAssetReferenceFilter().depends) ] +QueryAssetIdFilter = Annotated[_AssetIdFilter, Depends(_AssetIdFilter(AssetEvent.asset_id).depends)] + + +QuerySourceDagIdFilter = Annotated[ + _SourceDagIdFilter, Depends(_SourceDagIdFilter(AssetEvent.source_dag_id).depends) +] +QuerySourceTaskIdFilter = Annotated[ + _SourceTaskIdFilter, Depends(_SourceTaskIdFilter(AssetEvent.source_task_id).depends) +] +QuerySourceRunIdFilter = Annotated[ + _SourceRunIdFilter, Depends(_SourceRunIdFilter(AssetEvent.source_run_id).depends) +] +QuerySourceMapIndexFilter = Annotated[ + _SourceMapIndexFilter, Depends(_SourceMapIndexFilter(AssetEvent.source_map_index).depends) +] diff --git a/airflow/api_fastapi/core_api/datamodels/assets.py b/airflow/api_fastapi/core_api/datamodels/assets.py index 498dc4edae6fd..11f828eb9dbd7 100644 --- a/airflow/api_fastapi/core_api/datamodels/assets.py +++ b/airflow/api_fastapi/core_api/datamodels/assets.py @@ -19,7 +19,7 @@ from datetime import datetime -from pydantic import BaseModel +from pydantic import BaseModel, Field, model_validator class DagScheduleAssetReference(BaseModel): @@ -64,3 +64,45 @@ class AssetCollectionResponse(BaseModel): assets: list[AssetResponse] total_entries: int + + +class DagRunAssetReference(BaseModel): + """Serializable version of the DagRunAssetReference ORM SqlAlchemyModel.""" + + run_id: str + dag_id: str + execution_date: datetime = Field(alias="logical_date") + start_date: datetime + end_date: datetime + state: str + data_interval_start: datetime + data_interval_end: datetime + + +class AssetEventResponse(BaseModel): + """Asset event serializer for responses.""" + + id: int + asset_id: int + asset_uri: str + extra: dict | None = None + source_task_id: str | None = None + source_dag_id: str | None = None + source_run_id: str | None = None + source_map_index: int + created_dagruns: list[DagRunAssetReference] + timestamp: datetime + + @model_validator(mode="before") + def rename_uri_to_asset_uri(cls, values): + """Rename 'uri' to 'asset_uri' during serialization.""" + if hasattr(values, "uri") and values.uri: + values.asset_uri = values.uri + return values + + +class AssetEventCollectionResponse(BaseModel): + """Asset collection response.""" + + asset_events: list[AssetEventResponse] + total_entries: int diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index f1dd5d6ba4332..0648605a987f4 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -3180,6 +3180,106 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' + /public/assets/events: + get: + tags: + - Asset + summary: Get Asset Events + description: Get assets events. + operationId: get_asset_events + 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: timestamp + title: Order By + - name: asset_id + in: query + required: false + schema: + anyOf: + - type: integer + - type: 'null' + title: Asset Id + - name: source_dag_id + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + title: Source Dag Id + - name: source_task_id + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + title: Source Task Id + - name: source_run_id + in: query + required: false + schema: + anyOf: + - type: string + - type: 'null' + title: Source Run Id + - name: source_map_index + in: query + required: false + schema: + anyOf: + - type: integer + - type: 'null' + title: Source Map Index + responses: + '200': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/AssetEventCollectionResponse' + '401': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Unauthorized + '403': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Forbidden + '404': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Not Found + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' components: schemas: AppBuilderMenuItemResponse: @@ -3259,6 +3359,75 @@ components: - total_entries title: AssetCollectionResponse description: Asset collection response. + AssetEventCollectionResponse: + properties: + asset_events: + items: + $ref: '#/components/schemas/AssetEventResponse' + type: array + title: Asset Events + total_entries: + type: integer + title: Total Entries + type: object + required: + - asset_events + - total_entries + title: AssetEventCollectionResponse + description: Asset collection response. + AssetEventResponse: + properties: + id: + type: integer + title: Id + asset_id: + type: integer + title: Asset Id + asset_uri: + type: string + title: Asset Uri + extra: + anyOf: + - type: object + - type: 'null' + title: Extra + source_task_id: + anyOf: + - type: string + - type: 'null' + title: Source Task Id + source_dag_id: + anyOf: + - type: string + - type: 'null' + title: Source Dag Id + source_run_id: + anyOf: + - type: string + - type: 'null' + title: Source Run Id + source_map_index: + type: integer + title: Source Map Index + created_dagruns: + items: + $ref: '#/components/schemas/DagRunAssetReference' + type: array + title: Created Dagruns + timestamp: + type: string + format: date-time + title: Timestamp + type: object + required: + - id + - asset_id + - asset_uri + - source_map_index + - created_dagruns + - timestamp + title: AssetEventResponse + description: Asset event serializer for responses. AssetResponse: properties: id: @@ -4230,6 +4399,49 @@ components: - latest_dag_processor_heartbeat title: DagProcessorInfoSchema description: Schema for DagProcessor info. + DagRunAssetReference: + properties: + run_id: + type: string + title: Run Id + dag_id: + type: string + title: Dag Id + logical_date: + type: string + format: date-time + title: Logical Date + start_date: + type: string + format: date-time + title: Start Date + end_date: + type: string + format: date-time + title: End Date + state: + type: string + title: State + data_interval_start: + type: string + format: date-time + title: Data Interval Start + data_interval_end: + type: string + format: date-time + title: Data Interval End + type: object + required: + - run_id + - dag_id + - logical_date + - start_date + - end_date + - state + - data_interval_start + - data_interval_end + title: DagRunAssetReference + description: Serializable version of the DagRunAssetReference ORM SqlAlchemyModel. DagRunState: type: string enum: diff --git a/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow/api_fastapi/core_api/routes/public/assets.py index 7ec7012de8d37..b3e6e35e5571a 100644 --- a/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/airflow/api_fastapi/core_api/routes/public/assets.py @@ -26,15 +26,25 @@ from airflow.api_fastapi.common.db.common import get_session, paginated_select from airflow.api_fastapi.common.parameters import ( QueryAssetDagIdPatternSearch, + QueryAssetIdFilter, QueryLimit, QueryOffset, + QuerySourceDagIdFilter, + QuerySourceMapIndexFilter, + QuerySourceRunIdFilter, + QuerySourceTaskIdFilter, QueryUriPatternSearch, SortParam, ) from airflow.api_fastapi.common.router import AirflowRouter -from airflow.api_fastapi.core_api.datamodels.assets import AssetCollectionResponse, AssetResponse +from airflow.api_fastapi.core_api.datamodels.assets import ( + AssetCollectionResponse, + AssetEventCollectionResponse, + AssetEventResponse, + AssetResponse, +) from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc -from airflow.models.asset import AssetModel +from airflow.models.asset import AssetEvent, AssetModel assets_router = AirflowRouter(tags=["Asset"], prefix="/assets") @@ -50,7 +60,9 @@ def get_assets( dag_ids: QueryAssetDagIdPatternSearch, order_by: Annotated[ SortParam, - Depends(SortParam(["id", "uri", "created_at", "updated_at"], AssetModel).dynamic_depends()), + Depends( + SortParam(["id", "uri", "created_at", "updated_at", "timestamp"], AssetModel).dynamic_depends() + ), ], session: Annotated[Session, Depends(get_session)], ) -> AssetCollectionResponse: @@ -65,7 +77,53 @@ def get_assets( ) assets = session.scalars(assets_select).all() + print(f"assets are {assets}") + assets = [AssetResponse.model_validate(asset, from_attributes=True) for asset in assets] + print(f"updated assetss are {assets}") return AssetCollectionResponse( - assets=[AssetResponse.model_validate(asset, from_attributes=True) for asset in assets], + assets=assets, + total_entries=total_entries, + ) + + +@assets_router.get( + "/events", + responses=create_openapi_http_exception_doc([401, 403, 404]), +) +async def get_asset_events( + limit: QueryLimit, + offset: QueryOffset, + order_by: Annotated[ + SortParam, + Depends( + SortParam( + ["timestamp", "source_dag_id", "source_task_id", "source_run_id", "source_map_index"], + AssetEvent, + ).dynamic_depends("timestamp") + ), + ], + asset_id: QueryAssetIdFilter, + source_dag_id: QuerySourceDagIdFilter, + source_task_id: QuerySourceTaskIdFilter, + source_run_id: QuerySourceRunIdFilter, + source_map_index: QuerySourceMapIndexFilter, + session: Annotated[Session, Depends(get_session)], +) -> AssetEventCollectionResponse: + """Get assets events.""" + assets_event_select, total_entries = paginated_select( + select(AssetEvent), + filters=[asset_id, source_dag_id, source_task_id, source_run_id, source_map_index], + order_by=order_by, + offset=offset, + limit=limit, + session=session, + ) + + assets_events = session.scalars(assets_event_select).all() + + return AssetEventCollectionResponse( + asset_events=[ + AssetEventResponse.model_validate(asset, from_attributes=True) for asset in assets_events + ], total_entries=total_entries, ) diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index 5077e55c3902d..5e31d47b2949d 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -67,6 +67,50 @@ export const UseAssetServiceGetAssetsKeyFn = ( useAssetServiceGetAssetsKey, ...(queryKey ?? [{ dagIds, limit, offset, orderBy, uriPattern }]), ]; +export type AssetServiceGetAssetEventsDefaultResponse = Awaited< + ReturnType +>; +export type AssetServiceGetAssetEventsQueryResult< + TData = AssetServiceGetAssetEventsDefaultResponse, + TError = unknown, +> = UseQueryResult; +export const useAssetServiceGetAssetEventsKey = "AssetServiceGetAssetEvents"; +export const UseAssetServiceGetAssetEventsKeyFn = ( + { + assetId, + limit, + offset, + orderBy, + sourceDagId, + sourceMapIndex, + sourceRunId, + sourceTaskId, + }: { + assetId?: number; + limit?: number; + offset?: number; + orderBy?: string; + sourceDagId?: string; + sourceMapIndex?: number; + sourceRunId?: string; + sourceTaskId?: string; + } = {}, + queryKey?: Array, +) => [ + useAssetServiceGetAssetEventsKey, + ...(queryKey ?? [ + { + assetId, + limit, + offset, + orderBy, + sourceDagId, + sourceMapIndex, + sourceRunId, + sourceTaskId, + }, + ]), +]; export type DashboardServiceHistoricalMetricsDefaultResponse = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index 9354ff89bed55..37e82d84e51dd 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -83,6 +83,66 @@ export const prefetchUseAssetServiceGetAssets = ( queryFn: () => AssetService.getAssets({ dagIds, limit, offset, orderBy, uriPattern }), }); +/** + * Get Asset Events + * Get assets events. + * @param data The data for the request. + * @param data.limit + * @param data.offset + * @param data.orderBy + * @param data.assetId + * @param data.sourceDagId + * @param data.sourceTaskId + * @param data.sourceRunId + * @param data.sourceMapIndex + * @returns AssetEventCollectionResponse Successful Response + * @throws ApiError + */ +export const prefetchUseAssetServiceGetAssetEvents = ( + queryClient: QueryClient, + { + assetId, + limit, + offset, + orderBy, + sourceDagId, + sourceMapIndex, + sourceRunId, + sourceTaskId, + }: { + assetId?: number; + limit?: number; + offset?: number; + orderBy?: string; + sourceDagId?: string; + sourceMapIndex?: number; + sourceRunId?: string; + sourceTaskId?: string; + } = {}, +) => + queryClient.prefetchQuery({ + queryKey: Common.UseAssetServiceGetAssetEventsKeyFn({ + assetId, + limit, + offset, + orderBy, + sourceDagId, + sourceMapIndex, + sourceRunId, + sourceTaskId, + }), + queryFn: () => + AssetService.getAssetEvents({ + assetId, + limit, + offset, + orderBy, + sourceDagId, + sourceMapIndex, + sourceRunId, + sourceTaskId, + }), + }); /** * Historical Metrics * Return cluster activity historical metrics. diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 0e43d66d1397a..ce23255b264b0 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -113,6 +113,75 @@ export const useAssetServiceGetAssets = < }) as TData, ...options, }); +/** + * Get Asset Events + * Get assets events. + * @param data The data for the request. + * @param data.limit + * @param data.offset + * @param data.orderBy + * @param data.assetId + * @param data.sourceDagId + * @param data.sourceTaskId + * @param data.sourceRunId + * @param data.sourceMapIndex + * @returns AssetEventCollectionResponse Successful Response + * @throws ApiError + */ +export const useAssetServiceGetAssetEvents = < + TData = Common.AssetServiceGetAssetEventsDefaultResponse, + TError = unknown, + TQueryKey extends Array = unknown[], +>( + { + assetId, + limit, + offset, + orderBy, + sourceDagId, + sourceMapIndex, + sourceRunId, + sourceTaskId, + }: { + assetId?: number; + limit?: number; + offset?: number; + orderBy?: string; + sourceDagId?: string; + sourceMapIndex?: number; + sourceRunId?: string; + sourceTaskId?: string; + } = {}, + queryKey?: TQueryKey, + options?: Omit, "queryKey" | "queryFn">, +) => + useQuery({ + queryKey: Common.UseAssetServiceGetAssetEventsKeyFn( + { + assetId, + limit, + offset, + orderBy, + sourceDagId, + sourceMapIndex, + sourceRunId, + sourceTaskId, + }, + queryKey, + ), + queryFn: () => + AssetService.getAssetEvents({ + assetId, + limit, + offset, + orderBy, + sourceDagId, + sourceMapIndex, + sourceRunId, + sourceTaskId, + }) as TData, + ...options, + }); /** * Historical Metrics * Return cluster activity historical metrics. diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index 90d14c8eb6ca9..d0655e110f8f4 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -98,6 +98,75 @@ export const useAssetServiceGetAssetsSuspense = < }) as TData, ...options, }); +/** + * Get Asset Events + * Get assets events. + * @param data The data for the request. + * @param data.limit + * @param data.offset + * @param data.orderBy + * @param data.assetId + * @param data.sourceDagId + * @param data.sourceTaskId + * @param data.sourceRunId + * @param data.sourceMapIndex + * @returns AssetEventCollectionResponse Successful Response + * @throws ApiError + */ +export const useAssetServiceGetAssetEventsSuspense = < + TData = Common.AssetServiceGetAssetEventsDefaultResponse, + TError = unknown, + TQueryKey extends Array = unknown[], +>( + { + assetId, + limit, + offset, + orderBy, + sourceDagId, + sourceMapIndex, + sourceRunId, + sourceTaskId, + }: { + assetId?: number; + limit?: number; + offset?: number; + orderBy?: string; + sourceDagId?: string; + sourceMapIndex?: number; + sourceRunId?: string; + sourceTaskId?: string; + } = {}, + queryKey?: TQueryKey, + options?: Omit, "queryKey" | "queryFn">, +) => + useSuspenseQuery({ + queryKey: Common.UseAssetServiceGetAssetEventsKeyFn( + { + assetId, + limit, + offset, + orderBy, + sourceDagId, + sourceMapIndex, + sourceRunId, + sourceTaskId, + }, + queryKey, + ), + queryFn: () => + AssetService.getAssetEvents({ + assetId, + limit, + offset, + orderBy, + sourceDagId, + sourceMapIndex, + sourceRunId, + sourceTaskId, + }) as TData, + ...options, + }); /** * Historical Metrics * Return cluster activity historical metrics. diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index 36a2b26ebd289..22d8ff0f7570a 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -127,6 +127,114 @@ export const $AssetCollectionResponse = { description: "Asset collection response.", } as const; +export const $AssetEventCollectionResponse = { + properties: { + asset_events: { + items: { + $ref: "#/components/schemas/AssetEventResponse", + }, + type: "array", + title: "Asset Events", + }, + total_entries: { + type: "integer", + title: "Total Entries", + }, + }, + type: "object", + required: ["asset_events", "total_entries"], + title: "AssetEventCollectionResponse", + description: "Asset collection response.", +} as const; + +export const $AssetEventResponse = { + properties: { + id: { + type: "integer", + title: "Id", + }, + asset_id: { + type: "integer", + title: "Asset Id", + }, + asset_uri: { + type: "string", + title: "Asset Uri", + }, + extra: { + anyOf: [ + { + type: "object", + }, + { + type: "null", + }, + ], + title: "Extra", + }, + source_task_id: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Source Task Id", + }, + source_dag_id: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Source Dag Id", + }, + source_run_id: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Source Run Id", + }, + source_map_index: { + type: "integer", + title: "Source Map Index", + }, + created_dagruns: { + items: { + $ref: "#/components/schemas/DagRunAssetReference", + }, + type: "array", + title: "Created Dagruns", + }, + timestamp: { + type: "string", + format: "date-time", + title: "Timestamp", + }, + }, + type: "object", + required: [ + "id", + "asset_id", + "asset_uri", + "source_map_index", + "created_dagruns", + "timestamp", + ], + title: "AssetEventResponse", + description: "Asset event serializer for responses.", +} as const; + export const $AssetResponse = { properties: { id: { @@ -1669,6 +1777,62 @@ export const $DagProcessorInfoSchema = { description: "Schema for DagProcessor info.", } as const; +export const $DagRunAssetReference = { + properties: { + run_id: { + type: "string", + title: "Run Id", + }, + dag_id: { + type: "string", + title: "Dag Id", + }, + logical_date: { + type: "string", + format: "date-time", + title: "Logical Date", + }, + start_date: { + type: "string", + format: "date-time", + title: "Start Date", + }, + end_date: { + type: "string", + format: "date-time", + title: "End Date", + }, + state: { + type: "string", + title: "State", + }, + data_interval_start: { + type: "string", + format: "date-time", + title: "Data Interval Start", + }, + data_interval_end: { + type: "string", + format: "date-time", + title: "Data Interval End", + }, + }, + type: "object", + required: [ + "run_id", + "dag_id", + "logical_date", + "start_date", + "end_date", + "state", + "data_interval_start", + "data_interval_end", + ], + title: "DagRunAssetReference", + description: + "Serializable version of the DagRunAssetReference ORM SqlAlchemyModel.", +} as const; + export const $DagRunState = { type: "string", enum: ["queued", "running", "success", "failed"], diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index b2ab26989fbdb..7c7a6bf1d205b 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -7,6 +7,8 @@ import type { NextRunAssetsResponse, GetAssetsData, GetAssetsResponse, + GetAssetEventsData, + GetAssetEventsResponse, HistoricalMetricsData, HistoricalMetricsResponse, RecentDagRunsData, @@ -163,6 +165,46 @@ export class AssetService { }, }); } + + /** + * Get Asset Events + * Get assets events. + * @param data The data for the request. + * @param data.limit + * @param data.offset + * @param data.orderBy + * @param data.assetId + * @param data.sourceDagId + * @param data.sourceTaskId + * @param data.sourceRunId + * @param data.sourceMapIndex + * @returns AssetEventCollectionResponse Successful Response + * @throws ApiError + */ + public static getAssetEvents( + data: GetAssetEventsData = {}, + ): CancelablePromise { + return __request(OpenAPI, { + method: "GET", + url: "/public/assets/events", + query: { + limit: data.limit, + offset: data.offset, + order_by: data.orderBy, + asset_id: data.assetId, + source_dag_id: data.sourceDagId, + source_task_id: data.sourceTaskId, + source_run_id: data.sourceRunId, + source_map_index: data.sourceMapIndex, + }, + errors: { + 401: "Unauthorized", + 403: "Forbidden", + 404: "Not Found", + 422: "Validation Error", + }, + }); + } } export class DashboardService { diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index c609269739b2a..011a9b3eba806 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -37,6 +37,32 @@ export type AssetCollectionResponse = { total_entries: number; }; +/** + * Asset collection response. + */ +export type AssetEventCollectionResponse = { + asset_events: Array; + total_entries: number; +}; + +/** + * Asset event serializer for responses. + */ +export type AssetEventResponse = { + id: number; + asset_id: number; + asset_uri: string; + extra?: { + [key: string]: unknown; + } | null; + source_task_id?: string | null; + source_dag_id?: string | null; + source_run_id?: string | null; + source_map_index: number; + created_dagruns: Array; + timestamp: string; +}; + /** * Asset serializer for responses. */ @@ -349,6 +375,20 @@ export type DagProcessorInfoSchema = { latest_dag_processor_heartbeat: string | null; }; +/** + * Serializable version of the DagRunAssetReference ORM SqlAlchemyModel. + */ +export type DagRunAssetReference = { + run_id: string; + dag_id: string; + logical_date: string; + start_date: string; + end_date: string; + state: string; + data_interval_start: string; + data_interval_end: string; +}; + /** * All possible states that a DagRun can be in. * @@ -819,6 +859,19 @@ export type GetAssetsData = { export type GetAssetsResponse = AssetCollectionResponse; +export type GetAssetEventsData = { + assetId?: number | null; + limit?: number; + offset?: number; + orderBy?: string; + sourceDagId?: string | null; + sourceMapIndex?: number | null; + sourceRunId?: string | null; + sourceTaskId?: string | null; +}; + +export type GetAssetEventsResponse = AssetEventCollectionResponse; + export type HistoricalMetricsData = { endDate: string; startDate: string; @@ -1281,6 +1334,33 @@ export type $OpenApiTs = { }; }; }; + "/public/assets/events": { + get: { + req: GetAssetEventsData; + res: { + /** + * Successful Response + */ + 200: AssetEventCollectionResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Not Found + */ + 404: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; + }; "/ui/dashboard/historical_metrics_data": { get: { req: HistoricalMetricsData; diff --git a/tests/api_fastapi/core_api/routes/public/test_assets.py b/tests/api_fastapi/core_api/routes/public/test_assets.py index c6d085d060399..8c084b810ed9d 100644 --- a/tests/api_fastapi/core_api/routes/public/test_assets.py +++ b/tests/api_fastapi/core_api/routes/public/test_assets.py @@ -19,9 +19,12 @@ import pytest from airflow.models import DagModel -from airflow.models.asset import AssetModel, DagScheduleAssetReference, TaskOutletAssetReference +from airflow.models.asset import AssetEvent, AssetModel, DagScheduleAssetReference, TaskOutletAssetReference +from airflow.models.dagrun import DagRun from airflow.utils import timezone from airflow.utils.session import provide_session +from airflow.utils.state import DagRunState +from airflow.utils.types import DagRunType from tests_common.test_utils.config import conf_vars from tests_common.test_utils.db import clear_db_assets @@ -50,6 +53,59 @@ def _create_provided_asset(session, asset: AssetModel) -> None: session.commit() +def _create_assets_events(session, num: int = 2) -> None: + default_time = "2020-06-11T18:00:00+00:00" + assets_events = [ + AssetEvent( + id=i, + asset_id=i, + extra={"foo": "bar"}, + source_task_id="source_task_id", + source_dag_id="source_dag_id", + source_run_id=f"source_run_id_{i}", + timestamp=timezone.parse(default_time), + ) + for i in range(1, 1 + num) + ] + session.add_all(assets_events) + session.commit() + + +def _create_provided_asset_event(session, asset_event: AssetEvent) -> None: + session.add(asset_event) + session.commit() + + +def _create_dag_run(session, num: int = 2): + default_time = "2020-06-11T18:00:00+00:00" + dag_runs = [ + DagRun( + dag_id="source_dag_id", + run_id=f"source_run_id_{i}", + run_type=DagRunType.MANUAL, + execution_date=timezone.parse(default_time), + start_date=timezone.parse(default_time), + data_interval=(timezone.parse(default_time), timezone.parse(default_time)), + external_trigger=True, + state=DagRunState.SUCCESS, + ) + for i in range(1, 1 + num) + ] + for dag_run in dag_runs: + dag_run.end_date = timezone.parse(default_time) + session.add_all(dag_runs) + session.commit() + + +def _create_asset_dag_run(session, num: int = 2): + for i in range(1, 1 + num): + dag_run = session.query(DagRun).filter_by(run_id=f"source_run_id_{i}").first() + asset_event = session.query(AssetEvent).filter_by(id=i).first() + if dag_run and asset_event: + dag_run.consumed_asset_events.append(asset_event) + session.commit() + + class TestAssets: default_time = "2020-06-11T18:00:00+00:00" @@ -68,6 +124,22 @@ def create_assets(self, session, num: int = 2): def create_provided_asset(self, session, asset: AssetModel): _create_provided_asset(session=session, asset=asset) + @provide_session + def create_assets_events(self, session, num: int = 2): + _create_assets_events(session=session, num=num) + + @provide_session + def create_provided_asset_event(self, session, asset_event: AssetEvent): + _create_provided_asset_event(session=session, asset_event=asset_event) + + @provide_session + def create_dag_run(self, session, num: int = 2): + _create_dag_run(num=num, session=session) + + @provide_session + def create_asset_dag_run(self, session, num: int = 2): + _create_asset_dag_run(num=num, session=session) + class TestGetAssets(TestAssets): def test_should_respond_200(self, test_client, session): @@ -239,3 +311,74 @@ def test_should_return_conf_max_if_req_max_above_conf(self, test_client): assert response.status_code == 200 assert len(response.json()["assets"]) == 150 + + +class TestGetAssetsEvents(TestAssets): + def test_should_respond_200(self, test_client, session): + self.create_assets() + self.create_assets_events() + self.create_dag_run() + self.create_asset_dag_run() + assets = session.query(AssetEvent).all() + assert len(assets) == 2 + response = test_client.get("/public/assets/events") + assert response.status_code == 200 + response_data = response.json() + assert response_data == { + "asset_events": [ + { + "id": 1, + "asset_id": 1, + "uri": "s3://bucket/key/1", + "extra": {"foo": "bar"}, + "source_task_id": "source_task_id", + "source_dag_id": "source_dag_id", + "source_run_id": "source_run_id_1", + "source_map_index": -1, + "created_dagruns": [ + { + "run_id": "source_run_id_1", + "dag_id": "source_dag_id", + "logical_date": "2020-06-11T18:00:00Z", + "start_date": "2020-06-11T18:00:00Z", + "end_date": "2020-06-11T18:00:00Z", + "state": "success", + "data_interval_start": "2020-06-11T18:00:00Z", + "data_interval_end": "2020-06-11T18:00:00Z", + } + ], + "timestamp": "2020-06-11T18:00:00Z", + }, + { + "id": 2, + "asset_id": 2, + "uri": "s3://bucket/key/2", + "extra": {"foo": "bar"}, + "source_task_id": "source_task_id", + "source_dag_id": "source_dag_id", + "source_run_id": "source_run_id_2", + "source_map_index": -1, + "created_dagruns": [ + { + "run_id": "source_run_id_2", + "dag_id": "source_dag_id", + "logical_date": "2020-06-11T18:00:00Z", + "start_date": "2020-06-11T18:00:00Z", + "end_date": "2020-06-11T18:00:00Z", + "state": "success", + "data_interval_start": "2020-06-11T18:00:00Z", + "data_interval_end": "2020-06-11T18:00:00Z", + } + ], + "timestamp": "2020-06-11T18:00:00Z", + }, + ], + "total_entries": 2, + } + + def test_order_by_raises_400_for_invalid_attr(self, test_client, session): + response = test_client.get("/public/assets/events?order_by=fake") + + assert response.status_code == 400 + msg = "Ordering with 'fake' is disallowed or the attribute does not exist on the model" + assert response.json()["detail"] == msg From b808895e327bdcce39bfa8403968bcb679e19486 Mon Sep 17 00:00:00 2001 From: vatsrahul1001 Date: Mon, 11 Nov 2024 18:32:07 +0530 Subject: [PATCH 13/19] fixing test response --- tests/api_fastapi/core_api/routes/public/test_assets.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/api_fastapi/core_api/routes/public/test_assets.py b/tests/api_fastapi/core_api/routes/public/test_assets.py index 8c084b810ed9d..19c9f89b3b635 100644 --- a/tests/api_fastapi/core_api/routes/public/test_assets.py +++ b/tests/api_fastapi/core_api/routes/public/test_assets.py @@ -329,7 +329,7 @@ def test_should_respond_200(self, test_client, session): { "id": 1, "asset_id": 1, - "uri": "s3://bucket/key/1", + "asset_uri": "s3://bucket/key/1", "extra": {"foo": "bar"}, "source_task_id": "source_task_id", "source_dag_id": "source_dag_id", @@ -352,7 +352,7 @@ def test_should_respond_200(self, test_client, session): { "id": 2, "asset_id": 2, - "uri": "s3://bucket/key/2", + "asset_uri": "s3://bucket/key/2", "extra": {"foo": "bar"}, "source_task_id": "source_task_id", "source_dag_id": "source_dag_id", From 52939e5d9fc8564a23a53bc89675b118dabeb501 Mon Sep 17 00:00:00 2001 From: vatsrahul1001 Date: Tue, 12 Nov 2024 11:17:41 +0530 Subject: [PATCH 14/19] Adding tests for filtering --- .../core_api/routes/public/test_assets.py | 24 ++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/tests/api_fastapi/core_api/routes/public/test_assets.py b/tests/api_fastapi/core_api/routes/public/test_assets.py index 19c9f89b3b635..669363f5461de 100644 --- a/tests/api_fastapi/core_api/routes/public/test_assets.py +++ b/tests/api_fastapi/core_api/routes/public/test_assets.py @@ -27,7 +27,7 @@ from airflow.utils.types import DagRunType from tests_common.test_utils.config import conf_vars -from tests_common.test_utils.db import clear_db_assets +from tests_common.test_utils.db import clear_db_assets, clear_db_runs pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode] @@ -112,9 +112,11 @@ class TestAssets: @pytest.fixture(autouse=True) def setup(self) -> None: clear_db_assets() + clear_db_runs() def teardown_method(self) -> None: clear_db_assets() + clear_db_runs() @provide_session def create_assets(self, session, num: int = 2): @@ -382,3 +384,23 @@ def test_order_by_raises_400_for_invalid_attr(self, test_client, session): assert response.status_code == 400 msg = "Ordering with 'fake' is disallowed or the attribute does not exist on the model" assert response.json()["detail"] == msg + + @pytest.mark.parametrize( + "filter_type, filter_value, total_entries", + [ + ("asset_id", "2", 1), + ("source_dag_id", "source_dag_id", 2), + ("source_task_id", "source_task_id", 2), + ("source_run_id", "source_run_id_1", 1), + ("source_map_index", "-1", 2), + ], + ) + @provide_session + def test_filter_events_by_asset_id(self, test_client, filter_type, filter_value, total_entries, session): + self.create_assets() + self.create_assets_events() + self.create_dag_run() + self.create_asset_dag_run() + response = test_client.get(f"/public/assets/events?{filter_type}={filter_value}") + assert response.status_code == 200 + assert response.json()["total_entries"] == total_entries From 255dc815b727c5f7554121d7871f55e380dd535e Mon Sep 17 00:00:00 2001 From: vatsrahul1001 Date: Wed, 13 Nov 2024 15:30:47 +0530 Subject: [PATCH 15/19] address review comments --- .../api_fastapi/core_api/datamodels/assets.py | 2 +- .../core_api/openapi/v1-generated.yaml | 2 +- .../core_api/routes/public/__init__.py | 1 - .../core_api/routes/public/assets.py | 11 ++++++-- airflow/ui/openapi-gen/queries/prefetch.ts | 2 +- airflow/ui/openapi-gen/queries/queries.ts | 2 +- airflow/ui/openapi-gen/queries/suspense.ts | 2 +- .../ui/openapi-gen/requests/services.gen.ts | 2 +- .../core_api/routes/public/test_assets.py | 25 +++++++++++-------- 9 files changed, 30 insertions(+), 19 deletions(-) diff --git a/airflow/api_fastapi/core_api/datamodels/assets.py b/airflow/api_fastapi/core_api/datamodels/assets.py index 11f828eb9dbd7..0beeead27bc9a 100644 --- a/airflow/api_fastapi/core_api/datamodels/assets.py +++ b/airflow/api_fastapi/core_api/datamodels/assets.py @@ -95,7 +95,7 @@ class AssetEventResponse(BaseModel): @model_validator(mode="before") def rename_uri_to_asset_uri(cls, values): - """Rename 'uri' to 'asset_uri' during serialization.""" + """Rename 'uri' to 'asset_uri' during serialization to match legacy response.""" if hasattr(values, "uri") and values.uri: values.asset_uri = values.uri return values diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 671712f67bb9f..71f5d3d5432cf 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -3232,7 +3232,7 @@ paths: tags: - Asset summary: Get Asset Events - description: Get assets events. + description: Get asset events. operationId: get_asset_events parameters: - name: limit diff --git a/airflow/api_fastapi/core_api/routes/public/__init__.py b/airflow/api_fastapi/core_api/routes/public/__init__.py index 17ebacd199322..cdb1ed375a6a3 100644 --- a/airflow/api_fastapi/core_api/routes/public/__init__.py +++ b/airflow/api_fastapi/core_api/routes/public/__init__.py @@ -60,4 +60,3 @@ public_router.include_router(dag_stats_router) public_router.include_router(assets_router) public_router.include_router(xcom_router) -public_router.include_router(assets_router) diff --git a/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow/api_fastapi/core_api/routes/public/assets.py index 8ce9fa7686e7f..4e5faa686dc4f 100644 --- a/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/airflow/api_fastapi/core_api/routes/public/assets.py @@ -92,7 +92,14 @@ async def get_asset_events( SortParam, Depends( SortParam( - ["timestamp", "source_dag_id", "source_task_id", "source_run_id", "source_map_index"], + [ + "asset_id", + "source_task_id", + "source_dag_id", + "source_run_id", + "source_map_index", + "timestamp", + ], AssetEvent, ).dynamic_depends("timestamp") ), @@ -104,7 +111,7 @@ async def get_asset_events( source_map_index: QuerySourceMapIndexFilter, session: Annotated[Session, Depends(get_session)], ) -> AssetEventCollectionResponse: - """Get assets events.""" + """Get asset events.""" assets_event_select, total_entries = paginated_select( select(AssetEvent), filters=[asset_id, source_dag_id, source_task_id, source_run_id, source_map_index], diff --git a/airflow/ui/openapi-gen/queries/prefetch.ts b/airflow/ui/openapi-gen/queries/prefetch.ts index 7710851cf49c7..b334546dd8c07 100644 --- a/airflow/ui/openapi-gen/queries/prefetch.ts +++ b/airflow/ui/openapi-gen/queries/prefetch.ts @@ -86,7 +86,7 @@ export const prefetchUseAssetServiceGetAssets = ( }); /** * Get Asset Events - * Get assets events. + * Get asset events. * @param data The data for the request. * @param data.limit * @param data.offset diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 172514ffa1c99..d5ecb4aa9ffff 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -116,7 +116,7 @@ export const useAssetServiceGetAssets = < }); /** * Get Asset Events - * Get assets events. + * Get asset events. * @param data The data for the request. * @param data.limit * @param data.offset diff --git a/airflow/ui/openapi-gen/queries/suspense.ts b/airflow/ui/openapi-gen/queries/suspense.ts index 59587065ae818..2e831cae96e8a 100644 --- a/airflow/ui/openapi-gen/queries/suspense.ts +++ b/airflow/ui/openapi-gen/queries/suspense.ts @@ -101,7 +101,7 @@ export const useAssetServiceGetAssetsSuspense = < }); /** * Get Asset Events - * Get assets events. + * Get asset events. * @param data The data for the request. * @param data.limit * @param data.offset diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index 7a78230a90eea..32adcace5d1a2 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -172,7 +172,7 @@ export class AssetService { /** * Get Asset Events - * Get assets events. + * Get asset events. * @param data The data for the request. * @param data.limit * @param data.offset diff --git a/tests/api_fastapi/core_api/routes/public/test_assets.py b/tests/api_fastapi/core_api/routes/public/test_assets.py index caf15982999fe..dd34318cf4414 100644 --- a/tests/api_fastapi/core_api/routes/public/test_assets.py +++ b/tests/api_fastapi/core_api/routes/public/test_assets.py @@ -26,6 +26,7 @@ from airflow.utils.state import DagRunState from airflow.utils.types import DagRunType +from tests_common.test_utils.api_connexion_utils import assert_401 from tests_common.test_utils.db import clear_db_assets, clear_db_runs pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode] @@ -305,7 +306,7 @@ def test_should_respect_page_size_limit_default(self, test_client): assert len(response.json()["assets"]) == 100 -class TestGetAssetsEvents(TestAssets): +class TestGetAssetEvents(TestAssets): def test_should_respond_200(self, test_client, session): self.create_assets() self.create_assets_events() @@ -368,6 +369,16 @@ def test_should_respond_200(self, test_client, session): "total_entries": 2, } + @provide_session + def test_filtering(self, test_client, filter_type, filter_value, total_entries, session): + self.create_assets() + self.create_assets_events() + self.create_dag_run() + self.create_asset_dag_run() + response = test_client.get(f"/public/assets/events?{filter_type}={filter_value}") + assert response.status_code == 200 + assert response.json()["total_entries"] == total_entries + def test_order_by_raises_400_for_invalid_attr(self, test_client, session): response = test_client.get("/public/assets/events?order_by=fake") @@ -385,12 +396,6 @@ def test_order_by_raises_400_for_invalid_attr(self, test_client, session): ("source_map_index", "-1", 2), ], ) - @provide_session - def test_filter_events_by_asset_id(self, test_client, filter_type, filter_value, total_entries, session): - self.create_assets() - self.create_assets_events() - self.create_dag_run() - self.create_asset_dag_run() - response = test_client.get(f"/public/assets/events?{filter_type}={filter_value}") - assert response.status_code == 200 - assert response.json()["total_entries"] == total_entries + def test_should_raises_401_unauthenticated(self, test_client, session): + response = test_client.get("/public/assets/events") + assert_401(response) From 1b84e0998ceea89d6a22ee607b919f03bbe4d5e5 Mon Sep 17 00:00:00 2001 From: vatsrahul1001 Date: Wed, 13 Nov 2024 16:55:50 +0530 Subject: [PATCH 16/19] fixing test parametrize --- .../core_api/routes/public/test_assets.py | 25 ++++++++----------- 1 file changed, 10 insertions(+), 15 deletions(-) diff --git a/tests/api_fastapi/core_api/routes/public/test_assets.py b/tests/api_fastapi/core_api/routes/public/test_assets.py index dd34318cf4414..94609fee82206 100644 --- a/tests/api_fastapi/core_api/routes/public/test_assets.py +++ b/tests/api_fastapi/core_api/routes/public/test_assets.py @@ -26,7 +26,6 @@ from airflow.utils.state import DagRunState from airflow.utils.types import DagRunType -from tests_common.test_utils.api_connexion_utils import assert_401 from tests_common.test_utils.db import clear_db_assets, clear_db_runs pytestmark = [pytest.mark.db_test, pytest.mark.skip_if_database_isolation_mode] @@ -369,6 +368,16 @@ def test_should_respond_200(self, test_client, session): "total_entries": 2, } + @pytest.mark.parametrize( + "filter_type, filter_value, total_entries", + [ + ("asset_id", "2", 1), + ("source_dag_id", "source_dag_id", 2), + ("source_task_id", "source_task_id", 2), + ("source_run_id", "source_run_id_1", 1), + ("source_map_index", "-1", 2), + ], + ) @provide_session def test_filtering(self, test_client, filter_type, filter_value, total_entries, session): self.create_assets() @@ -385,17 +394,3 @@ def test_order_by_raises_400_for_invalid_attr(self, test_client, session): assert response.status_code == 400 msg = "Ordering with 'fake' is disallowed or the attribute does not exist on the model" assert response.json()["detail"] == msg - - @pytest.mark.parametrize( - "filter_type, filter_value, total_entries", - [ - ("asset_id", "2", 1), - ("source_dag_id", "source_dag_id", 2), - ("source_task_id", "source_task_id", 2), - ("source_run_id", "source_run_id_1", 1), - ("source_map_index", "-1", 2), - ], - ) - def test_should_raises_401_unauthenticated(self, test_client, session): - response = test_client.get("/public/assets/events") - assert_401(response) From 7eba6d471a1b5367ad7b311cbc93ddcfab171474 Mon Sep 17 00:00:00 2001 From: vatsrahul1001 Date: Wed, 13 Nov 2024 22:15:02 +0530 Subject: [PATCH 17/19] address review comments --- airflow/api_fastapi/core_api/datamodels/assets.py | 13 +++---------- .../api_fastapi/core_api/openapi/v1-generated.yaml | 8 ++++---- .../api_fastapi/core_api/routes/public/assets.py | 5 +++-- airflow/ui/openapi-gen/requests/schemas.gen.ts | 8 ++++---- airflow/ui/openapi-gen/requests/types.gen.ts | 4 ++-- 5 files changed, 16 insertions(+), 22 deletions(-) diff --git a/airflow/api_fastapi/core_api/datamodels/assets.py b/airflow/api_fastapi/core_api/datamodels/assets.py index 0beeead27bc9a..12ecb84569350 100644 --- a/airflow/api_fastapi/core_api/datamodels/assets.py +++ b/airflow/api_fastapi/core_api/datamodels/assets.py @@ -19,7 +19,7 @@ from datetime import datetime -from pydantic import BaseModel, Field, model_validator +from pydantic import BaseModel, Field class DagScheduleAssetReference(BaseModel): @@ -84,7 +84,7 @@ class AssetEventResponse(BaseModel): id: int asset_id: int - asset_uri: str + uri: str extra: dict | None = None source_task_id: str | None = None source_dag_id: str | None = None @@ -93,16 +93,9 @@ class AssetEventResponse(BaseModel): created_dagruns: list[DagRunAssetReference] timestamp: datetime - @model_validator(mode="before") - def rename_uri_to_asset_uri(cls, values): - """Rename 'uri' to 'asset_uri' during serialization to match legacy response.""" - if hasattr(values, "uri") and values.uri: - values.asset_uri = values.uri - return values - class AssetEventCollectionResponse(BaseModel): - """Asset collection response.""" + """Asset event collection response.""" asset_events: list[AssetEventResponse] total_entries: int diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index f98ca622b5e73..3416ed8786348 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -3577,7 +3577,7 @@ components: - asset_events - total_entries title: AssetEventCollectionResponse - description: Asset collection response. + description: Asset event collection response. AssetEventResponse: properties: id: @@ -3586,9 +3586,9 @@ components: asset_id: type: integer title: Asset Id - asset_uri: + uri: type: string - title: Asset Uri + title: Uri extra: anyOf: - type: object @@ -3625,7 +3625,7 @@ components: required: - id - asset_id - - asset_uri + - uri - source_map_index - created_dagruns - timestamp diff --git a/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow/api_fastapi/core_api/routes/public/assets.py index dcee33000e95b..00b649784376e 100644 --- a/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/airflow/api_fastapi/core_api/routes/public/assets.py @@ -21,7 +21,7 @@ from fastapi import Depends from sqlalchemy import select -from sqlalchemy.orm import Session +from sqlalchemy.orm import Session, subqueryload from airflow.api_fastapi.common.db.common import get_session, paginated_select from airflow.api_fastapi.common.parameters import ( @@ -85,7 +85,7 @@ def get_assets( "/events", responses=create_openapi_http_exception_doc([401, 403, 404]), ) -async def get_asset_events( +def get_asset_events( limit: QueryLimit, offset: QueryOffset, order_by: Annotated[ @@ -120,6 +120,7 @@ async def get_asset_events( session=session, ) + assets_event_select = assets_event_select.options(subqueryload(AssetEvent.created_dagruns)) assets_events = session.scalars(assets_event_select).all() return AssetEventCollectionResponse( diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index 89a48825ec9bd..3cefea463a2a4 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -144,7 +144,7 @@ export const $AssetEventCollectionResponse = { type: "object", required: ["asset_events", "total_entries"], title: "AssetEventCollectionResponse", - description: "Asset collection response.", + description: "Asset event collection response.", } as const; export const $AssetEventResponse = { @@ -157,9 +157,9 @@ export const $AssetEventResponse = { type: "integer", title: "Asset Id", }, - asset_uri: { + uri: { type: "string", - title: "Asset Uri", + title: "Uri", }, extra: { anyOf: [ @@ -226,7 +226,7 @@ export const $AssetEventResponse = { required: [ "id", "asset_id", - "asset_uri", + "uri", "source_map_index", "created_dagruns", "timestamp", diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index db709a16339ad..97b56016fc59f 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -38,7 +38,7 @@ export type AssetCollectionResponse = { }; /** - * Asset collection response. + * Asset event collection response. */ export type AssetEventCollectionResponse = { asset_events: Array; @@ -51,7 +51,7 @@ export type AssetEventCollectionResponse = { export type AssetEventResponse = { id: number; asset_id: number; - asset_uri: string; + uri: string; extra?: { [key: string]: unknown; } | null; From a7630333bed86e5cf1f9a2365e56d9315b633092 Mon Sep 17 00:00:00 2001 From: vatsrahul1001 Date: Wed, 13 Nov 2024 23:16:48 +0530 Subject: [PATCH 18/19] address review comments --- .../api_fastapi/core_api/datamodels/assets.py | 2 +- .../core_api/openapi/v1-generated.yaml | 2 +- .../ui/openapi-gen/requests/schemas.gen.ts | 3 +- airflow/ui/openapi-gen/requests/types.gen.ts | 2 +- .../core_api/routes/public/test_assets.py | 43 ++++++++++++++----- 5 files changed, 37 insertions(+), 15 deletions(-) diff --git a/airflow/api_fastapi/core_api/datamodels/assets.py b/airflow/api_fastapi/core_api/datamodels/assets.py index 12ecb84569350..85e41ff7b5698 100644 --- a/airflow/api_fastapi/core_api/datamodels/assets.py +++ b/airflow/api_fastapi/core_api/datamodels/assets.py @@ -67,7 +67,7 @@ class AssetCollectionResponse(BaseModel): class DagRunAssetReference(BaseModel): - """Serializable version of the DagRunAssetReference ORM SqlAlchemyModel.""" + """DAGRun serializer for asset responses.""" run_id: str dag_id: str diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 3416ed8786348..11cbee290acd4 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -4730,7 +4730,7 @@ components: - data_interval_start - data_interval_end title: DagRunAssetReference - description: Serializable version of the DagRunAssetReference ORM SqlAlchemyModel. + description: DAGRun serializer for asset responses. DagRunState: type: string enum: diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index 3cefea463a2a4..ef710ae05a3ef 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -1942,8 +1942,7 @@ export const $DagRunAssetReference = { "data_interval_end", ], title: "DagRunAssetReference", - description: - "Serializable version of the DagRunAssetReference ORM SqlAlchemyModel.", + description: "DAGRun serializer for asset responses.", } as const; export const $DagRunState = { diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 97b56016fc59f..e3e5e8cd322d0 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -411,7 +411,7 @@ export type DagProcessorInfoSchema = { }; /** - * Serializable version of the DagRunAssetReference ORM SqlAlchemyModel. + * DAGRun serializer for asset responses. */ export type DagRunAssetReference = { run_id: string; diff --git a/tests/api_fastapi/core_api/routes/public/test_assets.py b/tests/api_fastapi/core_api/routes/public/test_assets.py index 94609fee82206..b56c2b8b34498 100644 --- a/tests/api_fastapi/core_api/routes/public/test_assets.py +++ b/tests/api_fastapi/core_api/routes/public/test_assets.py @@ -321,7 +321,7 @@ def test_should_respond_200(self, test_client, session): { "id": 1, "asset_id": 1, - "asset_uri": "s3://bucket/key/1", + "uri": "s3://bucket/key/1", "extra": {"foo": "bar"}, "source_task_id": "source_task_id", "source_dag_id": "source_dag_id", @@ -344,7 +344,7 @@ def test_should_respond_200(self, test_client, session): { "id": 2, "asset_id": 2, - "asset_uri": "s3://bucket/key/2", + "uri": "s3://bucket/key/2", "extra": {"foo": "bar"}, "source_task_id": "source_task_id", "source_dag_id": "source_dag_id", @@ -369,22 +369,22 @@ def test_should_respond_200(self, test_client, session): } @pytest.mark.parametrize( - "filter_type, filter_value, total_entries", + "params, total_entries", [ - ("asset_id", "2", 1), - ("source_dag_id", "source_dag_id", 2), - ("source_task_id", "source_task_id", 2), - ("source_run_id", "source_run_id_1", 1), - ("source_map_index", "-1", 2), + ({"asset_id": "2"}, 1), + ({"source_dag_id": "source_dag_id"}, 2), + ({"source_task_id": "source_task_id"}, 2), + ({"source_run_id": "source_run_id_1"}, 1), + ({"source_map_index": "-1"}, 2), ], ) @provide_session - def test_filtering(self, test_client, filter_type, filter_value, total_entries, session): + def test_filtering(self, test_client, params, total_entries, session): self.create_assets() self.create_assets_events() self.create_dag_run() self.create_asset_dag_run() - response = test_client.get(f"/public/assets/events?{filter_type}={filter_value}") + response = test_client.get("/public/assets/events", params=params) assert response.status_code == 200 assert response.json()["total_entries"] == total_entries @@ -394,3 +394,26 @@ def test_order_by_raises_400_for_invalid_attr(self, test_client, session): assert response.status_code == 400 msg = "Ordering with 'fake' is disallowed or the attribute does not exist on the model" assert response.json()["detail"] == msg + + @pytest.mark.parametrize( + "params, expected_asset_uris", + [ + # Limit test data + ({"limit": "1"}, ["s3://bucket/key/1"]), + ({"limit": "100"}, [f"s3://bucket/key/{i}" for i in range(1, 101)]), + # Offset test data + ({"offset": "1"}, [f"s3://bucket/key/{i}" for i in range(2, 102)]), + ({"offset": "3"}, [f"s3://bucket/key/{i}" for i in range(4, 104)]), + ], + ) + def test_limit_and_offset(self, test_client, params, expected_asset_uris): + self.create_assets(num=110) + self.create_assets_events(num=110) + self.create_dag_run(num=110) + self.create_asset_dag_run(num=110) + + response = test_client.get("/public/assets/events", params=params) + + assert response.status_code == 200 + asset_uris = [asset["uri"] for asset in response.json()["asset_events"]] + assert asset_uris == expected_asset_uris From 18124fa05623228e9348d4d404707819e242aa8c Mon Sep 17 00:00:00 2001 From: vatsrahul1001 Date: Thu, 14 Nov 2024 14:11:03 +0530 Subject: [PATCH 19/19] removing http 401 and 403 as its now added in root router in #43932 --- airflow/api_fastapi/core_api/routes/public/assets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow/api_fastapi/core_api/routes/public/assets.py b/airflow/api_fastapi/core_api/routes/public/assets.py index 731b67821da22..67218c471615f 100644 --- a/airflow/api_fastapi/core_api/routes/public/assets.py +++ b/airflow/api_fastapi/core_api/routes/public/assets.py @@ -86,7 +86,7 @@ def get_assets( @assets_router.get( "/events", - responses=create_openapi_http_exception_doc([401, 403, 404]), + responses=create_openapi_http_exception_doc([404]), ) def get_asset_events( limit: QueryLimit,