Merge branch 'master' into hughhhh/zagreb-embed-page

tests/unit_tests/utils/oauth2_tests.py

@@ -137,6 +137,7 @@ def test_refresh_oauth2_token_deletes_token_on_oauth2_exception(
         refresh_oauth2_token(DUMMY_OAUTH2_CONFIG, 1, 1, db_engine_spec, token)
 
     db.session.delete.assert_called_with(token)
+    db.session.flush.assert_called_once()
 
 
 def test_refresh_oauth2_token_keeps_token_on_other_exception(
@@ -338,6 +339,45 @@ def test_encode_decode_oauth2_state(
     assert decoded["user_id"] == 2
 
 
+def test_get_oauth2_access_token_lock_not_acquired_no_error_log(
+    mocker: MockerFixture,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """
+    Test that when a distributed lock can't be acquired, no error is logged and
+    the function returns None instead of raising.
+
+    This scenario occurs when a dashboard with multiple charts from the same
+    OAuth2-enabled DB has an expired token: simultaneous requests compete for
+    the lock, and only the first one wins. The rest should silently return None.
+    """
+    import logging
+
+    from superset.exceptions import AcquireDistributedLockFailedException
+
+    mocker.patch("time.sleep")  # avoid backoff delays in tests
+
+    db = mocker.patch("superset.utils.oauth2.db")
+    db_engine_spec = mocker.MagicMock()
+    token = mocker.MagicMock()
+    token.access_token = "access-token"  # noqa: S105
+    token.access_token_expiration = datetime(2024, 1, 1)
+    token.refresh_token = "refresh-token"  # noqa: S105
+    db.session.query().filter_by().one_or_none.return_value = token
+
+    mocker.patch(
+        "superset.utils.oauth2.refresh_oauth2_token",
+        side_effect=AcquireDistributedLockFailedException("Lock not available"),
+    )
+
+    with freeze_time("2024-01-02"):
+        with caplog.at_level(logging.DEBUG):
+            result = get_oauth2_access_token({}, 1, 1, db_engine_spec)
+
+    assert result is None
+    assert not any(record.levelno >= logging.ERROR for record in caplog.records)
+
+
 def test_get_oauth2_redirect_uri_from_config(mocker: MockerFixture) -> None:
     """
     Test that get_oauth2_redirect_uri returns the configured value when set.

tests/unit_tests/utils/test_database.py

@@ -0,0 +1,53 @@
+# 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.
+"""Tests for superset.utils.database module."""
+
+import pytest
+from sqlalchemy import Sequence
+from sqlalchemy.dialects import mysql, postgresql
+from sqlalchemy.schema import CreateSequence
+from sqlalchemy.sql.compiler import DDLCompiler
+
+from superset.utils.database import apply_mariadb_ddl_fix
+
+
+@pytest.fixture(scope="module", autouse=True)
+def setup_mariadb_ddl_fix():
+    """Apply MariaDB DDL fix once per module before tests run."""
+    apply_mariadb_ddl_fix()
+
+
+def test_mariadb_nocycle_fix_applied():
+    """Test that 'NO CYCLE' is replaced with 'NOCYCLE' for MariaDB dialect."""
+    dialect = mysql.dialect()
+    dialect.name = "mariadb"
+    ddl_compiler = DDLCompiler(dialect, None)
+    seq = Sequence("test_seq", cycle=False)
+
+    result = ddl_compiler.visit_create_sequence(CreateSequence(seq))
+    assert "NOCYCLE" in result
+    assert "NO CYCLE" not in result
+
+
+def test_nocycle_fix_not_applied_for_postgresql():
+    """Test that 'NO CYCLE' is NOT replaced for PostgreSQL dialect."""
+    dialect = postgresql.dialect()
+    compiler = DDLCompiler(dialect, None)
+    seq = Sequence("test_seq", cycle=False)
+
+    result = compiler.visit_create_sequence(CreateSequence(seq))
+    assert "NO CYCLE" in result

tests/unit_tests/utils/test_impersonation_cache_key.py

@@ -0,0 +1,107 @@
+# 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)

tests/unit_tests/utils/webdriver_test.py

@@ -640,6 +640,114 @@ class TestWebDriverPlaywrightErrorHandling:
         mock_button.click.assert_called_once()
         mock_close_button.click.assert_called_once()
 
+    @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", True)
+    @patch("superset.utils.webdriver.sync_playwright")
+    @patch("superset.utils.webdriver.app")
+    def test_uses_wait_for_function_to_detect_spinners(
+        self, mock_app, mock_sync_playwright
+    ):
+        """wait_for_function polls for spinner absence rather than snapshotting."""
+        mock_user = MagicMock()
+        mock_user.username = "test_user"
+        mock_app.config = {
+            "WEBDRIVER_OPTION_ARGS": [],
+            "WEBDRIVER_WINDOW": {"pixel_density": 1},
+            "SCREENSHOT_PLAYWRIGHT_DEFAULT_TIMEOUT": 30000,
+            "SCREENSHOT_PLAYWRIGHT_WAIT_EVENT": "networkidle",
+            "SCREENSHOT_SELENIUM_HEADSTART": 0,
+            "SCREENSHOT_SELENIUM_ANIMATION_WAIT": 0,
+            "SCREENSHOT_REPLACE_UNEXPECTED_ERRORS": False,
+            "SCREENSHOT_TILED_ENABLED": False,
+            "SCREENSHOT_LOCATE_WAIT": 10,
+            "SCREENSHOT_LOAD_WAIT": 60,
+            "SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE": 10,
+            "SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE": 10,
+        }
+
+        mock_playwright_instance = MagicMock()
+        mock_browser = MagicMock()
+        mock_context = MagicMock()
+        mock_page = MagicMock()
+        mock_element = MagicMock()
+
+        mock_sync_playwright.return_value.__enter__.return_value = (
+            mock_playwright_instance
+        )
+        mock_playwright_instance.chromium.launch.return_value = mock_browser
+        mock_browser.new_context.return_value = mock_context
+        mock_context.new_page.return_value = mock_page
+        mock_page.locator.return_value = mock_element
+        mock_element.screenshot.return_value = b"screenshot"
+
+        with patch.object(WebDriverPlaywright, "auth", return_value=mock_context):
+            driver = WebDriverPlaywright("chrome")
+            driver.get_screenshot("http://example.com", "test-element", mock_user)
+
+        mock_page.wait_for_function.assert_called_once_with(
+            "() => document.querySelectorAll('.loading').length === 0",
+            timeout=60 * 1000,
+        )
+        # Guard against reintroducing the old snapshot-based approach
+        loading_locator_calls = [
+            c for c in mock_page.locator.call_args_list if c.args == (".loading",)
+        ]
+        assert loading_locator_calls == []
+
+    @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", True)
+    @patch("superset.utils.webdriver.sync_playwright")
+    @patch("superset.utils.webdriver.logger")
+    @patch("superset.utils.webdriver.app")
+    def test_spinner_timeout_logs_warning_and_raises(
+        self, mock_app, mock_logger, mock_sync_playwright
+    ):
+        """Spinner timeout is logged as a warning and re-raised."""
+        from superset.utils.webdriver import PlaywrightTimeout
+
+        mock_user = MagicMock()
+        mock_user.username = "test_user"
+        mock_app.config = {
+            "WEBDRIVER_OPTION_ARGS": [],
+            "WEBDRIVER_WINDOW": {"pixel_density": 1},
+            "SCREENSHOT_PLAYWRIGHT_DEFAULT_TIMEOUT": 30000,
+            "SCREENSHOT_PLAYWRIGHT_WAIT_EVENT": "networkidle",
+            "SCREENSHOT_SELENIUM_HEADSTART": 0,
+            "SCREENSHOT_SELENIUM_ANIMATION_WAIT": 0,
+            "SCREENSHOT_REPLACE_UNEXPECTED_ERRORS": False,
+            "SCREENSHOT_TILED_ENABLED": False,
+            "SCREENSHOT_LOCATE_WAIT": 10,
+            "SCREENSHOT_LOAD_WAIT": 60,
+            "SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE": 10,
+            "SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE": 10,
+        }
+
+        mock_playwright_instance = MagicMock()
+        mock_browser = MagicMock()
+        mock_context = MagicMock()
+        mock_page = MagicMock()
+        mock_element = MagicMock()
+
+        mock_sync_playwright.return_value.__enter__.return_value = (
+            mock_playwright_instance
+        )
+        mock_playwright_instance.chromium.launch.return_value = mock_browser
+        mock_browser.new_context.return_value = mock_context
+        mock_context.new_page.return_value = mock_page
+        mock_page.locator.return_value = mock_element
+
+        timeout = PlaywrightTimeout()
+        mock_page.wait_for_function.side_effect = timeout
+
+        with patch.object(WebDriverPlaywright, "auth", return_value=mock_context):
+            driver = WebDriverPlaywright("chrome")
+            with pytest.raises(PlaywrightTimeout) as exc_info:
+                driver.get_screenshot("http://example.com", "test-element", mock_user)
+
+        assert exc_info.value is timeout
+        mock_logger.warning.assert_any_call(
+            "Timed out waiting for charts to load at url %s",
+            "http://example.com",
+        )
+
     @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", True)
     @patch("superset.utils.webdriver.sync_playwright")
     @patch("superset.utils.webdriver.logger")