feat: current_user_rls_rules Jinja macro (#33614)

tests/unit_tests/jinja_context_test.py

@@ -31,7 +31,12 @@ from sqlalchemy.dialects.postgresql import dialect
 
 from superset import app
 from superset.commands.dataset.exceptions import DatasetNotFoundError
-from superset.connectors.sqla.models import SqlaTable, SqlMetric, TableColumn
+from superset.connectors.sqla.models import (
+    RowLevelSecurityFilter,
+    SqlaTable,
+    SqlMetric,
+    TableColumn,
+)
 from superset.exceptions import SupersetTemplateException
 from superset.jinja_context import (
     dataset_macro,
@@ -46,6 +51,7 @@ from superset.jinja_context import (
 from superset.models.core import Database
 from superset.models.slice import Slice
 from superset.utils import json
+from tests.unit_tests.conftest import with_feature_flags
 
 
 def test_filter_values_adhoc_filters() -> None:
@@ -355,16 +361,29 @@ def test_safe_proxy_nested_lambda() -> None:
         safe_proxy(func, {"foo": lambda: "bar"})
 
 
-def test_user_macros(mocker: MockerFixture):
+@pytest.mark.parametrize(
+    "add_to_cache_keys,mock_cache_key_wrapper_call_count",
+    [
+        (True, 4),
+        (False, 0),
+    ],
+)
+def test_user_macros(
+    mocker: MockerFixture,
+    add_to_cache_keys: bool,
+    mock_cache_key_wrapper_call_count: int,
+):
     """
     Test all user macros:
         - ``current_user_id``
         - ``current_username``
         - ``current_user_email``
         - ``current_user_roles``
+        - ``current_user_rls_rules``
     """
     mock_g = mocker.patch("superset.utils.core.g")
     mock_get_user_roles = mocker.patch("superset.security_manager.get_user_roles")
+    mock_get_user_rls = mocker.patch("superset.security_manager.get_rls_filters")
     mock_cache_key_wrapper = mocker.patch(
         "superset.jinja_context.ExtraCache.cache_key_wrapper"
     )
@@ -372,36 +391,20 @@ def test_user_macros(mocker: MockerFixture):
     mock_g.user.username = "my_username"
     mock_g.user.email = "my_email@test.com"
     mock_get_user_roles.return_value = [Role(name="my_role1"), Role(name="my_role2")]
-    cache = ExtraCache()
-    assert cache.current_user_id() == 1
-    assert cache.current_username() == "my_username"
-    assert cache.current_user_email() == "my_email@test.com"
-    assert cache.current_user_roles() == ["my_role1", "my_role2"]
-    assert mock_cache_key_wrapper.call_count == 4
+    mock_get_user_rls.return_value = [
+        RowLevelSecurityFilter(group_key="test", clause="1=1"),
+        RowLevelSecurityFilter(group_key="other_test", clause="product_id=1"),
+    ]
+    cache = ExtraCache(table=mocker.MagicMock())
+    assert cache.current_user_id(add_to_cache_keys) == 1
+    assert cache.current_username(add_to_cache_keys) == "my_username"
+    assert cache.current_user_email(add_to_cache_keys) == "my_email@test.com"
+    assert cache.current_user_roles(add_to_cache_keys) == ["my_role1", "my_role2"]
+    assert mock_cache_key_wrapper.call_count == mock_cache_key_wrapper_call_count
 
-    mock_get_user_roles.return_value = []
-    assert cache.current_user_roles() is None
-
-
-def test_user_macros_without_cache_key_inclusion(mocker: MockerFixture):
-    """
-    Test all user macros with ``add_to_cache_keys`` set to ``False``.
-    """
-    mock_g = mocker.patch("superset.utils.core.g")
-    mock_get_user_roles = mocker.patch("superset.security_manager.get_user_roles")
-    mock_cache_key_wrapper = mocker.patch(
-        "superset.jinja_context.ExtraCache.cache_key_wrapper"
-    )
-    mock_g.user.id = 1
-    mock_g.user.username = "my_username"
-    mock_g.user.email = "my_email@test.com"
-    mock_get_user_roles.return_value = [Role(name="my_role1"), Role(name="my_role2")]
-    cache = ExtraCache()
-    assert cache.current_user_id(False) == 1
-    assert cache.current_username(False) == "my_username"
-    assert cache.current_user_email(False) == "my_email@test.com"
-    assert cache.current_user_roles(False) == ["my_role1", "my_role2"]
-    assert mock_cache_key_wrapper.call_count == 0
+    # Testing {{ current_user_rls_rules() }} macro isolated and always without
+    # the param because it does not support it to avoid shared cache.
+    assert cache.current_user_rls_rules() == ["1=1", "product_id=1"]
 
 
 def test_user_macros_without_user_info(mocker: MockerFixture):
@@ -410,11 +413,55 @@ def test_user_macros_without_user_info(mocker: MockerFixture):
     """
     mock_g = mocker.patch("superset.utils.core.g")
     mock_g.user = None
+    cache = ExtraCache(table=mocker.MagicMock())
+    assert cache.current_user_id() is None
+    assert cache.current_username() is None
+    assert cache.current_user_email() is None
+    assert cache.current_user_roles() is None
+    assert cache.current_user_rls_rules() is None
+
+
+def test_current_user_rls_rules_with_no_table(mocker: MockerFixture):
+    """
+    Test the ``current_user_rls_rules`` macro when no table is provided.
+    """
+    mock_g = mocker.patch("superset.utils.core.g")
+    mock_get_user_rls = mocker.patch("superset.security_manager.get_rls_filters")
+    mock_is_guest_user = mocker.patch("superset.security_manager.is_guest_user")
+    mock_cache_key_wrapper = mocker.patch(
+        "superset.jinja_context.ExtraCache.cache_key_wrapper"
+    )
+    mock_g.user.id = 1
+    mock_g.user.username = "my_username"
+    mock_g.user.email = "my_email@test.com"
     cache = ExtraCache()
-    assert cache.current_user_id() == None  # noqa: E711
-    assert cache.current_username() == None  # noqa: E711
-    assert cache.current_user_email() == None  # noqa: E711
-    assert cache.current_user_roles() == None  # noqa: E711
+    assert cache.current_user_rls_rules() is None
+    assert mock_cache_key_wrapper.call_count == 0
+    assert mock_get_user_rls.call_count == 0
+    assert mock_is_guest_user.call_count == 0
+
+
+@with_feature_flags(EMBEDDED_SUPERSET=True)
+def test_current_user_rls_rules_guest_user(mocker: MockerFixture):
+    """
+    Test the ``current_user_rls_rules`` with an embedded user.
+    """
+    mock_g = mocker.patch("superset.utils.core.g")
+    mock_gg = mocker.patch("superset.tasks.utils.g")
+    mock_ggg = mocker.patch("superset.security.manager.g")
+    mock_get_user_rls = mocker.patch("superset.security_manager.get_guest_rls_filters")
+    mock_user = mocker.MagicMock()
+    mock_user.username = "my_username"
+    mock_user.is_guest_user = True
+    mock_user.is_anonymous = False
+    mock_g.user = mock_gg.user = mock_ggg.user = mock_user
+
+    mock_get_user_rls.return_value = [
+        {"group_key": "test", "clause": "1=1"},
+        {"group_key": "other_test", "clause": "product_id=1"},
+    ]
+    cache = ExtraCache(table=mocker.MagicMock())
+    assert cache.current_user_rls_rules() == ["1=1", "product_id=1"]
 
 
 def test_where_in() -> None: