mirror of
https://github.com/apache/superset.git
synced 2026-05-08 01:15:46 +00:00
108 lines
4.0 KiB
Python
108 lines
4.0 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
from typing import Any
|
|
from unittest.mock import patch
|
|
|
|
from flask_appbuilder.security.sqla.models import User
|
|
|
|
from superset.models.core import Database
|
|
from superset.utils.cache_keys import add_impersonation_cache_key_if_needed
|
|
from superset.utils.core import override_user
|
|
|
|
|
|
def _flag(name: str):
|
|
"""Build a feature-flag side_effect that returns True only for ``name``."""
|
|
|
|
def side_effect(feature=None):
|
|
return feature == name
|
|
|
|
return side_effect
|
|
|
|
|
|
def _run(database: Database) -> dict[str, Any]:
|
|
"""Run the helper against a fresh dict and return that dict."""
|
|
cache_dict: dict[str, Any] = {}
|
|
add_impersonation_cache_key_if_needed(database, cache_dict)
|
|
return cache_dict
|
|
|
|
|
|
def test_no_per_user_caching_yields_no_key():
|
|
database = Database(database_name="d", sqlalchemy_uri="sqlite://")
|
|
with override_user(User(username="u")):
|
|
assert "impersonation_key" not in _run(database)
|
|
|
|
|
|
@patch("superset.utils.cache_keys.feature_flag_manager")
|
|
def test_cache_query_by_user_adds_username(feature_flag_mock):
|
|
feature_flag_mock.is_feature_enabled.side_effect = _flag("CACHE_QUERY_BY_USER")
|
|
database = Database(database_name="d", sqlalchemy_uri="sqlite://")
|
|
with override_user(User(username="alice")):
|
|
assert _run(database)["impersonation_key"] == "alice"
|
|
|
|
|
|
@patch("superset.utils.cache_keys.feature_flag_manager")
|
|
def test_cache_query_by_user_distinct_per_user(feature_flag_mock):
|
|
feature_flag_mock.is_feature_enabled.side_effect = _flag("CACHE_QUERY_BY_USER")
|
|
database = Database(database_name="d", sqlalchemy_uri="sqlite://")
|
|
with override_user(User(username="alice")):
|
|
key_a = _run(database)["impersonation_key"]
|
|
with override_user(User(username="bob")):
|
|
key_b = _run(database)["impersonation_key"]
|
|
assert key_a != key_b
|
|
|
|
|
|
@patch("superset.utils.cache_keys.feature_flag_manager")
|
|
def test_cache_impersonation_requires_database_flag(feature_flag_mock):
|
|
"""
|
|
CACHE_IMPERSONATION alone is not enough; ``database.impersonate_user`` must
|
|
also be set on the database for the per-user key to apply.
|
|
"""
|
|
feature_flag_mock.is_feature_enabled.side_effect = _flag("CACHE_IMPERSONATION")
|
|
|
|
db_no_impersonation = Database(database_name="d", sqlalchemy_uri="sqlite://")
|
|
db_with_impersonation = Database(
|
|
database_name="d", sqlalchemy_uri="sqlite://", impersonate_user=True
|
|
)
|
|
|
|
with override_user(User(username="alice")):
|
|
assert "impersonation_key" not in _run(db_no_impersonation)
|
|
assert _run(db_with_impersonation)["impersonation_key"] == "alice"
|
|
|
|
|
|
def test_per_user_caching_in_extra_json_enables_key():
|
|
database = Database(
|
|
database_name="d",
|
|
sqlalchemy_uri="sqlite://",
|
|
extra='{"per_user_caching": true}',
|
|
)
|
|
with override_user(User(username="alice")):
|
|
assert _run(database)["impersonation_key"] == "alice"
|
|
|
|
|
|
def test_no_user_yields_no_key(app_context): # noqa: ARG001
|
|
"""
|
|
With no logged-in user, the engine spec returns None even when per-user
|
|
caching is enabled — there's no identity to key on.
|
|
"""
|
|
database = Database(
|
|
database_name="d",
|
|
sqlalchemy_uri="sqlite://",
|
|
extra='{"per_user_caching": true}',
|
|
)
|
|
# No override_user — g.user is unset
|
|
assert "impersonation_key" not in _run(database)
|