(superset) 01/01: fix(Jinja): Extra cache keys for Jinja columns

[di...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

diegopucci pushed a commit to branch geido/fix/jinja-columns-cache-keys
in repository https://gitbox.apache.org/repos/asf/superset.git

commit d9c1823e65a50b13709fe7ea85d433989d97168d
Author: Diego Pucci <di...@gmail.com>
AuthorDate: Fri Oct 25 16:04:14 2024 +0300

    fix(Jinja): Extra cache keys for Jinja columns
---
 superset/connectors/sqla/models.py           |  6 ++++
 tests/integration_tests/sqla_models_tests.py | 49 ++++++++++++++++++++++++++++
 2 files changed, 55 insertions(+)

diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index e8aa0d705b..c85695efbe 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -1980,6 +1980,12 @@ class SqlaTable(
             templatable_statements.append(extras["where"])
         if "having" in extras:
             templatable_statements.append(extras["having"])
+        if "columns" in query_obj:
+            templatable_statements += [
+                c["sqlExpression"]
+                for c in query_obj["columns"]
+                if c.get("sqlExpression")
+            ]
         if self.is_rls_supported:
             templatable_statements += [
                 f.clause for f in security_manager.get_rls_filters(self)
diff --git a/tests/integration_tests/sqla_models_tests.py b/tests/integration_tests/sqla_models_tests.py
index 2d7f6bf041..bd774afa4b 100644
--- a/tests/integration_tests/sqla_models_tests.py
+++ b/tests/integration_tests/sqla_models_tests.py
@@ -911,6 +911,55 @@ def test_extra_cache_keys_in_sql_expression(
     assert extra_cache_keys == expected_cache_keys
 
 
+@pytest.mark.usefixtures("app_context")
+@pytest.mark.parametrize(
+    "sql_expression,expected_cache_keys,has_extra_cache_keys",
+    [
+        ("'{{ current_username() }}'", ["abc"], True),
+        ("(user != 'abc')", [], False),
+    ],
+)
+@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="abc@test.com")
+def test_extra_cache_keys_in_columns(
+    mock_user_email,
+    mock_username,
+    mock_user_id,
+    sql_expression,
+    expected_cache_keys,
+    has_extra_cache_keys,
+):
+    table = SqlaTable(
+        table_name="test_has_no_extra_cache_keys_table",
+        sql="SELECT 'abc' as user",
+        database=get_example_database(),
+    )
+    base_query_obj = {
+        "granularity": None,
+        "from_dttm": None,
+        "to_dttm": None,
+        "groupby": [],
+        "metrics": [],
+        "is_timeseries": False,
+        "filter": [],
+    }
+
+    query_obj = dict(
+        **base_query_obj,
+        columns=[
+            {
+                "expressionType": "SQL",
+                "sqlExpression": sql_expression,
+            }
+        ],
+    )
+
+    extra_cache_keys = table.get_extra_cache_keys(query_obj)
+    assert table.has_extra_cache_key_calls(query_obj) == has_extra_cache_keys
+    assert extra_cache_keys == expected_cache_keys
+
+
 @pytest.mark.usefixtures("app_context")
 @pytest.mark.parametrize(
     "row,dimension,result",