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

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
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 203 additions and 10 deletions

View File

@ -1046,14 +1046,14 @@ class DashboardRestApi(BaseSupersetModelRestApi):
cache_dashboard_screenshot.delay(
username=get_current_user(),
guest_token=g.user.guest_token
if isinstance(g.user, GuestUser)
if get_current_user() and isinstance(g.user, GuestUser)
else None,
dashboard_id=dashboard.id,
dashboard_url=dashboard_url,
cache_key=cache_key,
force=True,
thumb_size=thumb_size,
window_size=window_size,
cache_key=cache_key,
)
return self.response(
202,

View File

@ -67,6 +67,7 @@ from superset.security.guest_token import (
GuestUser,
)
from superset.sql_parse import extract_tables_from_jinja_sql, Table
from superset.tasks.utils import get_current_user
from superset.utils import json
from superset.utils.core import (
DatasourceName,
@ -2639,8 +2640,12 @@ class SupersetSecurityManager( # pylint: disable=too-many-public-methods
if not is_feature_enabled("EMBEDDED_SUPERSET"):
return False
if not user:
if not get_current_user():
return False
user = g.user
return hasattr(user, "is_guest_user") and user.is_guest_user
def get_current_guest_user_if_guest(self) -> Optional[GuestUser]:

View File

@ -114,10 +114,10 @@ def cache_dashboard_screenshot( # pylint: disable=too-many-arguments
dashboard_id: int,
dashboard_url: str,
force: bool = True,
cache_key: Optional[str] = None,
guest_token: Optional[GuestToken] = None,
thumb_size: Optional[WindowSize] = None,
window_size: Optional[WindowSize] = None,
cache_key: Optional[str] = None,
) -> None:
# pylint: disable=import-outside-toplevel
from superset.models.dashboard import Dashboard

View File

@ -22,11 +22,14 @@ from typing import TYPE_CHECKING
from flask import current_app
from superset import security_manager
from superset.tasks.types import ExecutorType
from superset.tasks.utils import get_current_user, get_executor
from superset.utils.core import override_user
from superset.utils.hashing import md5_sha_from_str
if TYPE_CHECKING:
from superset.connectors.sqla.models import BaseDatasource, SqlaTable
from superset.models.dashboard import Dashboard
from superset.models.slice import Slice
@ -49,8 +52,46 @@ def _adjust_string_for_executor(
return unique_string
def _adjust_string_with_rls(
unique_string: str,
datasources: list[SqlaTable | None] | set[BaseDatasource],
executor: str,
) -> str:
"""
Add the RLS filters to the unique string based on current executor.
"""
user = (
security_manager.find_user(executor)
or security_manager.get_current_guest_user_if_guest()
)
if user:
stringified_rls = ""
with override_user(user):
for datasource in datasources:
if (
datasource
and hasattr(datasource, "is_rls_supported")
and datasource.is_rls_supported
):
rls_filters = datasource.get_sqla_row_level_filters()
if len(rls_filters) > 0:
stringified_rls += (
f"{str(datasource.id)}\t"
+ "\t".join([str(f) for f in rls_filters])
+ "\n"
)
if stringified_rls:
unique_string = f"{unique_string}\n{stringified_rls}"
return unique_string
def get_dashboard_digest(dashboard: Dashboard) -> str:
config = current_app.config
datasources = dashboard.datasources
executor_type, executor = get_executor(
executor_types=config["THUMBNAIL_EXECUTE_AS"],
model=dashboard,
@ -65,19 +106,25 @@ def get_dashboard_digest(dashboard: Dashboard) -> str:
)
unique_string = _adjust_string_for_executor(unique_string, executor_type, executor)
unique_string = _adjust_string_with_rls(unique_string, datasources, executor)
return md5_sha_from_str(unique_string)
def get_chart_digest(chart: Slice) -> str:
config = current_app.config
datasource = chart.datasource
executor_type, executor = get_executor(
executor_types=config["THUMBNAIL_EXECUTE_AS"],
model=chart,
current_user=get_current_user(),
)
if func := config["THUMBNAIL_CHART_DIGEST_FUNC"]:
return func(chart, executor_type, executor)
unique_string = f"{chart.params or ''}.{executor}"
unique_string = _adjust_string_for_executor(unique_string, executor_type, executor)
unique_string = _adjust_string_with_rls(unique_string, [datasource], executor)
return md5_sha_from_str(unique_string)

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 = (