QT/superset2
1
0
Fork 0
mirror of https://github.com/apache/superset.git synced 2026-04-19 08:04:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(Digest): Add RLS at digest generation for Charts and Dashboards (#30336)

Browse Source 
Co-authored-by: Ville Brofeldt <33317356+villebro@users.noreply.github.com>
This commit is contained in:
Geido
2024-09-24 15:39:11 +02:00
committed by GitHub
parent cc1bb69671
commit de3af85ee1
5 changed files with 203 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -18,14 +18,15 @@ from __future__ import annotations
from contextlib import nullcontext
from typing import Any, TYPE_CHECKING
from unittest.mock import patch
from unittest.mock import MagicMock, patch, PropertyMock
import pytest
from flask_appbuilder.security.sqla.models import User
from superset.connectors.sqla.models import BaseDatasource, SqlaTable
from superset.tasks.exceptions import ExecutorNotFoundError
from superset.tasks.types import ExecutorType
from superset.utils.core import override_user
from superset.utils.core import DatasourceType, override_user
if TYPE_CHECKING:
from superset.models.dashboard import Dashboard
@@ -62,14 +63,28 @@ def CUSTOM_CHART_FUNC(
return f"{chart.id}.{executor_type.value}.{executor}"
def prepare_datasource_mock(
datasource_conf: dict[str, Any], spec: type[BaseDatasource | SqlaTable]
) -> BaseDatasource | SqlaTable:
datasource = MagicMock(spec=spec)
datasource.id = 1
datasource.type = DatasourceType.TABLE
datasource.is_rls_supported = datasource_conf.get("is_rls_supported", False)
datasource.get_sqla_row_level_filters = datasource_conf.get(
"get_sqla_row_level_filters", MagicMock(return_value=[])
)
return datasource
@pytest.mark.parametrize(
"dashboard_overrides,execute_as,has_current_user,use_custom_digest,expected_result",
"dashboard_overrides,execute_as,has_current_user,use_custom_digest,rls_datasources,expected_result",
[
(
None,
[ExecutorType.SELENIUM],
False,
False,
[],
"71452fee8ffbd8d340193d611bcd4559",
),
(
@@ -77,6 +92,7 @@ def CUSTOM_CHART_FUNC(
[ExecutorType.CURRENT_USER],
True,
False,
[],
"209dc060ac19271b8708731e3b8280f5",
),
(
@@ -86,6 +102,7 @@ def CUSTOM_CHART_FUNC(
[ExecutorType.CURRENT_USER],
True,
False,
[],
"209dc060ac19271b8708731e3b8280f5",
),
(
@@ -95,6 +112,7 @@ def CUSTOM_CHART_FUNC(
[ExecutorType.CURRENT_USER],
True,
False,
[],
"06a4144466dbd5ffad0c3c2225e96296",
),
(
@@ -104,6 +122,7 @@ def CUSTOM_CHART_FUNC(
[ExecutorType.CURRENT_USER],
True,
False,
[],
"a823ece9563895ccb14f3d9095e84f7a",
),
(
@@ -113,6 +132,7 @@ def CUSTOM_CHART_FUNC(
[ExecutorType.CURRENT_USER],
True,
False,
[],
"33c5475f92a904925ab3ef493526e5b5",
),
(
@@ -122,6 +142,7 @@ def CUSTOM_CHART_FUNC(
[ExecutorType.CURRENT_USER],
True,
False,
[],
"cec57345e6402c0d4b3caee5cfaa0a03",
),
(
@@ -131,20 +152,68 @@ def CUSTOM_CHART_FUNC(
[ExecutorType.CURRENT_USER],
True,
False,
[],
"5380dcbe94621a0759b09554404f3d02",
),
(
None,
[ExecutorType.CURRENT_USER],
True,
False,
[
{
"is_rls_supported": True,
"get_sqla_row_level_filters": MagicMock(return_value=["filter1"]),
}
],
"4138959f275c1991466cafcfb190fd72",
),
(
None,
[ExecutorType.CURRENT_USER],
True,
"1.current_user.1",
False,
[
{
"is_rls_supported": True,
"get_sqla_row_level_filters": MagicMock(
return_value=["filter1", "filter2"]
),
},
{
"is_rls_supported": True,
"get_sqla_row_level_filters": MagicMock(
return_value=["filter3", "filter4"]
),
},
],
"80d3bfcc7144bccdba8c718cf49b6420",
),
(
None,
[ExecutorType.CURRENT_USER],
True,
False,
[
{
"is_rls_supported": False,
"get_sqla_row_level_filters": MagicMock(return_value=[]),
},
{
"is_rls_supported": True,
"get_sqla_row_level_filters": MagicMock(
return_value=["filter1", "filter2"]
),
},
],
"e8fc68cd5aba22a5f1acf06164bfc0f4",
),
(
None,
[ExecutorType.CURRENT_USER],
False,
False,
[],
ExecutorNotFoundError(),
),
],
@@ -154,22 +223,32 @@ def test_dashboard_digest(
execute_as: list[ExecutorType],
has_current_user: bool,
use_custom_digest: bool,
rls_datasources: list[dict[str, Any]],
expected_result: str | Exception,
) -> None:
from superset import app
from superset import app, security_manager
from superset.models.dashboard import Dashboard
from superset.models.slice import Slice
from superset.thumbnails.digest import get_dashboard_digest
# Prepare dashboard and slices
kwargs = {
**_DEFAULT_DASHBOARD_KWARGS,
**(dashboard_overrides or {}),
}
slices = [Slice(**slice_kwargs) for slice_kwargs in kwargs.pop("slices")]
dashboard = Dashboard(**kwargs, slices=slices)
# Mock datasources with RLS
datasources = []
for rls_source in rls_datasources:
datasource = prepare_datasource_mock(rls_source, BaseDatasource)
datasources.append(datasource)
user: User | None = None
if has_current_user:
user = User(id=1, username="1")
func = CUSTOM_DASHBOARD_FUNC if use_custom_digest else None
with (
@@ -180,6 +259,13 @@ def test_dashboard_digest(
"THUMBNAIL_DASHBOARD_DIGEST_FUNC": func,
},
),
patch.object(
type(dashboard),
"datasources",
new_callable=PropertyMock,
return_value=datasources,
),
patch.object(security_manager, "find_user", return_value=user),
override_user(user),
):
cm = (
@@ -192,13 +278,14 @@ def test_dashboard_digest(
@pytest.mark.parametrize(
"chart_overrides,execute_as,has_current_user,use_custom_digest,expected_result",
"chart_overrides,execute_as,has_current_user,use_custom_digest,rls_datasource,expected_result",
[
(
None,
[ExecutorType.SELENIUM],
False,
False,
None,
"47d852b5c4df211c115905617bb722c1",
),
(
@@ -206,6 +293,7 @@ def test_dashboard_digest(
[ExecutorType.CURRENT_USER],
True,
False,
None,
"4f8109d3761e766e650af514bb358f10",
),
(
@@ -213,13 +301,50 @@ def test_dashboard_digest(
[ExecutorType.CURRENT_USER],
True,
True,
None,
"2.current_user.1",
),
(
None,
[ExecutorType.CURRENT_USER],
True,
False,
{
"is_rls_supported": True,
"get_sqla_row_level_filters": MagicMock(return_value=["filter1"]),
},
"61e70336c27eb97fb050328a0b050373",
),
(
None,
[ExecutorType.CURRENT_USER],
True,
False,
{
"is_rls_supported": True,
"get_sqla_row_level_filters": MagicMock(
return_value=["filter1", "filter2"]
),
},
"95c7cefde8cb519f005f33bfb33cb196",
),
(
None,
[ExecutorType.CURRENT_USER],
True,
False,
{
"is_rls_supported": False,
"get_sqla_row_level_filters": MagicMock(return_value=[]),
},
"4f8109d3761e766e650af514bb358f10",
),
(
None,
[ExecutorType.CURRENT_USER],
False,
False,
None,
ExecutorNotFoundError(),
),
],
@@ -229,20 +354,29 @@ def test_chart_digest(
execute_as: list[ExecutorType],
has_current_user: bool,
use_custom_digest: bool,
rls_datasource: dict[str, Any] | None,
expected_result: str | Exception,
) -> None:
from superset import app
from superset import app, security_manager
from superset.models.slice import Slice
from superset.thumbnails.digest import get_chart_digest
# Mock datasource with RLS if provided
datasource = None
if rls_datasource:
datasource = prepare_datasource_mock(rls_datasource, SqlaTable)
# Prepare chart with the datasource in the constructor
kwargs = {
**_DEFAULT_CHART_KWARGS,
**(chart_overrides or {}),
}
chart = Slice(**kwargs)
user: User | None = None
if has_current_user:
user = User(id=1, username="1")
func = CUSTOM_CHART_FUNC if use_custom_digest else None
with (
@@ -253,6 +387,13 @@ def test_chart_digest(
"THUMBNAIL_CHART_DIGEST_FUNC": func,
},
),
patch.object(
type(chart),
"datasource",
new_callable=PropertyMock,
return_value=datasource,
),
patch.object(security_manager, "find_user", return_value=user),
override_user(user),
):
cm = (