Skip to content

Commit

Permalink
feat(jinja): current_user_email macro (#27197)
Browse files Browse the repository at this point in the history
  • Loading branch information
Vitor-Avila committed Feb 23, 2024
1 parent 11760d3 commit 1d571ec
Show file tree
Hide file tree
Showing 5 changed files with 116 additions and 22 deletions.
23 changes: 19 additions & 4 deletions docs/docs/installation/sql-templating.mdx
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ In this section, we'll walkthrough the pre-defined Jinja macros in Superset.

**Current Username**

The `{{ current_username() }}` macro returns the username of the currently logged in user.
The `{{ current_username() }}` macro returns the `username` of the currently logged in user.

If you have caching enabled in your Superset configuration, then by default the `username` value will be used
by Superset when calculating the cache key. A cache key is a unique identifier that determines if there's a
Expand All @@ -189,19 +189,34 @@ cache key by adding the following parameter to your Jinja code:

**Current User ID**

The `{{ current_user_id() }}` macro returns the user_id of the currently logged in user.
The `{{ current_user_id() }}` macro returns the account ID of the currently logged in user.

If you have caching enabled in your Superset configuration, then by default the `user_id` value will be used
If you have caching enabled in your Superset configuration, then by default the account `id` value will be used
by Superset when calculating the cache key. A cache key is a unique identifier that determines if there's a
cache hit in the future and Superset can retrieve cached data.

You can disable the inclusion of the `user_id` value in the calculation of the
You can disable the inclusion of the account `id` value in the calculation of the
cache key by adding the following parameter to your Jinja code:

```
{{ current_user_id(add_to_cache_keys=False) }}
```

**Current User Email**

The `{{ current_user_email() }}` macro returns the email address of the currently logged in user.

If you have caching enabled in your Superset configuration, then by default the email address value will be used
by Superset when calculating the cache key. A cache key is a unique identifier that determines if there's a
cache hit in the future and Superset can retrieve cached data.

You can disable the inclusion of the email value in the calculation of the
cache key by adding the following parameter to your Jinja code:

```
{{ current_user_email(add_to_cache_keys=False) }}
```

**Custom URL Parameters**

The `{{ url_param('custom_variable') }}` macro lets you define arbitrary URL
Expand Down
36 changes: 28 additions & 8 deletions superset/jinja_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
from typing import Any, Callable, cast, Optional, TYPE_CHECKING, TypedDict, Union

import dateutil
from flask import current_app, g, has_request_context, request
from flask import current_app, has_request_context, request
from flask_babel import gettext as _
from jinja2 import DebugUndefined
from jinja2.sandbox import SandboxedEnvironment
Expand All @@ -36,7 +36,9 @@
from superset.extensions import feature_flag_manager
from superset.utils.core import (
convert_legacy_filters_into_adhoc,
get_user,
get_user_email,
get_user_id,
get_username,
merge_extra_filters,
)

Expand Down Expand Up @@ -85,6 +87,7 @@ class ExtraCache:
r"\{\{.*("
r"current_user_id\(.*\)|"
r"current_username\(.*\)|"
r"current_user_email\(.*\)|"
r"cache_key_wrapper\(.*\)|"
r"url_param\(.*\)"
r").*\}\}"
Expand All @@ -110,10 +113,10 @@ def current_user_id(self, add_to_cache_keys: bool = True) -> Optional[int]:
:returns: The user ID
"""

if user := get_user():
if user_id := get_user_id():
if add_to_cache_keys:
self.cache_key_wrapper(user.id)
return user.id
self.cache_key_wrapper(user_id)
return user_id
return None

def current_username(self, add_to_cache_keys: bool = True) -> Optional[str]:
Expand All @@ -124,10 +127,24 @@ def current_username(self, add_to_cache_keys: bool = True) -> Optional[str]:
:returns: The username
"""

if g.user and hasattr(g.user, "username"):
if username := get_username():
if add_to_cache_keys:
self.cache_key_wrapper(g.user.username)
return g.user.username
self.cache_key_wrapper(username)
return username
return None

def current_user_email(self, add_to_cache_keys: bool = True) -> Optional[str]:
"""
Return the email address of the user who is currently logged in.
:param add_to_cache_keys: Whether the value should be included in the cache key
:returns: The user email address
"""

if email_address := get_user_email():
if add_to_cache_keys:
self.cache_key_wrapper(email_address)
return email_address
return None

def cache_key_wrapper(self, key: Any) -> Any:
Expand Down Expand Up @@ -530,6 +547,9 @@ def set_context(self, **kwargs: Any) -> None:
"url_param": partial(safe_proxy, extra_cache.url_param),
"current_user_id": partial(safe_proxy, extra_cache.current_user_id),
"current_username": partial(safe_proxy, extra_cache.current_username),
"current_user_email": partial(
safe_proxy, extra_cache.current_user_email
),
"cache_key_wrapper": partial(safe_proxy, extra_cache.cache_key_wrapper),
"filter_values": partial(safe_proxy, extra_cache.filter_values),
"get_filters": partial(safe_proxy, extra_cache.get_filters),
Expand Down
13 changes: 13 additions & 0 deletions superset/utils/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -1384,6 +1384,19 @@ def get_user_id() -> int | None:
return None


def get_user_email() -> str | None:
"""
Get the email (if defined) associated with the current user.
:returns: The email
"""

try:
return g.user.email
except Exception: # pylint: disable=broad-except
return None


@contextmanager
def override_user(user: User | None, force: bool = True) -> Iterator[Any]:
"""
Expand Down
28 changes: 18 additions & 10 deletions tests/integration_tests/sqla_models_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -132,14 +132,15 @@ def test_db_column_types(self):
col = TableColumn(column_name="foo", type=str_type, table=tbl, is_dttm=True)
self.assertTrue(col.is_temporal)

@patch("superset.jinja_context.g")
def test_extra_cache_keys(self, flask_g):
flask_g.user.username = "abc"
@patch("superset.jinja_context.get_user_id", return_value=1)
@patch("superset.jinja_context.get_username", return_value="abc")
@patch("superset.jinja_context.get_user_email", return_value="[email protected]")
def test_extra_cache_keys(self, mock_user_email, mock_username, mock_user_id):
base_query_obj = {
"granularity": None,
"from_dttm": None,
"to_dttm": None,
"groupby": ["user"],
"groupby": ["id", "username", "email"],
"metrics": [],
"is_timeseries": False,
"filter": [],
Expand All @@ -148,19 +149,27 @@ def test_extra_cache_keys(self, flask_g):
# Table with Jinja callable.
table1 = SqlaTable(
table_name="test_has_extra_cache_keys_table",
sql="SELECT '{{ current_username() }}' as user",
sql="""
SELECT '{{ current_user_id() }}' as id,
SELECT '{{ current_username() }}' as username,
SELECT '{{ current_user_email() }}' as email,
""",
database=get_example_database(),
)

query_obj = dict(**base_query_obj, extras={})
extra_cache_keys = table1.get_extra_cache_keys(query_obj)
self.assertTrue(table1.has_extra_cache_key_calls(query_obj))
assert extra_cache_keys == ["abc"]
assert extra_cache_keys == [1, "abc", "[email protected]"]

# Table with Jinja callable disabled.
table2 = SqlaTable(
table_name="test_has_extra_cache_keys_disabled_table",
sql="SELECT '{{ current_username(False) }}' as user",
sql="""
SELECT '{{ current_user_id(False) }}' as id,
SELECT '{{ current_username(False) }}' as username,
SELECT '{{ current_user_email(False) }}' as email,
""",
database=get_example_database(),
)
query_obj = dict(**base_query_obj, extras={})
Expand Down Expand Up @@ -189,9 +198,8 @@ def test_extra_cache_keys(self, flask_g):
self.assertTrue(table3.has_extra_cache_key_calls(query_obj))
assert extra_cache_keys == ["abc"]

@patch("superset.jinja_context.g")
def test_jinja_metrics_and_calc_columns(self, flask_g):
flask_g.user.username = "abc"
@patch("superset.jinja_context.get_username", return_value="abc")
def test_jinja_metrics_and_calc_columns(self, mock_username):
base_query_obj = {
"granularity": None,
"from_dttm": None,
Expand Down
38 changes: 38 additions & 0 deletions tests/unit_tests/test_jinja_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
# under the License.
import json
from typing import Any
from unittest.mock import patch

import pytest
from sqlalchemy.dialects.postgresql import dialect
Expand Down Expand Up @@ -265,3 +266,40 @@ def func(input_: Any) -> Any:

with pytest.raises(SupersetTemplateException):
safe_proxy(func, {"foo": lambda: "bar"})


@patch("superset.jinja_context.ExtraCache.cache_key_wrapper")
@patch("superset.utils.core.g")
def test_user_macros(mock_flask_user, mock_cache_key_wrapper):
mock_flask_user.user.id = 1
mock_flask_user.user.username = "my_username"
mock_flask_user.user.email = "[email protected]"
cache = ExtraCache()
assert cache.current_user_id() == 1
assert cache.current_username() == "my_username"
assert cache.current_user_email() == "[email protected]"
assert mock_cache_key_wrapper.call_count == 3


@patch("superset.jinja_context.ExtraCache.cache_key_wrapper")
@patch("superset.utils.core.g")
def test_user_macros_without_cache_key_inclusion(
mock_flask_user, mock_cache_key_wrapper
):
mock_flask_user.user.id = 1
mock_flask_user.user.username = "my_username"
mock_flask_user.user.email = "[email protected]"
cache = ExtraCache()
assert cache.current_user_id(False) == 1
assert cache.current_username(False) == "my_username"
assert cache.current_user_email(False) == "[email protected]"
assert mock_cache_key_wrapper.call_count == 0


@patch("superset.utils.core.g")
def test_user_macros_without_user_info(mock_flask_user):
mock_flask_user.user = None
cache = ExtraCache()
assert cache.current_user_id() == None
assert cache.current_username() == None
assert cache.current_user_email() == None

0 comments on commit 1d571ec

Please sign in to comment.