QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-07-26 16:42:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

perf(dashboard): Batch RLS filter lookups for dashboard digest computation (#37941)

Browse Source
This commit is contained in:
Kamil Gabryjelski
2026-02-16 21:35:55 +01:00
committed by GitHub
parent 7b21979fa3
commit 3e3c9686de
4 changed files with 390 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

123
superset/security/manager.py
View File

@@ -75,6 +75,7 @@ from superset.utils.core import (
DatasourceName,
DatasourceType,
get_user_id,
get_username,
RowLevelSecurityFilterType,
)
from superset.utils.filters import get_dataset_access_filters
@@ -110,6 +111,16 @@ class DatabaseCatalogSchema(NamedTuple):
schema: str
class _RLSFilterRow(NamedTuple):
id: int
group_key: Optional[str]
clause: str
_RLSCacheKey = tuple[str, int | str]
_RLSCache = dict[_RLSCacheKey, list[SqlaQuery]]
class SupersetRoleApi(RoleApi):
"""
Overriding the RoleApi to be able to delete roles with permissions
@@ -2701,6 +2712,15 @@ class SupersetSecurityManager( # pylint: disable=too-many-public-methods
if not (hasattr(g, "user") and g.user is not None):
return []
# Check request-scoped cache. Username is included in the key to stay
# safe if override_user() is called with different users in one request.
cache: _RLSCache = getattr(g, "_rls_filter_cache", {})
username = get_username()
if username is not None:
cache_key: _RLSCacheKey = (username, table.id)
if cache_key in cache:
return cache[cache_key]
# pylint: disable=import-outside-toplevel
from superset.connectors.sqla.models import (
RLSFilterRoles,
@@ -2750,7 +2770,108 @@ class SupersetSecurityManager( # pylint: disable=too-many-public-methods
)
)
)
return query.all()
result = query.all()
# Store in request-scoped cache
if username is not None:
if not hasattr(g, "_rls_filter_cache"):
g._rls_filter_cache = {}
g._rls_filter_cache[(username, table.id)] = result
return result
def prefetch_rls_filters(self, table_ids: list[int | str]) -> None:
"""
Batch-fetches RLS filters for multiple tables in a single query and
populates the request-scoped cache used by get_rls_filters().
:param table_ids: List of table IDs to prefetch filters for
"""
if not (hasattr(g, "user") and g.user is not None):
return
username = get_username()
if username is None:
return
if not hasattr(g, "_rls_filter_cache"):
g._rls_filter_cache = {}
# Filter out already-cached table_ids
uncached_ids = [
tid for tid in table_ids if (username, tid) not in g._rls_filter_cache
]
if not uncached_ids:
return
# pylint: disable=import-outside-toplevel
from superset.connectors.sqla.models import (
RLSFilterRoles,
RLSFilterTables,
RowLevelSecurityFilter,
)
user_roles = [role.id for role in self.get_user_roles(g.user)]
regular_filter_roles = (
self.session.query(RLSFilterRoles.c.rls_filter_id)
.join(RowLevelSecurityFilter)
.filter(
RowLevelSecurityFilter.filter_type == RowLevelSecurityFilterType.REGULAR
)
.filter(RLSFilterRoles.c.role_id.in_(user_roles))
)
base_filter_roles = (
self.session.query(RLSFilterRoles.c.rls_filter_id)
.join(RowLevelSecurityFilter)
.filter(
RowLevelSecurityFilter.filter_type == RowLevelSecurityFilterType.BASE
)
.filter(RLSFilterRoles.c.role_id.in_(user_roles))
)
# Batch query: get (table_id, filter) pairs for all uncached tables
filter_table_pairs = (
self.session.query(
RLSFilterTables.c.table_id,
RowLevelSecurityFilter.id,
RowLevelSecurityFilter.group_key,
RowLevelSecurityFilter.clause,
)
.join(
RowLevelSecurityFilter,
RowLevelSecurityFilter.id == RLSFilterTables.c.rls_filter_id,
)
.filter(RLSFilterTables.c.table_id.in_(uncached_ids))
.filter(
or_(
and_(
RowLevelSecurityFilter.filter_type
== RowLevelSecurityFilterType.REGULAR,
RowLevelSecurityFilter.id.in_(regular_filter_roles),
),
and_(
RowLevelSecurityFilter.filter_type
== RowLevelSecurityFilterType.BASE,
RowLevelSecurityFilter.id.notin_(base_filter_roles),
),
)
)
.all()
)
# Group results by table_id, storing as named tuples so callers
# can access .id, .group_key, .clause (matching get_rls_filters output)
grouped: dict[int | str, list[SqlaQuery]] = defaultdict(list)
for row in filter_table_pairs:
table_id = row[0]
grouped[table_id].append(
_RLSFilterRow(id=row[1], group_key=row[2], clause=row[3])
)
# Populate cache for all uncached table_ids (including those with no filters)
for tid in uncached_ids:
g._rls_filter_cache[(username, tid)] = grouped.get(tid, [])
def get_rls_sorted(
self, table: "BaseDatasource | Explorable"

9
superset/thumbnails/digest.py
View File

@@ -70,6 +70,15 @@ def _adjust_string_with_rls(
if user:
stringified_rls = ""
with override_user(user):
# Prefetch RLS filters for all datasources in a single batch query
table_ids = [
datasource.id
for datasource in datasources
if datasource and getattr(datasource, "is_rls_supported", False)
]
if table_ids:
security_manager.prefetch_rls_filters(table_ids)
for datasource in datasources:
if datasource and getattr(datasource, "is_rls_supported", False):
rls_filters = datasource.get_sqla_row_level_filters()

210
tests/unit_tests/security/manager_test.py
View File

@@ -18,6 +18,7 @@
# pylint: disable=invalid-name, unused-argument, redefined-outer-name
import json # noqa: TID251
from types import SimpleNamespace
import pytest
from flask_appbuilder.security.sqla.models import Role, User
@@ -1223,3 +1224,212 @@ def test_get_rls_filters_uses_table_id_directly(
# If it uses table.id directly (correct behavior), it will complete successfully
result = sm.get_rls_filters(table)
assert isinstance(result, list)
def test_get_rls_filters_returns_cached_result(
mocker: MockerFixture,
app_context: None,
) -> None:
"""
Test that get_rls_filters() returns cached results on subsequent calls
for the same user and table, avoiding redundant DB queries.
"""
sm = SupersetSecurityManager(appbuilder)
mock_user = mocker.MagicMock()
mock_user.id = 1
mock_user.username = "admin"
mock_user.roles = [mocker.MagicMock(id=1)]
mock_g = SimpleNamespace(user=mock_user)
mocker.patch("superset.security.manager.g", new=mock_g)
mocker.patch("superset.security.manager.get_username", return_value="admin")
mocker.patch.object(sm, "get_user_roles", return_value=mock_user.roles)
table = mocker.MagicMock()
table.id = 42
# First call populates the cache
result1 = sm.get_rls_filters(table)
# Verify cache was populated keyed by (username, table_id)
assert ("admin", 42) in mock_g._rls_filter_cache
# Replace session query with something that would fail if called
mocker.patch.object(
sm.session,
"query",
side_effect=AssertionError("DB should not be queried on cache hit"),
)
# Second call should return cached result without querying DB
result2 = sm.get_rls_filters(table)
assert result1 == result2
def test_prefetch_rls_filters_populates_cache(
mocker: MockerFixture,
app_context: None,
) -> None:
"""
Test that prefetch_rls_filters() populates the cache for all provided
table_ids, including empty results for tables with no matching filters.
"""
sm = SupersetSecurityManager(appbuilder)
mock_user = mocker.MagicMock()
mock_user.id = 1
mock_user.username = "admin"
mock_user.roles = [mocker.MagicMock(id=10)]
mock_g = SimpleNamespace(user=mock_user)
mocker.patch("superset.security.manager.g", new=mock_g)
mocker.patch("superset.security.manager.get_username", return_value="admin")
mocker.patch.object(sm, "get_user_roles", return_value=mock_user.roles)
# Mock the batch query to return filters for table 1 but not table 2
mock_query = mocker.MagicMock()
mock_query.join.return_value = mock_query
mock_query.filter.return_value = mock_query
mock_query.all.return_value = [
(1, 100, "group_a", "id > 0"), # table_id=1
(1, 101, None, "active = 1"), # table_id=1
]
mocker.patch.object(sm.session, "query", return_value=mock_query)
sm.prefetch_rls_filters([1, 2])
# Table 1 should have 2 filters with named attribute access
cached = mock_g._rls_filter_cache[("admin", 1)]
assert len(cached) == 2
assert cached[0].id == 100
assert cached[0].group_key == "group_a"
assert cached[0].clause == "id > 0"
assert cached[1].id == 101
assert cached[1].group_key is None
assert cached[1].clause == "active = 1"
# Table 2 should have empty list
assert mock_g._rls_filter_cache[("admin", 2)] == []
def test_prefetch_rls_filters_skips_cached_ids(
mocker: MockerFixture,
app_context: None,
) -> None:
"""
Test that prefetch_rls_filters() skips table_ids already in cache
and returns early when all ids are cached.
"""
sm = SupersetSecurityManager(appbuilder)
mock_user = mocker.MagicMock()
mock_user.id = 1
mock_user.username = "admin"
mock_user.roles = [mocker.MagicMock(id=10)]
mock_g = SimpleNamespace(
user=mock_user,
_rls_filter_cache={("admin", 1): [(100, "group_a", "id > 0")]},
)
mocker.patch("superset.security.manager.g", new=mock_g)
mocker.patch("superset.security.manager.get_username", return_value="admin")
mocker.patch.object(sm, "get_user_roles", return_value=mock_user.roles)
# If it queries the DB, this will fail
mocker.patch.object(
sm.session,
"query",
side_effect=AssertionError("DB should not be queried for cached ids"),
)
# All ids already cached -> should return immediately
sm.prefetch_rls_filters([1])
def test_prefetch_rls_filters_no_user(
mocker: MockerFixture,
app_context: None,
) -> None:
"""
Test that prefetch_rls_filters() returns early when no user is present.
"""
sm = SupersetSecurityManager(appbuilder)
mocker.patch("superset.security.manager.g", new=SimpleNamespace())
# Should not attempt any DB queries
mocker.patch.object(
sm.session,
"query",
side_effect=AssertionError("DB should not be queried without a user"),
)
sm.prefetch_rls_filters([1, 2])
def test_get_rls_filters_cache_works_for_guest_user(
mocker: MockerFixture,
app_context: None,
) -> None:
"""
Test that get_rls_filters() caches results for guest users
using the same (username, table_id) cache key as regular users.
"""
sm = SupersetSecurityManager(appbuilder)
mock_guest = mocker.MagicMock()
mock_guest.username = "guest_user"
mock_guest.roles = [mocker.MagicMock(id=99)]
mock_g = SimpleNamespace(user=mock_guest)
mocker.patch("superset.security.manager.g", new=mock_g)
mocker.patch("superset.security.manager.get_username", return_value="guest_user")
mocker.patch.object(sm, "get_user_roles", return_value=mock_guest.roles)
table = mocker.MagicMock()
table.id = 42
# First call runs the query
result1 = sm.get_rls_filters(table)
# Verify cache was populated with (username, table_id) key
assert ("guest_user", 42) in mock_g._rls_filter_cache
# Replace session query to detect if it's called again
mocker.patch.object(
sm.session,
"query",
side_effect=AssertionError("DB should not be queried on cache hit"),
)
# Second call should use cache
result2 = sm.get_rls_filters(table)
assert result1 == result2
def test_prefetch_rls_filters_works_for_guest_user(
mocker: MockerFixture,
app_context: None,
) -> None:
"""
Test that prefetch_rls_filters() works for guest users using the
same (username, table_id) cache key as regular users.
"""
sm = SupersetSecurityManager(appbuilder)
mock_guest = mocker.MagicMock()
mock_guest.username = "guest_user"
mock_guest.roles = [mocker.MagicMock(id=99)]
mock_g = SimpleNamespace(user=mock_guest)
mocker.patch("superset.security.manager.g", new=mock_g)
mocker.patch("superset.security.manager.get_username", return_value="guest_user")
mocker.patch.object(sm, "get_user_roles", return_value=mock_guest.roles)
# Mock the batch query returning no filters
mock_query = mocker.MagicMock()
mock_query.join.return_value = mock_query
mock_query.filter.return_value = mock_query
mock_query.all.return_value = []
mocker.patch.object(sm.session, "query", return_value=mock_query)
sm.prefetch_rls_filters([10, 20])
# Cache should be populated with (username, table_id) keys and empty lists
assert mock_g._rls_filter_cache[("guest_user", 10)] == []
assert mock_g._rls_filter_cache[("guest_user", 20)] == []

49
tests/unit_tests/thumbnails/test_digest.py
View File

@@ -433,3 +433,52 @@ def test_chart_digest(
)
with cm:
assert get_chart_digest(chart=chart) == expected_result
def test_dashboard_digest_prefetches_rls_filters(
app_context: None,
) -> None:
"""
Test that _adjust_string_with_rls calls prefetch_rls_filters with
table IDs from RLS-supporting datasources before iterating.
"""
from superset import security_manager
from superset.models.dashboard import Dashboard
from superset.models.slice import Slice
from superset.thumbnails.digest import get_dashboard_digest
kwargs = {**_DEFAULT_DASHBOARD_KWARGS}
slices = [Slice(**slice_kwargs) for slice_kwargs in kwargs.pop("slices")]
dashboard = Dashboard(**kwargs, slices=slices)
datasources = []
for ds_id, rls_supported in [(10, True), (20, True), (30, False)]:
ds = MagicMock(spec=BaseDatasource)
ds.id = ds_id
ds.is_rls_supported = rls_supported
ds.get_sqla_row_level_filters = MagicMock(return_value=[])
datasources.append(ds)
user = User(id=1, username="1")
with (
patch.dict(
current_app.config,
{
"THUMBNAIL_EXECUTORS": [ExecutorType.CURRENT_USER],
"THUMBNAIL_DASHBOARD_DIGEST_FUNC": None,
},
),
patch.object(
type(dashboard),
"datasources",
new_callable=PropertyMock,
return_value=datasources,
),
patch.object(security_manager, "find_user", return_value=user),
patch.object(security_manager, "prefetch_rls_filters") as mock_prefetch,
override_user(user),
):
get_dashboard_digest(dashboard=dashboard)
# Should be called with only the RLS-supporting datasource IDs
mock_prefetch.assert_called_once_with([10, 20])