This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 73a4ffa3a53 Avoid a full asset table scan per request in simple auth
manager (#72934)
73a4ffa3a53 is described below
commit 73a4ffa3a530da280600a7f7c2d066fac1c9044b
Author: Wei Lee <[email protected]>
AuthorDate: Wed Oct 7 22:28:06 2026 +0900
Avoid a full asset table scan per request in simple auth manager (#72934)
Scoping asset API responses to what a user may read (#72682) gave
BaseAuthManager a default that loads the id, name and uri of every asset
and then asks is_authorized_asset about each row. Simple auth manager
grants asset access by role alone and never inspects the asset it is
asked about, so that loop re-derives a single constant answer once per
row, along with the name and uri it only loads to build the details it
then ignores.
The auth managers whose is_authorized_asset makes a remote call were
given batched overrides at the time; simple auth manager, which is the
default, was left on the generic path.
---
.../auth/managers/simple/simple_auth_manager.py | 30 +++++++++
.../managers/simple/test_simple_auth_manager.py | 74 ++++++++++++++++++++++
.../core_api/routes/public/test_assets.py | 8 ++-
.../api_fastapi/core_api/routes/ui/test_assets.py | 12 ++--
4 files changed, 117 insertions(+), 7 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/auth/managers/simple/simple_auth_manager.py
b/airflow-core/src/airflow/api_fastapi/auth/managers/simple/simple_auth_manager.py
index 8f15e734d2c..a030c9e034d 100644
---
a/airflow-core/src/airflow/api_fastapi/auth/managers/simple/simple_auth_manager.py
+++
b/airflow-core/src/airflow/api_fastapi/auth/managers/simple/simple_auth_manager.py
@@ -33,6 +33,7 @@ from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
+from sqlalchemy import select
from termcolor import colored
from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
@@ -41,8 +42,11 @@ from
airflow.api_fastapi.auth.managers.models.resource_details import AccessView
from airflow.api_fastapi.auth.managers.simple.user import SimpleAuthManagerUser
from airflow.api_fastapi.common.types import MenuItem
from airflow.configuration import AIRFLOW_HOME, conf
+from airflow.models.asset import AssetModel
+from airflow.utils.session import NEW_SESSION, provide_session
if TYPE_CHECKING:
+ from sqlalchemy.orm import Session
from starlette.middleware import _MiddlewareFactory
from airflow.api_fastapi.auth.managers.base_auth_manager import
ResourceMethod
@@ -309,6 +313,32 @@ class
SimpleAuthManager(BaseAuthManager[SimpleAuthManagerUser]):
user=user,
)
+ @provide_session
+ def get_authorized_assets(
+ self,
+ *,
+ user: SimpleAuthManagerUser,
+ method: ResourceMethod = "GET",
+ session: Session = NEW_SESSION,
+ ) -> set[int]:
+ """
+ Get the ids of the assets the user has access to.
+
+ Simple auth manager authorizes assets at the role level:
``is_authorized_asset`` ignores the
+ asset details, so one check decides the whole listing. The default
per-asset loop would
+ re-evaluate that same decision once per row.
+
+ This relies on ``is_authorized_asset`` not reading ``details``. If it
starts to, or a subclass
+ overrides ``filter_authorized_assets``, remove or adjust this override.
+
+ :param user: the user
+ :param method: the method to filter on
+ :param session: the session
+ """
+ if not self.is_authorized_asset(method=method, user=user):
+ return set()
+ return set(session.scalars(select(AssetModel.id)))
+
def is_authorized_pool(
self,
*,
diff --git
a/airflow-core/tests/unit/api_fastapi/auth/managers/simple/test_simple_auth_manager.py
b/airflow-core/tests/unit/api_fastapi/auth/managers/simple/test_simple_auth_manager.py
index 5cdb5248a12..1cda1cf9061 100644
---
a/airflow-core/tests/unit/api_fastapi/auth/managers/simple/test_simple_auth_manager.py
+++
b/airflow-core/tests/unit/api_fastapi/auth/managers/simple/test_simple_auth_manager.py
@@ -24,6 +24,7 @@ from urllib.parse import urlencode
import pytest
from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
+from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager
from airflow.api_fastapi.auth.managers.models.resource_details import (
AccessView,
ConnectionDetails,
@@ -35,8 +36,20 @@ from
airflow.api_fastapi.auth.managers.models.resource_details import (
from airflow.api_fastapi.auth.managers.simple.simple_auth_manager import
SimpleAuthManager
from airflow.api_fastapi.auth.managers.simple.user import SimpleAuthManagerUser
from airflow.api_fastapi.common.types import MenuItem
+from airflow.models.asset import AssetActive, AssetModel
from tests_common.test_utils.config import conf_vars
+from tests_common.test_utils.db import clear_db_assets
+
+ASSET_AUTHORIZATION_MATRIX = [
+ pytest.param("ADMIN", "GET", True, id="admin-get"),
+ pytest.param(None, "GET", False, id="no-role-get"),
+ pytest.param("VIEWER", "GET", True, id="viewer-get"),
+ pytest.param("VIEWER", "POST", False, id="viewer-post"),
+ pytest.param("USER", "GET", True, id="user-get"),
+ pytest.param("USER", "POST", False, id="user-post"),
+ pytest.param("OP", "POST", True, id="op-post"),
+]
class TestSimpleAuthManager:
@@ -524,3 +537,64 @@ class TestSimpleAuthManager:
password = SimpleAuthManager._generate_password()
assert len(password) == 16
assert set(password).issubset(alphabet)
+
+ @pytest.fixture
+ def clean_assets(self):
+ clear_db_assets()
+ yield
+ clear_db_assets()
+
+ @pytest.mark.db_test
+ @pytest.mark.parametrize(("role", "method", "expect_authorized"),
ASSET_AUTHORIZATION_MATRIX)
+ def test_get_authorized_assets(
+ self, auth_manager, session, clean_assets, role, method,
expect_authorized
+ ):
+ """Three assets in the table; one role-level check, not one per asset,
is what the override does."""
+ assets = [AssetModel(name=f"asset{i}", uri=f"s3://bucket/asset{i}",
group="asset") for i in range(3)]
+ session.add_all(assets)
+ # The third asset has no AssetActive row and must still be returned.
+ session.add_all(AssetActive.for_asset(asset) for asset in assets[:2])
+ session.commit()
+ expected_ids = {asset.id for asset in assets} if expect_authorized
else set()
+
+ user = SimpleAuthManagerUser(username="test", role=role)
+ with mock.patch.object(
+ auth_manager,
+ "is_authorized_asset",
+ wraps=auth_manager.is_authorized_asset,
+ ) as mock_is_authorized_asset:
+ result = auth_manager.get_authorized_assets(user=user,
method=method, session=session)
+
+ assert result == expected_ids
+ mock_is_authorized_asset.assert_called_once_with(method=method,
user=user)
+
+ @pytest.mark.db_test
+ def test_get_authorized_assets_empty_table(self, auth_manager, session,
clean_assets):
+ user = SimpleAuthManagerUser(username="test", role="ADMIN")
+ with mock.patch.object(
+ auth_manager,
+ "is_authorized_asset",
+ wraps=auth_manager.is_authorized_asset,
+ ) as mock_is_authorized_asset:
+ result = auth_manager.get_authorized_assets(user=user,
session=session)
+
+ assert result == set()
+ mock_is_authorized_asset.assert_called_once_with(method="GET",
user=user)
+
+ @pytest.mark.db_test
+ @pytest.mark.parametrize(("role", "method", "expect_authorized"),
ASSET_AUTHORIZATION_MATRIX)
+ def test_get_authorized_assets_matches_base_implementation(
+ self, auth_manager, session, clean_assets, role, method,
expect_authorized
+ ):
+ assets = [AssetModel(name=f"asset{i}", uri=f"s3://bucket/asset{i}",
group="asset") for i in range(3)]
+ session.add_all(assets)
+ session.add_all(AssetActive.for_asset(asset) for asset in assets[:2])
+ session.commit()
+
+ user = SimpleAuthManagerUser(username="test", role=role)
+ result = auth_manager.get_authorized_assets(user=user, method=method,
session=session)
+
+ assert result == BaseAuthManager.get_authorized_assets(
+ auth_manager, user=user, method=method, session=session
+ )
+ assert bool(result) is expect_authorized
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py
index a51a0ef1ae5..fb84a18d9c7 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py
@@ -418,7 +418,9 @@ class TestGetAssets(TestAssets):
"total_entries": 2,
}
-
@mock.patch("airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_assets")
+ @mock.patch(
+
"airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager.get_authorized_assets"
+ )
def test_should_return_only_assets_the_caller_may_read(
self, mock_get_authorized_assets, test_client, session
):
@@ -1018,7 +1020,9 @@ class TestGetAssetEventsPerDagScoping(TestAssets):
# The count must be scoped too, so the existence of hidden events does
not leak.
assert body["total_entries"] == len(expected_ids)
-
@mock.patch("airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_assets")
+ @mock.patch(
+
"airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager.get_authorized_assets"
+ )
def test_endpoint_returns_only_events_of_assets_the_caller_may_read(
self, mock_get_authorized_assets, test_client, session
):
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py
index e26dc173651..87f83aac375 100644
--- a/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py
+++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/ui/test_assets.py
@@ -124,7 +124,7 @@ class TestNextRunAssets:
}
@mock.patch(
-
"airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_assets",
+
"airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager.get_authorized_assets",
autospec=True,
)
def test_asset_expression_hides_assets_the_caller_may_not_read(
@@ -165,7 +165,7 @@ class TestNextRunAssets:
assert "s3://bucket/hidden" not in redacted
@mock.patch(
-
"airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_assets",
+
"airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager.get_authorized_assets",
autospec=True,
)
def test_should_return_only_assets_the_caller_may_read(
@@ -191,7 +191,7 @@ class TestNextRunAssets:
assert response.json()["scheduling_asset_count"] == 2
@mock.patch(
-
"airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_assets",
+
"airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager.get_authorized_assets",
autospec=True,
)
def test_scheduling_asset_count_is_unaffected_when_no_asset_is_readable(
@@ -216,7 +216,7 @@ class TestNextRunAssets:
assert response.json()["scheduling_asset_count"] == 2
@mock.patch(
-
"airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_assets",
+
"airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager.get_authorized_assets",
autospec=True,
)
def
test_partitioned_dag_should_hide_keys_of_assets_the_caller_may_not_read(
@@ -686,7 +686,9 @@ class TestGetAssetsUi:
assert body["total_entries"] == 1
assert body["assets"][0]["name"] == "ui_asset"
-
@mock.patch("airflow.api_fastapi.auth.managers.base_auth_manager.BaseAuthManager.get_authorized_assets")
+ @mock.patch(
+
"airflow.api_fastapi.auth.managers.simple.simple_auth_manager.SimpleAuthManager.get_authorized_assets"
+ )
def test_should_return_only_assets_the_caller_may_read(
self, mock_get_authorized_assets, test_client, session
):
