kaxil commented on code in PR #72909: URL: https://github.com/apache/airflow/pull/72909#discussion_r4007966739
########## airflow-core/src/airflow/api_fastapi/core_api/routes/public/dag_bundles.py: ########## @@ -0,0 +1,214 @@ +# 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 TYPE_CHECKING, Annotated + +from fastapi import Depends +from sqlalchemy import and_, func, select + +from airflow.api_fastapi.auth.managers.models.resource_details import AccessView, DagAccessEntity +from airflow.api_fastapi.common.db.common import SessionDep, paginated_select +from airflow.api_fastapi.common.parameters import QueryLimit, QueryOffset, SortParam +from airflow.api_fastapi.common.router import AirflowRouter +from airflow.api_fastapi.core_api.datamodels.dag_bundles import ( + DagBundleCollectionResponse, + DagBundleResponse, +) +from airflow.api_fastapi.core_api.security import ( + AuthManagerDep, + GetUserDep, + ReadableDagBundlesFilterDep, + requires_access_dag, +) +from airflow.configuration import conf +from airflow.models import DagModel +from airflow.models.dagbundle import DagBundleModel +from airflow.models.errors import ParseImportError + +if TYPE_CHECKING: + from collections.abc import Sequence + + from sqlalchemy.orm import Session + + from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager + from airflow.api_fastapi.auth.managers.models.base_user import BaseUser + +dag_bundles_router = AirflowRouter(tags=["Dag Bundle"], prefix="/dagBundles") + + +def _import_error_counts( + *, + bundle_names: Sequence[str], + readable_dag_ids: set[str], + auth_manager: BaseAuthManager, + user: BaseUser, + session: Session, +) -> dict[str, int] | None: + """ + Count the import errors per bundle that this caller is allowed to know about. + + Reproduces the two-part authorization of ``GET /importErrors`` rather than counting every row + for the bundle: an error in a file the caller can read no Dag in stays hidden, and an error in + a file with no registered Dag needs the admin-by-default ``IMPORT_ERRORS_ALL``, since the + file's existence would otherwise leak. Returns ``None`` when the caller may not read import + errors at all. + """ + if not auth_manager.authorize_view(access_view=AccessView.IMPORT_ERRORS, user=user): + return None + if not bundle_names: + return {} + + # Files -- keyed ``(relative_fileloc, bundle_name)`` -- in which the caller can read a Dag. + readable_files = ( + select(DagModel.relative_fileloc, DagModel.bundle_name) + .where( + DagModel.dag_id.in_(readable_dag_ids), + DagModel.bundle_name.in_(bundle_names), + ) + .distinct() + .subquery() + ) + counts: dict[str, int] = { + bundle_name: count + for bundle_name, count in session.execute( + select(ParseImportError.bundle_name, func.count()) + .join( + readable_files, + and_( + ParseImportError.filename == readable_files.c.relative_fileloc, + ParseImportError.bundle_name == readable_files.c.bundle_name, + ), + ) + .group_by(ParseImportError.bundle_name) + ).all() + if bundle_name is not None + } + + # Errors for files that never registered a Dag, added only where the caller holds the + # admin-by-default view for that bundle's team. Narrowed to this page's bundles, unlike the + # equivalent in ``import_error.py``: this endpoint polls, so an unbounded ``SELECT DISTINCT`` + # over ``dag`` every few seconds would be a real cost on a large deployment. + files_with_any_dags = ( + select(DagModel.relative_fileloc, DagModel.bundle_name) + .where(DagModel.bundle_name.in_(bundle_names)) Review Comment: Agreed, and it is a bit worse than the two lookups you name. `DagModel.bundle_name` is a plain `ForeignKey` with no index, so it is unindexed on Postgres, though MySQL auto-creates one for FK columns. `import_error` has no index on `bundle_name` or `filename` at all. The larger cost is upstream of the query you are pointing at. `get_authorized_dag_ids` materialises every `dag_id` on every request, and that set is then embedded as an `IN` list twice more, once in the bundle filter subquery and once in `readable_files`. On a large deployment one idle browser tab drives a full `dag` join plus two large `IN` clauses every ten seconds. `GET /dags` does the same work, but it is not on a self-refreshing timer, so this is the first endpoint where the cost becomes continuous background load rather than per-navigation. Follow-up sounds right to me: add the indexes, and skip computing unregistered-error counts for bundles whose errors the caller cannot read. Flag it if you would rather they block this PR. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
