From bdf8ce69abb7be43f96d3e1c0dca9bc05afcf3e5 Mon Sep 17 00:00:00 2001 From: madhushreeag Date: Tue, 11 Aug 2026 16:06:26 -0700 Subject: [PATCH] feat(Reports&Alerts): remove Selenium support, require Playwright (#43028) Co-authored-by: madhushree agarwal --- UPDATING.md | 26 + docs/static/feature-flags.json | 6 - pyproject.toml | 4 +- requirements/base.txt | 33 +- requirements/development.txt | 47 +- superset/config.py | 46 +- superset/mcp_service/screenshot/__init__.py | 19 +- .../screenshot/pooled_screenshot.py | 492 ------------- .../screenshot/webdriver_config.py | 143 ---- .../mcp_service/screenshot/webdriver_pool.py | 437 ------------ superset/tasks/cache.py | 45 +- superset/utils/machine_auth.py | 33 +- superset/utils/screenshots.py | 58 +- superset/utils/webdriver.py | 668 +----------------- tests/integration_tests/thumbnails_tests.py | 113 +-- .../utils/machine_auth_tests.py | 56 +- .../mcp_service/test_pooled_screenshot.py | 22 +- tests/unit_tests/tasks/test_cache.py | 51 +- tests/unit_tests/utils/screenshot_test.py | 81 +-- .../test_playwright_migration_working.py | 199 ------ tests/unit_tests/utils/webdriver_test.py | 554 +-------------- 21 files changed, 188 insertions(+), 2945 deletions(-) delete mode 100644 superset/mcp_service/screenshot/pooled_screenshot.py delete mode 100644 superset/mcp_service/screenshot/webdriver_config.py delete mode 100644 superset/mcp_service/screenshot/webdriver_pool.py delete mode 100644 tests/unit_tests/utils/test_playwright_migration_working.py diff --git a/UPDATING.md b/UPDATING.md index e4277ba33fe..199a1c33524 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -28,6 +28,32 @@ assists people when migrating to a new version. - [42393](https://github.com/apache/superset/pull/42393): Exported dataset YAML now carries a `uuid` for each metric and column so that custom folder assignments (which reference metrics/columns by UUID) survive an import into another workspace. This affects any export bundle that contains datasets, not just a dataset export: chart, dashboard, database and full-asset exports all embed the same dataset YAML, so a dashboard exported from this release also fails to import into an older one even though no dataset was exported directly. As with `folders` and `currency_code_column`, the affected `datasets/` files fail schema validation (`Unknown field: uuid`) when imported into Superset releases that predate this change; regenerate or hand-edit exports for older targets in mixed-version fleets. - [42087](https://github.com/apache/superset/pull/42087): Stored calculated-column and metric expressions are validated when a query is built, under the same sub-query policy already applied to adhoc expressions. Previously only the dataset update path checked them on save, so expressions written by v1 import, by dataset duplication, or before that check existed were never validated. Since `ALLOW_ADHOC_SUBQUERY` defaults to `False` (see [19242](https://github.com/apache/superset/pull/19242)), a dataset whose stored expression contains a sub-query works before upgrading and afterwards fails at chart render with `Custom SQL fields cannot contain sub-queries.` There is no migration step, and the error does not name the offending dataset column, so audit stored expressions before upgrading: either rewrite them without the sub-query, or set `ALLOW_ADHOC_SUBQUERY = True` to keep the previous behaviour for both stored and adhoc expressions. +### Selenium support removed — Playwright is now required for screenshots + +Selenium support has been removed. **Playwright is now required** for all +report and thumbnail screenshot generation. Install it with: + +```bash +pip install playwright && playwright install chromium +``` + +**Breaking config changes:** + +- `PLAYWRIGHT_REPORTS_AND_THUMBNAILS` feature flag removed (Playwright is the only backend now) +- `WEBDRIVER_TYPE` config key removed (Playwright always uses Chromium) +- `WEBDRIVER_CONFIGURATION` config key removed (Selenium-only) +- `SCREENSHOT_PAGE_LOAD_WAIT` config key removed (Selenium-only) +- `SCREENSHOT_SELENIUM_RETRIES` config key removed (Selenium-only) +- `SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE` config key removed (Selenium-only) +- `SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE` config key removed (Selenium-only) +- `MachineAuthProvider.authenticate_webdriver()` removed; use `authenticate_browser_context()` instead + +**What operators should do before upgrading:** + +1. Install Playwright: `pip install playwright && playwright install chromium` +2. Remove any references to the removed config keys from custom `superset_config.py` +3. If you subclassed `MachineAuthProvider`, remove any `authenticate_webdriver` override and migrate auth logic to `authenticate_browser_context` + ### Soft delete is on by default, and purging is live `SOFT_DELETE` now ships **on** (`DEFAULT_FEATURE_FLAGS`), so deleting a diff --git a/docs/static/feature-flags.json b/docs/static/feature-flags.json index 19171a6a06f..0a843f20293 100644 --- a/docs/static/feature-flags.json +++ b/docs/static/feature-flags.json @@ -221,12 +221,6 @@ "lifecycle": "testing", "description": "When impersonating a user, use the email prefix instead of username" }, - { - "name": "PLAYWRIGHT_REPORTS_AND_THUMBNAILS", - "default": false, - "lifecycle": "testing", - "description": "Replace Selenium with Playwright for reports and thumbnails. Supports deck.gl visualizations. Requires playwright pip package." - }, { "name": "RLS_IN_SQLLAB", "default": false, diff --git a/pyproject.toml b/pyproject.toml index 5b0e45c1c6a..5099bf0aa2e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -110,7 +110,7 @@ dependencies = [ "PyJWT>=2.4.0, <3.0", "redis>=5.0.0, <9.0", "rison>=2.0.1, <3.0", - "selenium>=4.46.0, <5.0", + "shillelagh[gsheetsapi]>=1.4.5, <2.0", "sshtunnel>=0.4.0, <0.5", "simplejson>=4.1.1", @@ -309,7 +309,7 @@ combine_as_imports = true include_trailing_comma = true line_length = 88 known_first_party = "superset, apache-superset-core, apache-superset-extensions-cli" -known_third_party = "alembic, apispec, backoff, celery, click, colorama, cron_descriptor, croniter, cryptography, dateutil, deprecation, flask, flask_appbuilder, flask_babel, flask_caching, flask_compress, flask_jwt_extended, flask_login, flask_migrate, flask_sqlalchemy, flask_talisman, flask_testing, flask_wtf, freezegun, geohash, geopy, holidays, humanize, isodate, jinja2, jwt, markdown, markupsafe, marshmallow, marshmallow-union, msgpack, nh3, numpy, pandas, parameterized, parsedatetime, pgsanity, polyline, rison, progress, pyarrow, sqlalchemy_bigquery, pyhive, pyparsing, pytest, pytest_mock, pytz, redis, requests, selenium, setuptools, shillelagh, simplejson, slack, sqlalchemy, sqlalchemy_utils, syntaqlite, typing_extensions, urllib3, werkzeug, wtforms, wtforms_json, yaml" +known_third_party = "alembic, apispec, backoff, celery, click, colorama, cron_descriptor, croniter, cryptography, dateutil, deprecation, flask, flask_appbuilder, flask_babel, flask_caching, flask_compress, flask_jwt_extended, flask_login, flask_migrate, flask_sqlalchemy, flask_talisman, flask_testing, flask_wtf, freezegun, geohash, geopy, holidays, humanize, isodate, jinja2, jwt, markdown, markupsafe, marshmallow, marshmallow-union, msgpack, nh3, numpy, pandas, parameterized, parsedatetime, pgsanity, polyline, rison, progress, pyarrow, sqlalchemy_bigquery, pyhive, pyparsing, pytest, pytest_mock, pytz, redis, requests, setuptools, shillelagh, simplejson, slack, sqlalchemy, sqlalchemy_utils, syntaqlite, typing_extensions, urllib3, werkzeug, wtforms, wtforms_json, yaml" multi_line_output = 3 order_by_type = false diff --git a/requirements/base.txt b/requirements/base.txt index 492323c6406..39e1ff1cab3 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -20,10 +20,8 @@ attrs==25.3.0 # via # cattrs # jsonschema - # outcome # referencing # requests-cache - # trio babel==2.17.0 # via flask-babel backoff==2.2.1 @@ -53,9 +51,7 @@ cattrs==25.1.1 celery==5.6.3 # via apache-superset (pyproject.toml) certifi==2026.5.20 - # via - # requests - # selenium + # via requests cffi==2.0.0 # via # cryptography @@ -174,8 +170,6 @@ greenlet==3.5.4 # sqlalchemy gunicorn==26.0.0 # via apache-superset (pyproject.toml) -h11==0.16.0 - # via wsproto hashids==1.3.1 # via apache-superset (pyproject.toml) holidays==0.102 @@ -186,7 +180,6 @@ idna==3.15 # via # email-validator # requests - # trio # url-normalize isodate==0.7.2 # via @@ -265,10 +258,6 @@ openpyxl==3.1.5 # via pandas ordered-set==4.1.0 # via flask-limiter -outcome==1.3.0.post0 - # via - # trio - # trio-websocket packaging==25.0 # via # apache-superset (pyproject.toml) @@ -332,8 +321,6 @@ pyopenssl==26.4.0 # shillelagh pyparsing==3.3.2 # via apache-superset (pyproject.toml) -pysocks==1.7.1 - # via urllib3 python-calamine==0.8.2 # via pandas python-dateutil==2.9.0.post0 @@ -380,8 +367,6 @@ rpds-py==0.25.0 # via # jsonschema # referencing -selenium==4.46.0 - # via apache-superset (pyproject.toml) setuptools==80.9.0 # via -r requirements/base.in shillelagh==1.4.5 @@ -396,10 +381,6 @@ six==1.17.0 # wtforms-json slack-sdk==3.43.0 # via apache-superset (pyproject.toml) -sniffio==1.3.1 - # via trio -sortedcontainers==2.4.0 - # via trio sqlalchemy==1.4.54 # via # apache-superset (pyproject.toml) @@ -426,12 +407,6 @@ sshtunnel==0.4.0 # via apache-superset (pyproject.toml) tabulate==0.10.0 # via apache-superset (pyproject.toml) -trio==0.33.0 - # via - # selenium - # trio-websocket -trio-websocket==0.12.2 - # via selenium typing-extensions==4.16.0 # via # apache-superset (pyproject.toml) @@ -443,7 +418,6 @@ typing-extensions==4.16.0 # pydantic-core # pyopenssl # referencing - # selenium # shillelagh # typing-inspection typing-inspection==0.4.2 @@ -461,7 +435,6 @@ urllib3==2.7.0 # -r requirements/base.in # requests # requests-cache - # selenium vine==5.1.0 # via # amqp @@ -471,8 +444,6 @@ watchdog==6.0.0 # via apache-superset (pyproject.toml) wcwidth==0.2.13 # via prompt-toolkit -websocket-client==1.8.0 - # via selenium werkzeug==3.1.6 # via # -r requirements/base.in @@ -483,8 +454,6 @@ werkzeug==3.1.6 # flask-login wrapt==1.17.2 # via deprecated -wsproto==1.2.0 - # via trio-websocket wtforms==3.2.2 # via # apache-superset (pyproject.toml) diff --git a/requirements/development.txt b/requirements/development.txt index f43ebf95070..1d01361ac4c 100644 --- a/requirements/development.txt +++ b/requirements/development.txt @@ -48,10 +48,8 @@ attrs==25.3.0 # cattrs # cyclopts # jsonschema - # outcome # referencing # requests-cache - # trio authlib==1.6.12 # via fastmcp-slim babel==2.17.0 @@ -122,7 +120,6 @@ certifi==2026.5.20 # httpcore # httpx # requests - # selenium cffi==2.0.0 # via # -c requirements/base-constraint.txt @@ -400,10 +397,8 @@ gunicorn==26.0.0 # apache-superset h11==0.16.0 # via - # -c requirements/base-constraint.txt # httpcore # uvicorn - # wsproto hashids==1.3.1 # via # -c requirements/base-constraint.txt @@ -434,7 +429,6 @@ idna==3.15 # email-validator # httpx # requests - # trio # url-normalize importlib-metadata==8.7.0 # via @@ -619,11 +613,6 @@ ordered-set==4.1.0 # flask-limiter orjson==3.11.9 # via trino -outcome==1.3.0.post0 - # via - # -c requirements/base-constraint.txt - # trio - # trio-websocket packaging==25.0 # via # -c requirements/base-constraint.txt @@ -803,10 +792,6 @@ pyparsing==3.3.2 # matplotlib pyperclip==1.10.0 # via fastmcp-slim -pysocks==1.7.1 - # via - # -c requirements/base-constraint.txt - # urllib3 pytest==7.4.4 # via # apache-superset @@ -934,10 +919,6 @@ s3transfer==0.16.0 # via boto3 secretstorage==3.5.0 # via keyring -selenium==4.46.0 - # via - # -c requirements/base-constraint.txt - # apache-superset semver==3.0.4 # via apache-superset-extensions-cli setuptools==80.9.0 @@ -968,14 +949,7 @@ slack-sdk==3.43.0 # -c requirements/base-constraint.txt # apache-superset sniffio==1.3.1 - # via - # -c requirements/base-constraint.txt - # anyio - # trio -sortedcontainers==2.4.0 - # via - # -c requirements/base-constraint.txt - # trio + # via anyio sqlalchemy==1.4.54 # via # -c requirements/base-constraint.txt @@ -1039,15 +1013,6 @@ tqdm==4.67.1 # prophet trino==0.338.0 # via apache-superset -trio==0.33.0 - # via - # -c requirements/base-constraint.txt - # selenium - # trio-websocket -trio-websocket==0.12.2 - # via - # -c requirements/base-constraint.txt - # selenium typing-extensions==4.16.0 # via # -c requirements/base-constraint.txt @@ -1067,7 +1032,6 @@ typing-extensions==4.16.0 # pydantic-core # pyopenssl # referencing - # selenium # shillelagh # starlette # typing-inspection @@ -1100,7 +1064,6 @@ urllib3==2.7.0 # docker # requests # requests-cache - # selenium uvicorn==0.37.0 # via # fastmcp-slim @@ -1124,10 +1087,6 @@ wcwidth==0.2.13 # via # -c requirements/base-constraint.txt # prompt-toolkit -websocket-client==1.8.0 - # via - # -c requirements/base-constraint.txt - # selenium websockets==15.0.1 # via fastmcp-slim werkzeug==3.1.6 @@ -1142,10 +1101,6 @@ wrapt==1.17.2 # via # -c requirements/base-constraint.txt # deprecated -wsproto==1.2.0 - # via - # -c requirements/base-constraint.txt - # trio-websocket wtforms==3.2.2 # via # -c requirements/base-constraint.txt diff --git a/superset/config.py b/superset/config.py index 34f27cd3e7e..41a2d1cdebb 100644 --- a/superset/config.py +++ b/superset/config.py @@ -795,10 +795,6 @@ DEFAULT_FEATURE_FLAGS: dict[str, bool] = { # When impersonating a user, use the email prefix instead of username # @lifecycle: testing "IMPERSONATE_WITH_EMAIL_PREFIX": False, - # Replace Selenium with Playwright for reports and thumbnails. - # Supports deck.gl visualizations. Requires playwright pip package. - # @lifecycle: testing - "PLAYWRIGHT_REPORTS_AND_THUMBNAILS": False, # Apply RLS rules to SQL Lab queries. Requires query parsing/manipulation. # May break queries or allow RLS bypass. Use with care! # @lifecycle: testing @@ -1311,24 +1307,12 @@ SUPERSET_CACHE_WARMUP_USER: str | None = None SCREENSHOT_LOCATE_WAIT = int(timedelta(seconds=10).total_seconds()) # Time before screenshot capture times out while waiting for chart readiness. SCREENSHOT_LOAD_WAIT = int(timedelta(minutes=1).total_seconds()) -# Maximum time (in seconds) selenium waits for an initial page navigation -# (driver.get) to complete. Without it the navigation blocks indefinitely when -# the target page never finishes loading (e.g. an unreachable WEBDRIVER_BASEURL), -# which leaves the report schedule stuck in the WORKING state. Set to None to -# disable (not recommended). -SCREENSHOT_PAGE_LOAD_WAIT = int(timedelta(minutes=2).total_seconds()) -# Selenium destroy retries -SCREENSHOT_SELENIUM_RETRIES = 5 -# Give selenium an headstart, in seconds +# Give the browser an initial headstart, in seconds SCREENSHOT_SELENIUM_HEADSTART = 3 # Wait for the chart animation, in seconds SCREENSHOT_SELENIUM_ANIMATION_WAIT = 5 # Replace unexpected errors in screenshots with real error messages SCREENSHOT_REPLACE_UNEXPECTED_ERRORS = False -# Max time to wait for error message modal to show up, in seconds -SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE = 5 -# Max time to wait for error message modal to close, in seconds -SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE = 5 # Event that Playwright waits for when loading a new page # Possible values: "load", "commit", "domcontentloaded", "networkidle" # Docs: https://playwright.dev/python/docs/api/class-page#page-goto-option-wait-until @@ -2604,18 +2588,6 @@ SLACK_API_RATE_LIMIT_RETRY_COUNT = 2 # patching code, consistent with the SMTP/CSV/screenshot timeouts. SLACK_API_TIMEOUT = 30 -# The webdriver to use for generating reports when using Selenium (not Playwright). -# This setting is ignored when PLAYWRIGHT_REPORTS_AND_THUMBNAILS is enabled, as -# Playwright always uses Chromium regardless of this value. -# Use one of the following: -# firefox -# Requires: geckodriver and firefox installations -# Limitations: can be buggy at times -# chrome: -# Requires: headless chrome -# Limitations: unable to generate screenshots of elements -WEBDRIVER_TYPE = "firefox" - # Window size - this will impact the rendering of the data WEBDRIVER_WINDOW = { "dashboard": (1600, 2000), @@ -2623,20 +2595,12 @@ WEBDRIVER_WINDOW = { "pixel_density": 1, } -# An optional override to the default auth hook used to provide auth to the offline -# webdriver (when using Selenium) or browser context (when using Playwright - see -# PLAYWRIGHT_REPORTS_AND_THUMBNAILS feature flag) +# An optional override to the default auth hook used to provide auth to the +# browser context when using Playwright WEBDRIVER_AUTH_FUNC = None -# Any config options to be passed as-is to the webdriver -WEBDRIVER_CONFIGURATION = { - "options": {"capabilities": {}, "preferences": {}, "binary_location": ""}, - "service": {"log_output": "/dev/null", "service_args": [], "port": 0, "env": {}}, -} - -# Additional args to be passed as arguments to the config object -# Note: If using Chrome, you'll want to add the "--marionette" arg. -WEBDRIVER_OPTION_ARGS = ["--headless"] +# Additional args to be passed to the Playwright browser launch. +WEBDRIVER_OPTION_ARGS: list[str] = [] # The base URL to query for accessing the user interface WEBDRIVER_BASEURL = "http://0.0.0.0:8080/" diff --git a/superset/mcp_service/screenshot/__init__.py b/superset/mcp_service/screenshot/__init__.py index 3531e6e2c58..fd3d127ce55 100644 --- a/superset/mcp_service/screenshot/__init__.py +++ b/superset/mcp_service/screenshot/__init__.py @@ -15,21 +15,4 @@ # specific language governing permissions and limitations # under the License. -"""Screenshot and WebDriver infrastructure for MCP service.""" - -from .pooled_screenshot import ( - PooledBaseScreenshot, - PooledChartScreenshot, - PooledDashboardScreenshot, - PooledExploreScreenshot, -) -from .webdriver_pool import get_webdriver_pool, WebDriverPool - -__all__ = [ - "PooledBaseScreenshot", - "PooledChartScreenshot", - "PooledDashboardScreenshot", - "PooledExploreScreenshot", - "WebDriverPool", - "get_webdriver_pool", -] +"""Screenshot infrastructure for MCP service.""" diff --git a/superset/mcp_service/screenshot/pooled_screenshot.py b/superset/mcp_service/screenshot/pooled_screenshot.py deleted file mode 100644 index 136806b4293..00000000000 --- a/superset/mcp_service/screenshot/pooled_screenshot.py +++ /dev/null @@ -1,492 +0,0 @@ -# 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. - -""" -Pooled screenshot implementations for improved performance -""" - -import logging -import time -from typing import Any - -from flask import current_app -from flask_appbuilder.security.sqla.models import User -from selenium.common.exceptions import TimeoutException, WebDriverException -from selenium.webdriver.common.by import By -from selenium.webdriver.support import expected_conditions -from selenium.webdriver.support.ui import WebDriverWait - -from superset.extensions import machine_auth_provider_factory -from superset.mcp_service.screenshot.webdriver_pool import get_webdriver_pool -from superset.mcp_service.utils.retry_utils import retry_screenshot_operation -from superset.utils.report_execution import ReportExecutionContext -from superset.utils.screenshots import BaseScreenshot, WindowSize - -logger = logging.getLogger(__name__) - - -class PooledBaseScreenshot(BaseScreenshot): - """ - Enhanced screenshot class that uses WebDriver pooling for better performance. - - Key improvements: - - Reuses WebDriver instances from a connection pool - - Faster screenshot generation (no browser startup/shutdown per request) - - Better resource management and automatic cleanup - - Health checking and recovery of WebDriver instances - """ - - def get_screenshot( - self, - user: User, - window_size: WindowSize | None = None, - log_context: str | None = None, - report_execution_context: ReportExecutionContext | None = None, - ) -> bytes | None: - """ - Generate screenshot using pooled WebDriver with retry logic for reliability. - - Args: - user: User context for authentication - window_size: Optional window size override - log_context: Accepted for signature compatibility with - BaseScreenshot; the pooled Selenium path does not emit the - per-tile readiness logs that use it. - report_execution_context: Accepted for BaseScreenshot compatibility. - - Returns: - Screenshot as PNG bytes or None if failed - """ - return retry_screenshot_operation( - self._get_screenshot_internal, user, window_size - ) - - def _get_screenshot_internal( - self, user: User, window_size: WindowSize | None = None - ) -> bytes | None: - """ - Internal screenshot generation method with pooled WebDriver. - This method is wrapped by retry logic in get_screenshot(). - """ - window_size = window_size or self.window_size - pool = get_webdriver_pool() - - # Use pooled WebDriver - with pool.get_driver(window_size, user.id) as driver: - try: - # Clear any existing cookies to ensure clean authentication - try: - driver.delete_all_cookies() - logger.debug("Cleared all cookies from WebDriver") - except Exception as e: - logger.warning("Failed to clear cookies: %s", e) - - # Authenticate the driver for this user - user_name = user.username if user else "None" - logger.debug("Authenticating WebDriver for user %s", user_name) - machine_auth_provider_factory.instance.authenticate_webdriver( - driver, user - ) - - # Navigate to the URL - logger.debug("Navigating to screenshot URL: %s", self.url) - driver.get(self.url) - - # Check if we were redirected to login (authentication failed) - current_url = driver.current_url - if "/login" in current_url: - msg = ( - f"Authentication failed - WebDriver redirected to login page: " - f"{current_url}" - ) - raise Exception(msg) - - # Take screenshot using the specific implementation - return self._take_screenshot(driver, user) - - except (TimeoutException, WebDriverException, OSError) as e: - # These are retryable exceptions - logger.warning("Retryable error taking screenshot: %s", e) - raise - except Exception as e: - logger.error("Non-retryable error taking screenshot: %s", e) - raise - - def _take_screenshot(self, driver: Any, user: User) -> bytes | None: - """ - Override this method in subclasses to implement specific screenshot logic. - - Args: - driver: Authenticated WebDriver instance - user: User context - - Returns: - Screenshot as PNG bytes - """ - raise NotImplementedError("Subclasses must implement _take_screenshot") - - -class PooledChartScreenshot(PooledBaseScreenshot): - """Pooled version of chart screenshot generation""" - - thumbnail_type: str = "chart" - element: str = "chart-container" - - def __init__( - self, - url: str, - digest: str | None, - window_size: WindowSize | None = None, - thumb_size: WindowSize | None = None, - ): - from superset.utils.urls import modify_url_query - from superset.utils.webdriver import ChartStandaloneMode - - # Chart reports are in standalone="true" mode - url = modify_url_query( - url, - standalone=ChartStandaloneMode.HIDE_NAV.value, - ) - super().__init__(url, digest) - self.window_size = window_size or (800, 600) - self.thumb_size = thumb_size or (400, 300) - - def _take_screenshot(self, driver: Any, user: User) -> bytes | None: - """Take screenshot of chart with standard Superset chart handling""" - try: - # Wait for page to load - selenium_headstart = current_app.config["SCREENSHOT_SELENIUM_HEADSTART"] - logger.debug("Sleeping for %s seconds", selenium_headstart) - time.sleep(selenium_headstart) - - # Wait for chart container - wait = WebDriverWait(driver, current_app.config["SCREENSHOT_LOCATE_WAIT"]) - element = wait.until( - expected_conditions.presence_of_element_located( - (By.CLASS_NAME, self.element) - ) - ) - - # Wait for chart containers to render - wait.until( - expected_conditions.visibility_of_all_elements_located( - (By.CLASS_NAME, "chart-container") - ) - ) - - # Wait for loading to complete - WebDriverWait(driver, current_app.config["SCREENSHOT_LOAD_WAIT"]).until_not( - expected_conditions.presence_of_all_elements_located( - (By.CLASS_NAME, "loading") - ) - ) - - # Wait for animations - animation_wait = current_app.config["SCREENSHOT_SELENIUM_ANIMATION_WAIT"] - logger.debug("Wait %s seconds for chart animation", animation_wait) - time.sleep(animation_wait) - - # Handle unexpected errors if configured - if current_app.config.get("SCREENSHOT_REPLACE_UNEXPECTED_ERRORS"): - from superset.utils.webdriver import WebDriverSelenium - - unexpected_errors = WebDriverSelenium.find_unexpected_errors(driver) - if unexpected_errors: - logger.warning( - "%s errors found in screenshot. ", len(unexpected_errors) - ) - - # Take screenshot - logger.debug("Taking PNG screenshot as user %s", user.username) - return element.screenshot_as_png - - except TimeoutException: - logger.exception("Timeout taking chart screenshot for URL: %s", self.url) - raise - except WebDriverException: - logger.exception("WebDriver error taking screenshot for URL: %s", self.url) - raise - - -class PooledExploreScreenshot(PooledBaseScreenshot): - """ - Pooled version of explore screenshot with UI hiding functionality. - - This class provides the same clean chart-only screenshots as the previous - implementation but with improved performance through WebDriver pooling. - """ - - thumbnail_type: str = "explore" - element: str = "chart-container" - - def __init__( - self, - url: str, - digest: str | None, - window_size: WindowSize | None = None, - thumb_size: WindowSize | None = None, - ): - super().__init__(url, digest) - self.window_size = window_size or (1600, 1200) - self.thumb_size = thumb_size or (800, 600) - - def _take_screenshot(self, driver: Any, user: User) -> bytes | None: - """ - Take screenshot of explore page with UI elements hidden for clean chart display. - - This implementation: - 1. Waits for chart to load - 2. Hides navigation, panels, and headers using JavaScript - 3. Takes screenshot of just the chart area - 4. Includes fallback strategies if chart container not found - """ - try: - # Give explore page time to initialize - logger.debug("Waiting for explore page to initialize") - time.sleep(3) - - # Wait for chart container to load and be visible - wait = WebDriverWait(driver, 45) - chart_element = wait.until( - expected_conditions.visibility_of_element_located( - (By.CLASS_NAME, "chart-container") - ) - ) - - # Wait for any loading indicators to disappear - WebDriverWait(driver, 30).until_not( - expected_conditions.presence_of_all_elements_located( - (By.CLASS_NAME, "loading") - ) - ) - - # Hide overlapping UI elements using JavaScript - self._hide_ui_elements(driver) - - # Wait for UI hiding animations to complete - time.sleep(2) - - # Take screenshot of just the chart area - img = chart_element.screenshot_as_png - logger.info("Successfully captured chart-container screenshot") - return img - - except TimeoutException as e: - logger.warning("Chart container not found, trying fallbacks: %s", e) - return self._fallback_screenshot(driver) - except Exception as e: - logger.error("Error taking explore screenshot: %s", e) - raise - - def _hide_ui_elements(self, driver: Any) -> None: - """Hide UI elements to show only the chart""" - hide_script = """ - // Hide the specific left panels by their exact class names - var dataSourcePanel = document.querySelector( - '.explore-column.data-source-selection' - ); - if (dataSourcePanel) { - dataSourcePanel.style.display = 'none'; - console.log('Hidden data source panel'); - } - - var controlsPanel = document.querySelector( - '.col-sm-3.explore-column.controls-column' - ); - if (controlsPanel) { - controlsPanel.style.display = 'none'; - console.log('Hidden controls panel'); - } - - // Also try alternate selectors for the controls - var controlsAlt = document.querySelector('.explore-column.controls-column'); - if (controlsAlt) { - controlsAlt.style.display = 'none'; - console.log('Hidden controls panel (alt)'); - } - - // Hide the main navigation header - var mainHeader = document.querySelector('header.top#main-menu'); - if (mainHeader) { - mainHeader.style.display = 'none'; - console.log('Hidden main navigation header'); - } - - // Hide the chart header with actions (title input and save button) - var chartHeader = document.querySelector('.header-with-actions'); - if (chartHeader) { - chartHeader.style.display = 'none'; - console.log('Hidden chart header with actions'); - } - - // Make sure the main chart area expands to use available space - var mainContent = document.querySelector('.main-explore-content'); - if (mainContent) { - mainContent.style.width = '100%'; - // Change from col-sm-7 to col-sm-12 - mainContent.className = 'main-explore-content col-sm-12'; - } - - // Ensure chart container is visible - var chartContainer = document.querySelector('.chart-container'); - if (chartContainer) { - chartContainer.style.position = 'relative'; - chartContainer.style.zIndex = '999'; - chartContainer.style.width = '100%'; - } - """ - - try: - driver.execute_script(hide_script) - logger.debug("Successfully executed UI hiding script") - except WebDriverException as e: - logger.warning("Failed to execute UI hiding script: %s", e) - - def _fallback_screenshot(self, driver: Any) -> bytes | None: - """Fallback screenshot strategies if chart container not found""" - try: - # Fallback 1: Try slice container - logger.debug("Trying slice_container fallback") - slice_element = WebDriverWait(driver, 10).until( - expected_conditions.visibility_of_element_located( - (By.CLASS_NAME, "slice_container") - ) - ) - img = slice_element.screenshot_as_png - logger.info("Successfully captured slice_container screenshot") - return img - - except TimeoutException: - try: - # Fallback 2: Try any chart-related container - logger.debug("Trying chart-related container fallback") - chart_containers = driver.find_elements( - By.CSS_SELECTOR, - ".chart, .slice-container, .chart-content, " - "[data-test='chart-container']", - ) - if chart_containers: - img = chart_containers[0].screenshot_as_png - logger.info( - "Successfully captured chart-related container screenshot" - ) - return img - - except Exception as e2: - logger.warning("Chart-related container fallback failed: %s", e2) - - # Fallback 3: Full page screenshot - logger.warning("Using full page screenshot as final fallback") - img = driver.get_screenshot_as_png() - logger.info("Using full page screenshot as fallback") - return img - - -class PooledDashboardScreenshot(PooledBaseScreenshot): - """Pooled version of dashboard screenshot generation""" - - thumbnail_type: str = "dashboard" - element: str = "standalone" - - def __init__( - self, - url: str, - digest: str | None, - window_size: WindowSize | None = None, - thumb_size: WindowSize | None = None, - ): - from superset.utils.urls import modify_url_query - from superset.utils.webdriver import DashboardStandaloneMode - - # Dashboard screenshots should always capture in standalone - url = modify_url_query( - url, - standalone=DashboardStandaloneMode.REPORT.value, - ) - super().__init__(url, digest) - self.window_size = window_size or (1600, 1200) - self.thumb_size = thumb_size or (800, 600) - - def _take_screenshot(self, driver: Any, user: User) -> bytes | None: - """Take screenshot of dashboard with standard Superset dashboard handling""" - try: - # Wait for page to load - selenium_headstart = current_app.config["SCREENSHOT_SELENIUM_HEADSTART"] - time.sleep(selenium_headstart) - - # Wait for dashboard element - wait = WebDriverWait(driver, current_app.config["SCREENSHOT_LOCATE_WAIT"]) - element = wait.until( - expected_conditions.presence_of_element_located( - (By.CLASS_NAME, self.element) - ) - ) - - # Wait for chart containers to render - try: - wait.until( - expected_conditions.visibility_of_all_elements_located( - (By.CLASS_NAME, "chart-container") - ) - ) - except TimeoutException: - # Fallback for empty dashboards - try: - WebDriverWait(driver, 0).until( - expected_conditions.visibility_of_all_elements_located( - (By.CLASS_NAME, "grid-container") - ) - ) - except TimeoutException: - logger.exception("Dashboard failed to load at URL: %s", self.url) - raise - - # Wait for loading to complete - WebDriverWait(driver, current_app.config["SCREENSHOT_LOAD_WAIT"]).until_not( - expected_conditions.presence_of_all_elements_located( - (By.CLASS_NAME, "loading") - ) - ) - - # Wait for animations - animation_wait = current_app.config["SCREENSHOT_SELENIUM_ANIMATION_WAIT"] - time.sleep(animation_wait) - - # Handle unexpected errors if configured - if current_app.config.get("SCREENSHOT_REPLACE_UNEXPECTED_ERRORS"): - from superset.utils.webdriver import WebDriverSelenium - - unexpected_errors = WebDriverSelenium.find_unexpected_errors(driver) - if unexpected_errors: - logger.warning( - "%s errors found in dashboard ", len(unexpected_errors) - ) - - # Take screenshot - logger.debug("Taking PNG dashboard screenshot as user %s", user.username) - return element.screenshot_as_png - - except TimeoutException: - logger.exception( - "Timeout taking dashboard screenshot for URL: %s", self.url - ) - raise - except WebDriverException: - logger.exception( - "WebDriver error taking dashboard screenshot for URL: %s", self.url - ) - raise diff --git a/superset/mcp_service/screenshot/webdriver_config.py b/superset/mcp_service/screenshot/webdriver_config.py deleted file mode 100644 index 997e355d7c6..00000000000 --- a/superset/mcp_service/screenshot/webdriver_config.py +++ /dev/null @@ -1,143 +0,0 @@ -# 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. - -""" -WebDriver pool configuration defaults for Superset MCP service -""" - -import logging -from typing import Any, Dict - -logger = logging.getLogger(__name__) - -# Default WebDriver pool configuration -DEFAULT_WEBDRIVER_POOL_CONFIG = { - # Maximum number of WebDriver instances to keep in the pool - "MAX_POOL_SIZE": 5, - # Maximum age of a WebDriver instance (in seconds) - # After this time, the driver will be destroyed and recreated - "MAX_AGE_SECONDS": 3600, # 1 hour - # Maximum number of times a WebDriver can be reused - # After this many uses, the driver will be destroyed and recreated - "MAX_USAGE_COUNT": 50, - # How long a WebDriver can sit idle before being destroyed (in seconds) - "IDLE_TIMEOUT_SECONDS": 300, # 5 minutes - # How often to perform health checks on WebDriver instances (in seconds) - "HEALTH_CHECK_INTERVAL": 60, # 1 minute -} - - -def configure_webdriver_pool(app_config: Dict[str, Any]) -> None: - """ - Configure WebDriver pool settings in Superset app config. - - This function adds WebDriver pool configuration to the Superset app config - if it doesn't already exist, using sensible defaults. - - Args: - app_config: The Superset application configuration dictionary - """ - if "WEBDRIVER_POOL" not in app_config: - app_config["WEBDRIVER_POOL"] = DEFAULT_WEBDRIVER_POOL_CONFIG.copy() - else: - # Merge with defaults for any missing keys - for key, default_value in DEFAULT_WEBDRIVER_POOL_CONFIG.items(): - if key not in app_config["WEBDRIVER_POOL"]: - app_config["WEBDRIVER_POOL"][key] = default_value - - -def get_pool_stats_endpoint() -> Any: - """ - Create a Flask endpoint to view WebDriver pool statistics. - - This function can be called to register a debugging endpoint - that shows the current state of the WebDriver pool. - - Returns: - Flask route function for pool statistics - """ - - def pool_stats() -> Any: - from flask import jsonify - - try: - from superset.mcp_service.screenshot.webdriver_pool import ( - get_webdriver_pool, - ) - - pool = get_webdriver_pool() - stats = pool.get_stats() - - return jsonify({"webdriver_pool": stats, "status": "healthy"}) - except Exception: - logger.exception("Failed to retrieve webdriver pool stats") - return jsonify( - {"error": "Failed to retrieve pool stats", "status": "error"} - ), 500 - - return pool_stats - - -# Performance tuning recommendations based on use case -PERFORMANCE_CONFIGS = { - "low_traffic": { - "MAX_POOL_SIZE": 2, - "MAX_AGE_SECONDS": 1800, # 30 minutes - "MAX_USAGE_COUNT": 25, - "IDLE_TIMEOUT_SECONDS": 180, # 3 minutes - "HEALTH_CHECK_INTERVAL": 120, # 2 minutes - }, - "medium_traffic": { - "MAX_POOL_SIZE": 5, - "MAX_AGE_SECONDS": 3600, # 1 hour - "MAX_USAGE_COUNT": 50, - "IDLE_TIMEOUT_SECONDS": 300, # 5 minutes - "HEALTH_CHECK_INTERVAL": 60, # 1 minute - }, - "high_traffic": { - "MAX_POOL_SIZE": 10, - "MAX_AGE_SECONDS": 7200, # 2 hours - "MAX_USAGE_COUNT": 100, - "IDLE_TIMEOUT_SECONDS": 600, # 10 minutes - "HEALTH_CHECK_INTERVAL": 30, # 30 seconds - }, - "development": { - "MAX_POOL_SIZE": 2, - "MAX_AGE_SECONDS": 900, # 15 minutes - "MAX_USAGE_COUNT": 10, - "IDLE_TIMEOUT_SECONDS": 120, # 2 minutes - "HEALTH_CHECK_INTERVAL": 30, # 30 seconds - }, -} - - -def configure_for_environment( - app_config: Dict[str, Any], environment: str = "medium_traffic" -) -> None: - """ - Configure WebDriver pool for specific environment/traffic levels. - - Args: - app_config: The Superset application configuration dictionary - environment: Environment type (low_traffic, medium_traffic, - high_traffic, development) - """ - if environment in PERFORMANCE_CONFIGS: - app_config["WEBDRIVER_POOL"] = PERFORMANCE_CONFIGS[environment].copy() - else: - # Fallback to default - configure_webdriver_pool(app_config) diff --git a/superset/mcp_service/screenshot/webdriver_pool.py b/superset/mcp_service/screenshot/webdriver_pool.py deleted file mode 100644 index ba1e2c887ed..00000000000 --- a/superset/mcp_service/screenshot/webdriver_pool.py +++ /dev/null @@ -1,437 +0,0 @@ -# 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. - -""" -WebDriver connection pooling for improved screenshot performance -""" - -import logging -import signal -import threading -import time -from contextlib import contextmanager -from dataclasses import dataclass -from queue import Empty, Full, Queue -from typing import Any, Dict, Generator - -from flask import current_app -from selenium.common.exceptions import WebDriverException -from selenium.webdriver.remote.webdriver import WebDriver - -from superset.utils.webdriver import WebDriverSelenium, WindowSize - -logger = logging.getLogger(__name__) - - -class WebDriverCreationError(Exception): - """Exception raised when WebDriver creation times out""" - - pass - - -def _timeout_handler(signum: int, frame: Any) -> None: - """Signal handler for WebDriver creation timeout""" - raise WebDriverCreationError("WebDriver creation timed out") - - -@dataclass -class PooledWebDriver: - """Wrapper for pooled WebDriver instance with metadata""" - - driver: WebDriver - created_at: float - last_used: float - window_size: WindowSize - user_id: int | None = None - is_healthy: bool = True - usage_count: int = 0 - - -class WebDriverPool: - """ - Connection pool for WebDriver instances to improve screenshot performance. - - Features: - - Reuses WebDriver instances across requests - - Automatic health checking and recovery - - TTL-based expiration to prevent memory leaks - - Thread-safe operations - - Per-user driver isolation for security - """ - - def __init__( - self, - max_pool_size: int = 5, - max_age_seconds: int = 3600, # 1 hour - max_usage_count: int = 50, # Recreate after 50 uses - idle_timeout_seconds: int = 300, # 5 minutes - health_check_interval: int = 60, # 1 minute - creation_timeout_seconds: int = 30, # SECURITY FIX: Timeout for driver creation - ): - self.max_pool_size = max_pool_size - self.max_age_seconds = max_age_seconds - self.max_usage_count = max_usage_count - self.idle_timeout_seconds = idle_timeout_seconds - self.health_check_interval = health_check_interval - self.creation_timeout_seconds = creation_timeout_seconds - - # Thread-safe pool management - self._pool: Queue[PooledWebDriver] = Queue(maxsize=max_pool_size) - self._active_drivers: Dict[int, PooledWebDriver] = {} - self._lock = threading.RLock() - self._last_health_check = time.time() - - # Pool statistics - self._stats = { - "created": 0, - "destroyed": 0, - "borrowed": 0, - "returned": 0, - "health_check_failures": 0, - "evictions": 0, - } - - def get_stats(self) -> Dict[str, Any]: - """Get pool statistics for monitoring""" - with self._lock: - return { - **self._stats, - "pool_size": self._pool.qsize(), - "active_count": len(self._active_drivers), - "max_pool_size": self.max_pool_size, - } - - def _create_driver( - self, window_size: WindowSize, user_id: int | None = None - ) -> PooledWebDriver: - """Create a new WebDriver instance with timeout protection""" - driver = None - old_handler = None - - try: - # SECURITY FIX: Set up timeout protection for driver creation - old_handler = signal.signal(signal.SIGALRM, _timeout_handler) - signal.alarm(self.creation_timeout_seconds) - - driver_type = current_app.config.get("WEBDRIVER_TYPE", "firefox") - selenium_driver = WebDriverSelenium(driver_type, window_size) - - # Create the actual WebDriver with timeout protection - driver = selenium_driver.create() - driver.set_window_size(*window_size) - - # Clear the alarm - creation successful - signal.alarm(0) - - pooled_driver = PooledWebDriver( - driver=driver, - created_at=time.time(), - last_used=time.time(), - window_size=window_size, - user_id=user_id, - is_healthy=True, - usage_count=0, - ) - - self._stats["created"] += 1 - logger.debug( - "Created new WebDriver instance for window size %s", window_size - ) - return pooled_driver - - except WebDriverCreationError: - logger.error( - "WebDriver creation timed out after %s seconds", - self.creation_timeout_seconds, - ) - if driver: - try: - driver.quit() - except Exception: - logger.debug("Failed to cleanup driver during timeout") - raise Exception("WebDriver creation timed out") from None - - except Exception as e: - logger.error("Failed to create WebDriver: %s", e) - if driver: - try: - driver.quit() - except Exception: - logger.debug("Failed to cleanup driver during error") - raise - - finally: - # Restore original signal handler and clear alarm - signal.alarm(0) - if old_handler is not None: - signal.signal(signal.SIGALRM, old_handler) - - def _is_driver_valid(self, pooled_driver: PooledWebDriver) -> bool: - """Check if a pooled driver is still valid for use""" - now = time.time() - - # Check age limit - if now - pooled_driver.created_at > self.max_age_seconds: - logger.debug("Driver expired due to age") - return False - - # Check usage count limit - if pooled_driver.usage_count >= self.max_usage_count: - logger.debug("Driver expired due to usage count") - return False - - # Check idle timeout - if now - pooled_driver.last_used > self.idle_timeout_seconds: - logger.debug("Driver expired due to idle timeout") - return False - - # Check if driver is healthy - if not pooled_driver.is_healthy: - logger.debug("Driver marked as unhealthy") - return False - - return True - - def _health_check_driver(self, pooled_driver: PooledWebDriver) -> bool: - """Perform health check on a WebDriver instance""" - try: - # Simple health check - try to get current URL - # This will fail if the driver is dead/hung - _ = pooled_driver.driver.current_url - pooled_driver.is_healthy = True - return True - except WebDriverException: - pooled_driver.is_healthy = False - self._stats["health_check_failures"] += 1 - logger.warning("WebDriver failed health check") - return False - except Exception as e: - pooled_driver.is_healthy = False - self._stats["health_check_failures"] += 1 - logger.warning("WebDriver health check error: %s", e) - return False - - def _destroy_driver(self, pooled_driver: PooledWebDriver) -> None: - """Safely destroy a WebDriver instance""" - try: - try: - pooled_driver.driver.close() - except Exception: # pylint: disable=broad-except # noqa: S110 - pass - pooled_driver.driver.quit() - self._stats["destroyed"] += 1 - logger.debug("Destroyed WebDriver instance") - except Exception as e: - logger.warning("Error destroying WebDriver: %s", e) - - def _cleanup_expired_drivers(self) -> None: - """Remove expired drivers from the pool""" - expired_drivers = [] - - # Check pool for expired drivers - while not self._pool.empty(): - try: - pooled_driver = self._pool.get_nowait() - if self._is_driver_valid(pooled_driver): - # Driver is still valid, put it back - self._pool.put_nowait(pooled_driver) - break - else: - # Driver is expired - expired_drivers.append(pooled_driver) - self._stats["evictions"] += 1 - except Empty: - break - except Full: - # Pool is full, stop checking - break - - # Destroy expired drivers - for pooled_driver in expired_drivers: - self._destroy_driver(pooled_driver) - - def _periodic_health_check(self) -> None: - """Perform periodic health checks if needed""" - now = time.time() - if now - self._last_health_check < self.health_check_interval: - return - - self._last_health_check = now - logger.debug("Performing periodic WebDriver pool health check") - - # Cleanup expired drivers - self._cleanup_expired_drivers() - - # Health check active drivers - unhealthy_drivers = [] - for driver_id, pooled_driver in self._active_drivers.items(): - if not self._health_check_driver(pooled_driver): - unhealthy_drivers.append(driver_id) - - # Remove unhealthy active drivers - for driver_id in unhealthy_drivers: - pooled_driver = self._active_drivers.pop(driver_id) - if pooled_driver: - self._destroy_driver(pooled_driver) - - @contextmanager - def get_driver( - self, window_size: WindowSize, user_id: int | None = None - ) -> Generator[WebDriver, None, None]: - """ - Context manager to get a WebDriver from the pool. - - Args: - window_size: Required window size for the driver - user_id: Optional user ID for driver isolation - - Yields: - WebDriver instance ready for use - """ - pooled_driver = None - driver_id = None - - try: - with self._lock: - # Periodic maintenance - self._periodic_health_check() - - # Try to get a driver from the pool - while not self._pool.empty(): - try: - candidate = self._pool.get_nowait() - - # Check if driver is valid and matches requirements - if ( - self._is_driver_valid(candidate) - and candidate.window_size == window_size - ): - # Update user_id for the reused driver - candidate.user_id = user_id - pooled_driver = candidate - break - else: - # Driver is invalid, destroy it - self._destroy_driver(candidate) - self._stats["evictions"] += 1 - except Empty: - break - - # If no suitable driver found, create a new one - if pooled_driver is None: - pooled_driver = self._create_driver(window_size, user_id) - - # Mark driver as in use - driver_id = id(pooled_driver.driver) - pooled_driver.last_used = time.time() - pooled_driver.usage_count += 1 - self._active_drivers[driver_id] = pooled_driver - self._stats["borrowed"] += 1 - - # Yield the driver for use - yield pooled_driver.driver - - except Exception as e: - # Mark driver as unhealthy if an error occurred - if pooled_driver: - pooled_driver.is_healthy = False - logger.error("Error using pooled WebDriver: %s", e) - raise - - finally: - # Return driver to pool or destroy if unhealthy - if pooled_driver and driver_id: - with self._lock: - self._active_drivers.pop(driver_id, None) - - if pooled_driver.is_healthy and self._is_driver_valid( - pooled_driver - ): - # Try to return to pool - try: - self._pool.put_nowait(pooled_driver) - self._stats["returned"] += 1 - logger.debug("Returned WebDriver to pool") - except Full: - # Pool is full, destroy the driver - self._destroy_driver(pooled_driver) - logger.debug("Pool full, destroyed WebDriver") - else: - # Driver is unhealthy or expired, destroy it - self._destroy_driver(pooled_driver) - logger.debug("Destroyed unhealthy/expired WebDriver") - - def shutdown(self) -> None: - """Shutdown the pool and destroy all drivers""" - with self._lock: - logger.info("Shutting down WebDriver pool") - - # Destroy all active drivers - for pooled_driver in self._active_drivers.values(): - self._destroy_driver(pooled_driver) - self._active_drivers.clear() - - # Destroy all pooled drivers - while not self._pool.empty(): - try: - pooled_driver = self._pool.get_nowait() - self._destroy_driver(pooled_driver) - except Empty: - break - - logger.info( - "WebDriver pool shutdown complete. Final stats: %s", self.get_stats() - ) - - -# Global pool instance -_global_pool: WebDriverPool | None = None -_pool_lock = threading.Lock() - - -def get_webdriver_pool() -> WebDriverPool: - """Get or create the global WebDriver pool""" - global _global_pool - - if _global_pool is None: - with _pool_lock: - if _global_pool is None: - # Get pool configuration from Flask config - config = current_app.config - pool_config = config.get("WEBDRIVER_POOL", {}) - - _global_pool = WebDriverPool( - max_pool_size=pool_config.get("MAX_POOL_SIZE", 5), - max_age_seconds=pool_config.get("MAX_AGE_SECONDS", 3600), - max_usage_count=pool_config.get("MAX_USAGE_COUNT", 50), - idle_timeout_seconds=pool_config.get("IDLE_TIMEOUT_SECONDS", 300), - health_check_interval=pool_config.get("HEALTH_CHECK_INTERVAL", 60), - ) - logger.info("Initialized global WebDriver pool") - - return _global_pool - - -def shutdown_webdriver_pool() -> None: - """Shutdown the global WebDriver pool""" - global _global_pool - - if _global_pool is not None: - with _pool_lock: - if _global_pool is not None: - _global_pool.shutdown() - _global_pool = None diff --git a/superset/tasks/cache.py b/superset/tasks/cache.py index c994e735dc3..c2a6d4e6d3b 100644 --- a/superset/tasks/cache.py +++ b/superset/tasks/cache.py @@ -38,7 +38,7 @@ from superset.tasks.native_filter_cache import ( get_eligible_native_filters, ) from superset.utils.date_parser import parse_human_datetime -from superset.utils.webdriver import WebDriverSelenium +from superset.utils.webdriver import _browser_manager, WebDriverPlaywright logger: logging.Logger = get_task_logger(__name__) logger.setLevel(logging.INFO) @@ -327,6 +327,28 @@ strategy_registry: dict[str, type[Strategy]] = { } +def _warmup_urls( + urls: list[str], + user: Any, + window: tuple[int, int], + results: dict[str, list[str]], +) -> None: + wd: WebDriverPlaywright = WebDriverPlaywright("", window) + try: + for url in urls: + try: + logger.info("Fetching %s", url) + screenshot = wd.get_screenshot(url, "grid-container", user=user) + if not screenshot: + raise RuntimeError("No screenshot captured") + results["success"].append(url) + except Exception: # noqa: BLE001 + logger.exception("Error warming up cache for %s", url) + results["errors"].append(url) + finally: + _browser_manager._cleanup() # noqa: SLF001 + + @celery_app.task(name="cache-warmup") def cache_warmup( strategy_name: str, *args: Any, **kwargs: Any @@ -397,21 +419,10 @@ def cache_warmup( return results - wd: WebDriverSelenium = WebDriverSelenium( - current_app.config["WEBDRIVER_TYPE"], user=user + _warmup_urls( + strategy.get_urls(), + user, + current_app.config["WEBDRIVER_WINDOW"]["dashboard"], + results, ) - - try: - for url in strategy.get_urls(): - try: - logger.info("Fetching %s", url) - wd.get_screenshot(url, "grid-container") - results["success"].append(url) - except Exception: # noqa: BLE001 - logger.exception("Error warming up cache for %s", url) - results["errors"].append(url) - finally: - # Ensure WebDriver is properly cleaned up - wd.destroy() - return results diff --git a/superset/utils/machine_auth.py b/superset/utils/machine_auth.py index 9f9d2fe79b6..d417bc586d3 100644 --- a/superset/utils/machine_auth.py +++ b/superset/utils/machine_auth.py @@ -23,7 +23,6 @@ from urllib.parse import urlparse from flask import current_app as app, Flask, request, Response, session from flask_login import login_user -from selenium.webdriver.remote.webdriver import WebDriver from werkzeug.http import parse_cookie from superset.utils.class_utils import load_class_from_name @@ -43,40 +42,14 @@ if TYPE_CHECKING: class MachineAuthProvider: def __init__( self, - auth_webdriver_func_override: Callable[ - [WebDriver | BrowserContext, User], WebDriver | BrowserContext - ] + auth_webdriver_func_override: Callable[[BrowserContext, User], BrowserContext] | None = None, ): - # This is here in order to allow for the authenticate_webdriver - # or authenticate_browser_context (if PLAYWRIGHT_REPORTS_AND_THUMBNAILS is - # enabled) func to be overridden via config, as opposed to the entire + # This is here in order to allow for the authenticate_browser_context + # func to be overridden via config, as opposed to the entire # provider implementation self._auth_webdriver_func_override = auth_webdriver_func_override - def authenticate_webdriver( - self, - driver: WebDriver, - user: User, - ) -> WebDriver: - """ - Default AuthDriverFuncType type that sets a session cookie flask-login style - :return: The WebDriver passed in (fluent) - """ - # Short-circuit this method if we have an override configured - if self._auth_webdriver_func_override: - return self._auth_webdriver_func_override(driver, user) - - # Setting cookies requires doing a request first - driver.get(headless_url("/login/")) - - cookies = self.get_cookies(user) - - for cookie_name, cookie_val in cookies.items(): - driver.add_cookie({"name": cookie_name, "value": cookie_val}) - - return driver - def authenticate_browser_context( self, browser_context: BrowserContext, diff --git a/superset/utils/screenshots.py b/superset/utils/screenshots.py index 5ebfb85125d..4fcd58c9c85 100644 --- a/superset/utils/screenshots.py +++ b/superset/utils/screenshots.py @@ -25,7 +25,7 @@ from typing import cast, TYPE_CHECKING, TypedDict from flask import current_app as app -from superset import feature_flag_manager, thumbnail_cache +from superset import thumbnail_cache from superset.distributed_lock import DistributedLock from superset.exceptions import ( LockAlreadyHeldException, @@ -40,22 +40,11 @@ from superset.utils.webdriver import ( DashboardStandaloneMode, WebDriverPlaywright, WebDriverProxy, - WebDriverSelenium, WindowSize, ) logger = logging.getLogger(__name__) -# Import Playwright availability and install message -try: - from superset.utils.webdriver import ( - PLAYWRIGHT_AVAILABLE, - PLAYWRIGHT_INSTALL_MESSAGE, - ) -except ImportError: - PLAYWRIGHT_AVAILABLE = False - PLAYWRIGHT_INSTALL_MESSAGE = "Playwright module not found" - DEFAULT_SCREENSHOT_WINDOW_SIZE = 800, 600 DEFAULT_SCREENSHOT_THUMBNAIL_SIZE = 400, 300 @@ -199,10 +188,6 @@ class ScreenshotCachePayload: class BaseScreenshot: - @property - def driver_type(self) -> str: - return app.config["WEBDRIVER_TYPE"] - url: str digest: str | None screenshot: bytes | None @@ -220,27 +205,10 @@ class BaseScreenshot: def driver( self, window_size: WindowSize | None = None, - user: User | None = None, - log_context: str | None = None, ) -> WebDriverProxy: window_size = window_size or self.window_size - if feature_flag_manager.is_feature_enabled("PLAYWRIGHT_REPORTS_AND_THUMBNAILS"): - # Try to use Playwright if available (supports WebGL/DeckGL, unlike Cypress) - if PLAYWRIGHT_AVAILABLE: - return WebDriverPlaywright(self.driver_type, window_size) - - # Playwright not available, falling back to Selenium - context_suffix = f" [{log_context}]" if log_context else "" - logger.info( - "PLAYWRIGHT_REPORTS_AND_THUMBNAILS enabled but Playwright not " - "installed. Falling back to Selenium (WebGL/Canvas charts may " - "not render correctly). %s%s", - PLAYWRIGHT_INSTALL_MESSAGE, - context_suffix, - ) - - # Use Selenium as default/fallback - return WebDriverSelenium(self.driver_type, window_size, user) + # Empty string for driver_type — unused by WebDriverPlaywright internals + return WebDriverPlaywright("", window_size) def get_screenshot( self, @@ -249,18 +217,14 @@ class BaseScreenshot: log_context: str | None = None, report_execution_context: ReportExecutionContext | None = None, ) -> bytes | None: - driver = self.driver(window_size, user, log_context=log_context) - try: - self.screenshot = driver.get_screenshot( - self.url, - self.element, - user, - log_context=log_context, - report_execution_context=report_execution_context, - ) - finally: - if isinstance(driver, WebDriverSelenium): - driver.destroy() + driver = self.driver(window_size) + self.screenshot = driver.get_screenshot( + self.url, + self.element, + user, + log_context=log_context, + report_execution_context=report_execution_context, + ) return self.screenshot def get_cache_key( diff --git a/superset/utils/webdriver.py b/superset/utils/webdriver.py index c16e71f256b..18944937636 100644 --- a/superset/utils/webdriver.py +++ b/superset/utils/webdriver.py @@ -22,29 +22,14 @@ import logging import time from abc import ABC, abstractmethod from enum import Enum -from time import sleep from typing import Any, TYPE_CHECKING from flask import current_app as app -from packaging import version -from selenium import __version__ as selenium_version -from selenium.common.exceptions import ( - StaleElementReferenceException, - TimeoutException, - WebDriverException, -) -from selenium.webdriver import chrome, firefox, FirefoxProfile -from selenium.webdriver.common.by import By -from selenium.webdriver.common.service import Service -from selenium.webdriver.remote.webdriver import WebDriver -from selenium.webdriver.support import expected_conditions as EC # noqa: N812 -from selenium.webdriver.support.ui import WebDriverWait from superset.extensions import machine_auth_provider_factory from superset.utils.report_execution import ( ReportExecutionContext, ) -from superset.utils.retries import retry_call from superset.utils.screenshot_utils import ( CHART_CONTAINER_READY_JS, CHART_CONTAINER_STATE_JS, @@ -94,39 +79,8 @@ except ImportError: def check_playwright_availability() -> bool: - """ - Lightweight check for Playwright availability. - - First checks if browser binary exists, falls back to launch test if needed. - """ - if sync_playwright is None: - return False - - try: - with sync_playwright() as p: - # First try lightweight check - just verify executable exists - try: - executable_path = p.chromium.executable_path - if executable_path: - return True - except Exception: - # Fall back to full launch test if executable_path fails - logger.debug( - "Executable path check failed, falling back to launch test" - ) - - # Fallback: actually launch browser to ensure it works - browser = p.chromium.launch(headless=True) - browser.close() - return True - except Exception as e: - logger.warning( - "Playwright module is installed but browser launch failed. " - "Run 'playwright install chromium' to install browser binaries. " - "Error: %s", - str(e), - ) - return False + """Check Playwright availability by verifying the module is importable.""" + return sync_playwright is not None PLAYWRIGHT_AVAILABLE = check_playwright_availability() @@ -174,28 +128,6 @@ _browser_manager = _PlaywrightBrowserManager() atexit.register(_browser_manager._cleanup) -def validate_webdriver_config() -> dict[str, Any]: - """ - Validate webdriver configuration and dependencies. - - Used to check migration status from Cypress to Playwright. - Returns a dictionary with the status of available webdrivers - and feature flags. - """ - from superset import feature_flag_manager - - return { - "selenium_available": True, # Always available as required dependency - "playwright_available": PLAYWRIGHT_AVAILABLE, - "playwright_feature_enabled": feature_flag_manager.is_feature_enabled( - "PLAYWRIGHT_REPORTS_AND_THUMBNAILS" - ), - "recommended_action": ( - PLAYWRIGHT_INSTALL_MESSAGE if not PLAYWRIGHT_AVAILABLE else None - ), - } - - class DashboardStandaloneMode(Enum): HIDE_NAV = 1 HIDE_NAV_AND_TITLE = 2 @@ -584,17 +516,19 @@ class WebDriverPlaywright(WebDriverProxy): ) context_suffix = f" [{log_context}]" if log_context else "" if not PLAYWRIGHT_AVAILABLE: - logger.info( - "Playwright not available - falling back to Selenium. " - "Note: WebGL/Canvas charts may not render correctly with Selenium. " - "%s%s", - PLAYWRIGHT_INSTALL_MESSAGE, - context_suffix, + raise RuntimeError( + f"Playwright is required for screenshots. " + f"{PLAYWRIGHT_INSTALL_MESSAGE}{context_suffix}" ) - return None browser_args = app.config["WEBDRIVER_OPTION_ARGS"] - browser = _browser_manager.get_browser(browser_args) + try: + browser = _browser_manager.get_browser(browser_args) + except Exception as ex: + raise RuntimeError( + f"Playwright is required for screenshots. " + f"{PLAYWRIGHT_INSTALL_MESSAGE}{context_suffix}" + ) from ex pixel_density = app.config["WEBDRIVER_WINDOW"].get("pixel_density", 1) viewport_height = self._window[1] viewport_width = self._window[0] @@ -1014,581 +948,3 @@ class WebDriverPlaywright(WebDriverProxy): finally: context.close() return img - - -class WebDriverSelenium(WebDriverProxy): - def __init__( - self, - driver_type: str, - window: WindowSize | None = None, - user: User | None = None, - ): - super().__init__(driver_type, window) - self._user = user - self._driver: WebDriver | None = None - - def __del__(self) -> None: - self._destroy() - - @property - def driver(self) -> WebDriver: - if not self._driver: - self._driver = self._create() - if not self._driver: - raise RuntimeError("WebDriver creation failed") - try: - self._driver.set_window_size(*self._window) - # Bound driver.get() so an unreachable page raises a - # TimeoutException instead of blocking the worker (and the - # report schedule) forever. - page_load_wait = app.config["SCREENSHOT_PAGE_LOAD_WAIT"] - if page_load_wait is not None: - self._driver.set_page_load_timeout(page_load_wait) - if self._user: - self._auth(self._user) - except Exception: - # A failure mid-setup (e.g. the new page-load timeout or auth - # raising) would otherwise leave a partially initialized, - # unauthenticated driver cached for reuse. Tear it down so the - # next access recreates it cleanly. - self._destroy() - raise - return self._driver - - def _create_firefox_driver( - self, pixel_density: float - ) -> tuple[type[WebDriver], type[Service], dict[str, Any]]: - """Create Firefox driver configuration.""" - options = firefox.options.Options() - profile = FirefoxProfile() - profile.set_preference("layout.css.devPixelsPerPx", str(pixel_density)) - options.profile = profile - return ( - firefox.webdriver.WebDriver, - firefox.service.Service, - {"options": options}, - ) - - def _create_chrome_driver( - self, pixel_density: float - ) -> tuple[type[WebDriver], type[Service], dict[str, Any]]: - """Create Chrome driver configuration.""" - options = chrome.options.Options() - options.add_argument(f"--force-device-scale-factor={pixel_density}") - options.add_argument(f"--window-size={self._window[0]},{self._window[1]}") - return ( - chrome.webdriver.WebDriver, - chrome.service.Service, - {"options": options}, - ) - - def _normalize_timeout_values(self, config: dict[str, Any]) -> dict[str, Any]: - """Convert timeout values to float for urllib3 2.x compatibility.""" - timeout_keys = [ - "timeout", - "connect_timeout", - "socket_timeout", - "read_timeout", - "page_load_timeout", - "implicit_wait", - "command_executor_timeout", - "connection_timeout", - ] - - for key, value in config.items(): - if any(timeout_key in key.lower() for timeout_key in timeout_keys): - if value is None or value == "None" or value == "null": - config[key] = None - else: - try: - config[key] = float(value) - except (ValueError, TypeError): - config[key] = None - logger.warning( - "Invalid timeout value for %s: %s, setting to None", - key, - value, - ) - return config - - def create(self) -> WebDriver: - """Create and return the WebDriver instance. - - This is the public interface for creating the driver. It wraps - the internal _create method for backward compatibility. - """ - return self._create() - - def destroy(self) -> None: - """Destroy the WebDriver instance. - - This is the public interface for cleanup. It wraps the internal - _destroy method and should be called when done with the driver. - """ - self._destroy() - - def _create(self) -> WebDriver: - pixel_density = app.config["WEBDRIVER_WINDOW"].get("pixel_density", 1) - - # Get driver class and initial kwargs based on driver type - if self._driver_type == "firefox": - driver_class, service_class, kwargs = self._create_firefox_driver( - pixel_density - ) - elif self._driver_type == "chrome": - driver_class, service_class, kwargs = self._create_chrome_driver( - pixel_density - ) - else: - raise Exception( # pylint: disable=broad-exception-raised - f"Webdriver name ({self._driver_type}) not supported" - ) - - # Add additional arguments from config - options = kwargs["options"] - for arg in list(app.config["WEBDRIVER_OPTION_ARGS"]): - options.add_argument(arg) - - # Fix timeout values for urllib3 2.x compatibility - webdriver_config = app.config["WEBDRIVER_CONFIGURATION"].copy() - webdriver_config = self._normalize_timeout_values(webdriver_config) - kwargs.update(webdriver_config) - - # Set the binary location if provided - # We need to pop it from the dict due to selenium_version < 4.10.0 - options.binary_location = webdriver_config.pop("binary_location", "") - - if version.parse(selenium_version) < version.parse("4.10.0"): - kwargs |= webdriver_config - else: - driver_opts = dict( - webdriver_config.get("options", {"capabilities": {}, "preferences": {}}) - ) - driver_srv = dict( - webdriver_config.get( - "service", - { - "log_output": "/dev/null", - "service_args": [], - "port": 0, - "env": {}, - }, - ) - ) - for name, value in driver_opts.get("capabilities", {}).items(): - options.set_capability(name, value) - if hasattr(options, "profile"): - for name, value in driver_opts.get("preferences", {}).items(): - options.profile.set_preference(str(name), value) - kwargs |= { - "options": options, - "service": service_class(**driver_srv), - } - - logger.debug("Init selenium driver") - return driver_class(**kwargs) - - def _auth(self, user: User) -> None: - """Authenticate the persistent driver in-place.""" - if self._driver is None: - raise RuntimeError("WebDriver is not initialized") - machine_auth_provider_factory.instance.authenticate_webdriver( - self._driver, user - ) - - def _destroy(self, tries: int = 2) -> None: - """Destroy the persistent driver""" - if not self._driver: - return - # This is some very flaky code in selenium. Hence the retries - # and catch-all exceptions - try: - retry_call(self._driver.close, max_tries=tries) - except Exception: # pylint: disable=broad-except # noqa: S110 - pass - try: - self._driver.quit() - except Exception: # pylint: disable=broad-except # noqa: S110 - pass - self._driver = None - - @staticmethod - def find_unexpected_errors( - driver: WebDriver, log_context: str | None = None - ) -> list[str]: - error_messages = [] - context_suffix = f" [{log_context}]" if log_context else "" - - try: - alert_divs = driver.find_elements(By.XPATH, "//div[@role = 'alert']") - logger.debug( - "%i alert elements have been found in the screenshot%s", - len(alert_divs), - context_suffix, - ) - - for alert_div in alert_divs: - # See More button - alert_div.find_element(By.XPATH, ".//*[@role = 'button']").click() - - # wait for modal to show up - modal = WebDriverWait( - driver, - app.config["SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE"], - ).until( - EC.visibility_of_any_elements_located( - (By.CLASS_NAME, "ant-modal-container") - ) - )[0] - - err_msg_div = modal.find_element(By.CLASS_NAME, "ant-modal-body") - - # collect error message - error_messages.append(err_msg_div.text) - - # close modal after collecting error messages - modal.find_element(By.CLASS_NAME, "ant-modal-close").click() - - # wait until the modal becomes invisible - WebDriverWait( - driver, - app.config["SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE"], - ).until(EC.invisibility_of_element(modal)) - - # Use HTML so that error messages are shown in the same style (color) - error_as_html = err_msg_div.get_attribute("innerHTML").replace( - "'", "\\'" - ) - - try: - # Even if some errors can't be updated in the screenshot, - # keep all the errors in the server log and do not fail the loop - driver.execute_script( - f"arguments[0].innerHTML = '{error_as_html}'", alert_div - ) - except WebDriverException: - logger.exception( - "Failed to update error messages using alert_div%s", - context_suffix, - ) - except WebDriverException: - logger.exception("Failed to capture unexpected errors%s", context_suffix) - - return error_messages - - def get_screenshot( # noqa: C901 - self, - url: str, - element_name: str, - user: User | None = None, - log_context: str | None = None, - report_execution_context: ReportExecutionContext | None = None, - ) -> bytes | None: - if report_execution_context: - log_context = report_execution_context.log_context - - def phase_timeout( - phase: str, - requested_seconds: float | None, - reserve_seconds: float = 0.0, - ) -> float: - if report_execution_context: - return report_execution_context.deadline.timeout_seconds( - phase, - requested_seconds=requested_seconds, - reserve_seconds=reserve_seconds, - ) - return float(requested_seconds or self._screenshot_load_wait) - - context_suffix = f" [{log_context}]" if log_context else "" - - # If a user is passed explicitly and differs from the stored user, - # update and re-authenticate - if user and user != self._user: - self._user = user - if self._driver: - self._destroy() - driver = self.driver - if report_execution_context: - driver.set_page_load_timeout( - phase_timeout( - "browser_navigation", - None, - report_execution_context.readiness_reserve_seconds, - ) - ) - driver.get(url) - img: bytes | None = None - selenium_headstart = app.config["SCREENSHOT_SELENIUM_HEADSTART"] - if report_execution_context: - selenium_headstart = min( - selenium_headstart, - phase_timeout( - "browser_headstart", - None, - report_execution_context.readiness_reserve_seconds, - ), - ) - logger.debug("Sleeping for %i seconds%s", selenium_headstart, context_suffix) - sleep(selenium_headstart) - - # WebDriver cleanup is intentionally not performed in this method. When the - # driver is used persistently (e.g., cache warmup), cleanup is handled - # externally via destroy(). When used for one-off screenshots, the caller or - # __del__ handles cleanup. - try: - try: - # page didn't load - logger.debug( - "Wait for the presence of %s at url: %s%s", - element_name, - url, - context_suffix, - ) - element = WebDriverWait( - driver, - phase_timeout( - "dashboard_mount", - self._screenshot_locate_wait, - ( - report_execution_context.readiness_reserve_seconds - if report_execution_context - else 0.0 - ), - ), - ).until(EC.presence_of_element_located((By.CLASS_NAME, element_name))) - except TimeoutException: - logger.warning( - "Selenium timed out requesting url %s%s", - url, - context_suffix, - exc_info=True, - ) - raise - - if report_execution_context and element_name in { - "standalone", - "chart-container", - }: - readiness_predicate = ( - REPORT_CHART_HOLDERS_READY_JS - if element_name == "standalone" - else CHART_CONTAINER_READY_JS - ) - readiness_timeout = phase_timeout( - "chart_readiness", - None, - report_execution_context.readiness_reserve_seconds, - ) - try: - WebDriverWait(driver, readiness_timeout).until( - lambda webdriver: webdriver.execute_script( - f"return ({readiness_predicate})()" - ) - ) - holder_states = ( - driver.execute_script( - f"return ({FIND_CHART_HOLDER_STATES_JS})()" - ) - if element_name == "standalone" - else [ - { - "chartId": report_execution_context.chart_id, - "state": "rendered", - } - ] - ) - ready_states = {"rendered", "empty", "error", "virtualized"} - deadline = report_execution_context.deadline - logger.info( - "report_readiness_ready url=%s expected_holders=%s " - "mounted_holders=%s ready_holders=%s elapsed_seconds=%s " - "remaining_seconds=%s%s", - url, - report_execution_context.expected_chart_count, - len(holder_states), - sum( - holder.get("state") in ready_states - for holder in holder_states - ), - f"{deadline.elapsed_seconds:.2f}", - f"{deadline.remaining_seconds:.2f}", - f" [{log_context}]" if log_context else "", - ) - except TimeoutException: - holder_states = ( - driver.execute_script( - f"return ({FIND_CHART_HOLDER_STATES_JS})()" - ) - if element_name == "standalone" - else [ - { - "chartId": report_execution_context.chart_id, - "state": "not_ready", - } - ] - ) - ready_states = {"rendered", "empty", "error", "virtualized"} - ready_holders = sum( - holder.get("state") in ready_states for holder in holder_states - ) - deadline = report_execution_context.deadline - logger.warning( - "report_readiness_terminal url=%s expected_holders=%s " - "mounted_holders=%s ready_holders=%s elapsed_seconds=%s " - "remaining_seconds=%s effective_wait_seconds=%.2f%s " - "terminal_reason=readiness_timeout states=%s; " - "aborting before capture or delivery", - url, - report_execution_context.expected_chart_count, - len(holder_states), - ready_holders, - f"{deadline.elapsed_seconds:.2f}", - f"{deadline.remaining_seconds:.2f}", - readiness_timeout, - f" [{log_context}]" if log_context else "", - holder_states, - ) - raise - else: - try: - # chart containers didn't render - logger.debug("Wait for chart containers to draw at url: %s", url) - WebDriverWait( - driver, - phase_timeout( - "chart_mount", - self._screenshot_locate_wait, - ( - report_execution_context.readiness_reserve_seconds - if report_execution_context - else 0.0 - ), - ), - ).until( - EC.visibility_of_all_elements_located( - (By.CLASS_NAME, "chart-container") - ) - ) - except TimeoutException: - if element_name == "standalone": - logger.info("Timeout Exception caught") - # Preserve support for empty dashboard thumbnails. Report - # dashboards use the positive holder gate above instead. - try: - WebDriverWait(driver, 0).until( - EC.visibility_of_all_elements_located( - (By.CLASS_NAME, "grid-container") - ) - ) - except Exception: - logger.warning( - "Selenium timed out waiting for dashboard to draw " - "at url %s", - url, - exc_info=True, - ) - raise - else: - logger.warning( - "Selenium timed out waiting for chart to draw at url %s", - url, - exc_info=True, - ) - raise - - try: - # charts took too long to load - logger.debug( - "Wait for loading element of charts to be gone at url: %s", - url, - ) - WebDriverWait( - driver, - phase_timeout( - "chart_readiness", - self._screenshot_load_wait, - ( - report_execution_context.readiness_reserve_seconds - if report_execution_context - else 0.0 - ), - ), - ).until_not( - EC.presence_of_all_elements_located((By.CLASS_NAME, "loading")) - ) - except TimeoutException: - logger.warning( - "Selenium timed out waiting for charts to load at url %s%s", - url, - context_suffix, - exc_info=True, - ) - raise - - selenium_animation_wait = app.config["SCREENSHOT_SELENIUM_ANIMATION_WAIT"] - if report_execution_context: - selenium_animation_wait = min( - selenium_animation_wait, - phase_timeout( - "chart_animation", - None, - report_execution_context.readiness_reserve_seconds, - ), - ) - logger.debug( - "Wait %i seconds for chart animation%s", - selenium_animation_wait, - context_suffix, - ) - sleep(selenium_animation_wait) - logger.debug( - "Taking a PNG screenshot of url %s as user %s%s", - url, - self._user.username if self._user else "None", - context_suffix, - ) - - if app.config["SCREENSHOT_REPLACE_UNEXPECTED_ERRORS"]: - unexpected_errors = WebDriverSelenium.find_unexpected_errors( - driver, log_context=log_context - ) - if unexpected_errors: - logger.warning( - "%i errors found in the screenshot. URL: %s. Errors are: %s%s", - len(unexpected_errors), - url, - unexpected_errors, - context_suffix, - ) - - if report_execution_context: - phase_timeout( - "screenshot_capture", - None, - report_execution_context.post_capture_reserve_seconds, - ) - img = element.screenshot_as_png - except TimeoutException: - # Already logged at WARNING in the inner handlers above - raise - except StaleElementReferenceException: - logger.warning( - "Selenium got a stale element while requesting url %s%s", - url, - context_suffix, - exc_info=True, - ) - raise - except WebDriverException: - logger.warning( - "Encountered an unexpected error when requesting url %s%s", - url, - context_suffix, - exc_info=True, - ) - raise - except Exception as ex: - logger.warning("exception in webdriver%s", context_suffix, exc_info=ex) - raise - return img diff --git a/tests/integration_tests/thumbnails_tests.py b/tests/integration_tests/thumbnails_tests.py index b29f6ce71b7..89b21e608a4 100644 --- a/tests/integration_tests/thumbnails_tests.py +++ b/tests/integration_tests/thumbnails_tests.py @@ -19,7 +19,7 @@ import urllib.request from unittest import skipUnless -from unittest.mock import ANY, call, MagicMock, patch +from unittest.mock import patch import pytest from flask_testing import LiveServerTestCase @@ -36,8 +36,6 @@ from superset.utils.screenshots import ( DashboardScreenshot, ScreenshotCachePayload, ) -from superset.utils.urls import get_url_path -from superset.utils.webdriver import WebDriverSelenium from tests.integration_tests.base_tests import SupersetTestCase from tests.integration_tests.conftest import with_feature_flags from tests.integration_tests.constants import ADMIN_USERNAME, ALPHA_USERNAME @@ -82,115 +80,6 @@ class TestThumbnailsSeleniumLive(LiveServerTestCase): assert response.getcode() == 202 -class TestWebDriverScreenshotErrorDetector(SupersetTestCase): - @patch("superset.utils.webdriver.WebDriverWait") - @patch("superset.utils.webdriver.firefox") - @patch("superset.utils.webdriver.WebDriverSelenium.find_unexpected_errors") - def test_not_call_find_unexpected_errors_if_feature_disabled( - self, mock_find_unexpected_errors, mock_firefox, mock_webdriver_wait - ): - webdriver_proxy = WebDriverSelenium("firefox") - user = security_manager.get_user_by_username(ADMIN_USERNAME) - url = get_url_path("Superset.dashboard", dashboard_id_or_slug=1) - webdriver_proxy.get_screenshot(url, "grid-container", user=user) - - assert not mock_find_unexpected_errors.called - - @patch("superset.utils.webdriver.WebDriverWait") - @patch("superset.utils.webdriver.firefox") - @patch("superset.utils.webdriver.WebDriverSelenium.find_unexpected_errors") - def test_call_find_unexpected_errors_if_feature_enabled( - self, mock_find_unexpected_errors, mock_firefox, mock_webdriver_wait - ): - app.config["SCREENSHOT_REPLACE_UNEXPECTED_ERRORS"] = True - webdriver_proxy = WebDriverSelenium("firefox") - user = security_manager.get_user_by_username(ADMIN_USERNAME) - url = get_url_path("Superset.dashboard", dashboard_id_or_slug=1) - webdriver_proxy.get_screenshot(url, "grid-container", user=user) - - assert mock_find_unexpected_errors.called - - app.config["SCREENSHOT_REPLACE_UNEXPECTED_ERRORS"] = False - - def test_find_unexpected_errors_no_alert(self): - webdriver = MagicMock() - - webdriver.find_elements.return_value = [] - - unexpected_errors = WebDriverSelenium.find_unexpected_errors(driver=webdriver) - assert len(unexpected_errors) == 0 - - assert "alert" in webdriver.find_elements.call_args_list[0][0][1] - - @patch("superset.utils.webdriver.WebDriverWait") - def test_find_unexpected_errors(self, mock_webdriver_wait): - webdriver = MagicMock() - alert_div = MagicMock() - - webdriver.find_elements.return_value = [alert_div] - alert_div.find_elements.return_value = MagicMock() - - unexpected_errors = WebDriverSelenium.find_unexpected_errors(driver=webdriver) - assert len(unexpected_errors) == 1 - - # attempt to find alerts - assert "alert" in webdriver.find_elements.call_args_list[0][0][1] - # attempt to click on "See more" buttons - assert "button" in alert_div.find_element.call_args_list[0][0][1] - # Wait for error modal to show up and to hide - assert 2 == len(mock_webdriver_wait.call_args_list) - # replace the text in alert div, eg, "unexpected errors" - assert alert_div == webdriver.execute_script.call_args_list[0][0][1] - - -class TestWebDriverSelenium(SupersetTestCase): - @patch("superset.utils.webdriver.WebDriverWait") - @patch("superset.utils.webdriver.firefox") - @patch("superset.utils.webdriver.sleep") - def test_screenshot_selenium_headstart( - self, mock_sleep, mock_webdriver, mock_webdriver_wait - ): - user = security_manager.get_user_by_username(ADMIN_USERNAME) - webdriver = WebDriverSelenium("firefox", user=user) - url = get_url_path("Superset.slice", slice_id=1, standalone="true") - app.config["SCREENSHOT_SELENIUM_HEADSTART"] = 5 - webdriver.get_screenshot(url, "chart-container") - assert mock_sleep.call_args_list[0] == call(5) - - @patch("superset.utils.webdriver.WebDriverWait") - @patch("superset.utils.webdriver.firefox") - def test_screenshot_selenium_locate_wait(self, mock_webdriver, mock_webdriver_wait): - app.config["SCREENSHOT_LOCATE_WAIT"] = 15 - user = security_manager.get_user_by_username(ADMIN_USERNAME) - webdriver = WebDriverSelenium("firefox", user=user) - url = get_url_path("Superset.slice", slice_id=1, standalone="true") - webdriver.get_screenshot(url, "chart-container") - assert mock_webdriver_wait.call_args_list[0] == call(ANY, 15) - - @patch("superset.utils.webdriver.WebDriverWait") - @patch("superset.utils.webdriver.firefox") - def test_screenshot_selenium_load_wait(self, mock_webdriver, mock_webdriver_wait): - app.config["SCREENSHOT_LOAD_WAIT"] = 15 - user = security_manager.get_user_by_username(ADMIN_USERNAME) - webdriver = WebDriverSelenium("firefox", user=user) - url = get_url_path("Superset.slice", slice_id=1, standalone="true") - webdriver.get_screenshot(url, "chart-container") - assert mock_webdriver_wait.call_args_list[2] == call(ANY, 15) - - @patch("superset.utils.webdriver.WebDriverWait") - @patch("superset.utils.webdriver.firefox") - @patch("superset.utils.webdriver.sleep") - def test_screenshot_selenium_animation_wait( - self, mock_sleep, mock_webdriver, mock_webdriver_wait - ): - user = security_manager.get_user_by_username(ADMIN_USERNAME) - webdriver = WebDriverSelenium("firefox", user=user) - url = get_url_path("Superset.slice", slice_id=1, standalone="true") - app.config["SCREENSHOT_SELENIUM_ANIMATION_WAIT"] = 4 - webdriver.get_screenshot(url, "chart-container") - assert mock_sleep.call_args_list[1] == call(4) - - class TestThumbnails(SupersetTestCase): mock_image = b"bytes mock image" digest_return_value = "foo_bar" diff --git a/tests/integration_tests/utils/machine_auth_tests.py b/tests/integration_tests/utils/machine_auth_tests.py index 40de236b6a9..bdbd62b90ab 100644 --- a/tests/integration_tests/utils/machine_auth_tests.py +++ b/tests/integration_tests/utils/machine_auth_tests.py @@ -15,9 +15,10 @@ # specific language governing permissions and limitations # under the License. -from unittest.mock import call, Mock, patch +from unittest.mock import MagicMock, patch from superset.extensions import machine_auth_provider_factory +from superset.utils.machine_auth import MachineAuthProvider from tests.integration_tests.base_tests import SupersetTestCase @@ -27,30 +28,35 @@ class MachineAuthProviderTests(SupersetTestCase): auth_cookies = machine_auth_provider_factory.instance.get_auth_cookies(user) assert auth_cookies["session"] is not None - @patch("superset.utils.machine_auth.MachineAuthProvider.get_auth_cookies") - def test_auth_driver_user(self, get_auth_cookies): + def test_authenticate_browser_context_sets_cookies(self): + """authenticate_browser_context navigates to login and sets auth cookies.""" user = self.get_user("admin") - driver = Mock() - get_auth_cookies.return_value = { - "session": "session_val", - "other_cookie": "other_val", - } - machine_auth_provider_factory.instance.authenticate_webdriver(driver, user) - driver.add_cookie.assert_has_calls( - [ - call({"name": "session", "value": "session_val"}), - call({"name": "other_cookie", "value": "other_val"}), - ] + provider = machine_auth_provider_factory.instance + + mock_context = MagicMock() + mock_page = MagicMock() + mock_context.new_page.return_value = mock_page + + with patch.object(provider, "get_cookies", return_value={"session": "abc123"}): + result = provider.authenticate_browser_context(mock_context, user) + + assert result is mock_context + mock_page.goto.assert_called_once() + mock_context.clear_cookies.assert_called_once() + mock_context.add_cookies.assert_called_once() + cookies_added = mock_context.add_cookies.call_args[0][0] + assert any( + c["name"] == "session" and c["value"] == "abc123" for c in cookies_added ) - @patch("superset.utils.machine_auth.request") - def test_auth_driver_request(self, request): - driver = Mock() - request.cookies = {"session": "session_val", "other_cookie": "other_val"} - machine_auth_provider_factory.instance.authenticate_webdriver(driver, None) - driver.add_cookie.assert_has_calls( - [ - call({"name": "session", "value": "session_val"}), - call({"name": "other_cookie", "value": "other_val"}), - ] - ) + def test_authenticate_browser_context_uses_override(self): + """authenticate_browser_context calls the override func when configured.""" + user = MagicMock() + mock_context = MagicMock() + mock_override = MagicMock(return_value=mock_context) + + provider = MachineAuthProvider(auth_webdriver_func_override=mock_override) + result = provider.authenticate_browser_context(mock_context, user) + + mock_override.assert_called_once_with(mock_context, user) + assert result is mock_context diff --git a/tests/unit_tests/mcp_service/test_pooled_screenshot.py b/tests/unit_tests/mcp_service/test_pooled_screenshot.py index da70bbacfd9..4b87c375017 100644 --- a/tests/unit_tests/mcp_service/test_pooled_screenshot.py +++ b/tests/unit_tests/mcp_service/test_pooled_screenshot.py @@ -17,20 +17,16 @@ from unittest.mock import MagicMock, patch -from superset.mcp_service.screenshot.pooled_screenshot import PooledBaseScreenshot +import pytest + +from superset.utils.screenshots import BaseScreenshot -@patch("superset.mcp_service.screenshot.pooled_screenshot.retry_screenshot_operation") -def test_get_screenshot_accepts_base_log_context( - mock_retry_screenshot_operation: MagicMock, -) -> None: - screenshot = PooledBaseScreenshot("http://example.com", "digest") +@patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", False) +def test_get_screenshot_raises_when_playwright_unavailable() -> None: + """get_screenshot raises RuntimeError via the real PLAYWRIGHT_AVAILABLE check.""" + screenshot = BaseScreenshot("http://example.com", "digest") user = MagicMock() - screenshot.get_screenshot(user, log_context="cache_key=abc") - - mock_retry_screenshot_operation.assert_called_once_with( - screenshot._get_screenshot_internal, # pylint: disable=protected-access - user, - None, - ) + with pytest.raises(RuntimeError, match="Playwright is required"): + screenshot.get_screenshot(user, log_context="cache_key=abc") diff --git a/tests/unit_tests/tasks/test_cache.py b/tests/unit_tests/tasks/test_cache.py index aebaa22725a..172b20009dc 100644 --- a/tests/unit_tests/tasks/test_cache.py +++ b/tests/unit_tests/tasks/test_cache.py @@ -20,7 +20,9 @@ from unittest import mock def _fake_app(config: Optional[dict[str, Any]] = None) -> mock.MagicMock: """Build a stand-in for ``current_app`` with a controllable config dict.""" - base: dict[str, Any] = {"WEBDRIVER_TYPE": "chrome"} + base: dict[str, Any] = { + "WEBDRIVER_WINDOW": {"dashboard": (1600, 1200)}, + } if config: base.update(config) app: mock.MagicMock = mock.MagicMock() @@ -32,7 +34,7 @@ def test_cache_warmup_unknown_strategy(app_context: None) -> None: """An unknown strategy name returns an explanatory message and warms nothing.""" from superset.tasks.cache import cache_warmup - with mock.patch("superset.tasks.cache.WebDriverSelenium") as mock_wd: + with mock.patch("superset.tasks.cache.WebDriverPlaywright") as mock_wd: result: dict[str, list[str]] | str = cache_warmup("does_not_exist") assert result == "No strategy does_not_exist found!" @@ -48,7 +50,7 @@ def test_cache_warmup_missing_config(app_context: None) -> None: "superset.tasks.cache.current_app", _fake_app({"SUPERSET_CACHE_WARMUP_USER": None}), ), - mock.patch("superset.tasks.cache.WebDriverSelenium") as mock_wd, + mock.patch("superset.tasks.cache.WebDriverPlaywright") as mock_wd, ): result: dict[str, list[str]] | str = cache_warmup("dummy") @@ -67,7 +69,7 @@ def test_cache_warmup_user_not_found(app_context: None) -> None: _fake_app({"SUPERSET_CACHE_WARMUP_USER": "bot"}), ), mock.patch("superset.tasks.cache.security_manager") as mock_sm, - mock.patch("superset.tasks.cache.WebDriverSelenium") as mock_wd, + mock.patch("superset.tasks.cache.WebDriverPlaywright") as mock_wd, ): mock_sm.find_user = mock.MagicMock(return_value=None) result: dict[str, list[str]] | str = cache_warmup("dummy") @@ -91,7 +93,7 @@ def test_cache_warmup_happy_path(app_context: None) -> None: _fake_app({"SUPERSET_CACHE_WARMUP_USER": "bot"}), ), mock.patch("superset.tasks.cache.security_manager") as mock_sm, - mock.patch("superset.tasks.cache.WebDriverSelenium") as mock_wd, + mock.patch("superset.tasks.cache.WebDriverPlaywright") as mock_wd, mock.patch("superset.tasks.cache.DummyStrategy.get_urls", return_value=urls), ): mock_sm.find_user = mock.MagicMock(return_value=user) @@ -99,9 +101,38 @@ def test_cache_warmup_happy_path(app_context: None) -> None: result: dict[str, list[str]] | str = cache_warmup("dummy") assert result == {"success": urls, "errors": []} - mock_wd.assert_called_once_with("chrome", user=user) + mock_wd.assert_called_once_with("", (1600, 1200)) assert driver.get_screenshot.call_count == len(urls) - driver.destroy.assert_called_once_with() + + +def test_cache_warmup_treats_none_screenshot_as_error(app_context: None) -> None: + """A URL whose screenshot returns None is recorded as an error.""" + from superset.tasks.cache import cache_warmup + + urls: list[str] = ["http://localhost/dash/ok", "http://localhost/dash/none"] + user: mock.MagicMock = mock.MagicMock() + + def side_effect(url: str, _element: str, **kwargs: Any) -> Any: + return None if url.endswith("none") else b"PNG" + + with ( + mock.patch( + "superset.tasks.cache.current_app", + _fake_app({"SUPERSET_CACHE_WARMUP_USER": "bot"}), + ), + mock.patch("superset.tasks.cache.security_manager") as mock_sm, + mock.patch("superset.tasks.cache.WebDriverPlaywright") as mock_wd, + mock.patch("superset.tasks.cache.DummyStrategy.get_urls", return_value=urls), + ): + mock_sm.find_user = mock.MagicMock(return_value=user) + driver: mock.MagicMock = mock_wd.return_value + driver.get_screenshot.side_effect = side_effect + result: dict[str, list[str]] | str = cache_warmup("dummy") + + assert result == { + "success": ["http://localhost/dash/ok"], + "errors": ["http://localhost/dash/none"], + } def test_cache_warmup_collects_errors_and_destroys(app_context: None) -> None: @@ -111,9 +142,10 @@ def test_cache_warmup_collects_errors_and_destroys(app_context: None) -> None: urls: list[str] = ["http://localhost/dash/ok", "http://localhost/dash/boom"] user: mock.MagicMock = mock.MagicMock() - def side_effect(url: str, _element: str) -> None: + def side_effect(url: str, _element: str, **kwargs: Any) -> bytes | None: if url.endswith("boom"): raise Exception("screenshot failed") + return b"PNG" with ( mock.patch( @@ -121,7 +153,7 @@ def test_cache_warmup_collects_errors_and_destroys(app_context: None) -> None: _fake_app({"SUPERSET_CACHE_WARMUP_USER": "bot"}), ), mock.patch("superset.tasks.cache.security_manager") as mock_sm, - mock.patch("superset.tasks.cache.WebDriverSelenium") as mock_wd, + mock.patch("superset.tasks.cache.WebDriverPlaywright") as mock_wd, mock.patch("superset.tasks.cache.DummyStrategy.get_urls", return_value=urls), ): mock_sm.find_user = mock.MagicMock(return_value=user) @@ -133,7 +165,6 @@ def test_cache_warmup_collects_errors_and_destroys(app_context: None) -> None: "success": ["http://localhost/dash/ok"], "errors": ["http://localhost/dash/boom"], } - driver.destroy.assert_called_once_with() def test_native_filter_options_strategy_returns_tasks_for_eligible_filters() -> None: diff --git a/tests/unit_tests/utils/screenshot_test.py b/tests/unit_tests/utils/screenshot_test.py index c9445d2165e..d6a26b21625 100644 --- a/tests/unit_tests/utils/screenshot_test.py +++ b/tests/unit_tests/utils/screenshot_test.py @@ -26,7 +26,6 @@ from superset.utils.hashing import hash_from_dict from superset.utils.screenshots import ( BaseScreenshot, ChartScreenshot, - DashboardScreenshot, ScreenshotCachePayload, ScreenshotCachePayloadType, ) @@ -271,7 +270,7 @@ class TestScreenshotCachePayloadGetImage: class TestBaseScreenshotDriverFallback: """Test BaseScreenshot.driver() fallback logic for Playwright migration.""" - @patch("superset.utils.screenshots.PLAYWRIGHT_AVAILABLE", True) + @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", True) @patch("superset.extensions.feature_flag_manager.is_feature_enabled") def test_driver_returns_playwright_when_feature_enabled_and_available( self, mock_feature_flag, screenshot_obj @@ -282,43 +281,8 @@ class TestBaseScreenshotDriverFallback: driver = screenshot_obj.driver() assert driver.__class__.__name__ == "WebDriverPlaywright" - mock_feature_flag.assert_called_once_with("PLAYWRIGHT_REPORTS_AND_THUMBNAILS") - @patch("superset.utils.screenshots.logger") - @patch("superset.utils.screenshots.PLAYWRIGHT_AVAILABLE", False) - @patch("superset.extensions.feature_flag_manager.is_feature_enabled") - def test_driver_falls_back_to_selenium_when_playwright_unavailable( - self, mock_feature_flag, mock_logger, screenshot_obj - ): - """Test driver() falls back to Selenium when Playwright unavailable.""" - mock_feature_flag.return_value = True - - driver = screenshot_obj.driver() - - assert driver.__class__.__name__ == "WebDriverSelenium" - # Should log the fallback message - mock_logger.info.assert_called_once() - log_call = mock_logger.info.call_args[0][0] - assert ( - "PLAYWRIGHT_REPORTS_AND_THUMBNAILS enabled but Playwright not installed" - in log_call - ) - assert "Falling back to Selenium" in log_call - assert "WebGL/Canvas charts may not render correctly" in log_call - - @patch("superset.extensions.feature_flag_manager.is_feature_enabled") - def test_driver_uses_selenium_when_feature_flag_disabled( - self, mock_feature_flag, screenshot_obj - ): - """Test driver() uses Selenium when feature flag disabled.""" - mock_feature_flag.return_value = False - - driver = screenshot_obj.driver() - - assert driver.__class__.__name__ == "WebDriverSelenium" - mock_feature_flag.assert_called_once_with("PLAYWRIGHT_REPORTS_AND_THUMBNAILS") - - @patch("superset.utils.screenshots.PLAYWRIGHT_AVAILABLE", True) + @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", True) @patch("superset.extensions.feature_flag_manager.is_feature_enabled") def test_driver_passes_window_size_to_playwright( self, mock_feature_flag, screenshot_obj @@ -332,20 +296,7 @@ class TestBaseScreenshotDriverFallback: assert driver._window == custom_window_size assert driver.__class__.__name__ == "WebDriverPlaywright" - @patch("superset.extensions.feature_flag_manager.is_feature_enabled") - def test_driver_passes_window_size_to_selenium( - self, mock_feature_flag, screenshot_obj - ): - """Test driver() passes window_size parameter to WebDriverSelenium.""" - mock_feature_flag.return_value = False - custom_window_size = (1200, 800) - - driver = screenshot_obj.driver(window_size=custom_window_size) - - assert driver._window == custom_window_size - assert driver.__class__.__name__ == "WebDriverSelenium" - - @patch("superset.utils.screenshots.PLAYWRIGHT_AVAILABLE", True) + @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", True) @patch("superset.extensions.feature_flag_manager.is_feature_enabled") def test_driver_uses_default_window_size_when_none_provided( self, mock_feature_flag, screenshot_obj @@ -360,9 +311,9 @@ class TestBaseScreenshotDriverFallback: class TestScreenshotSubclassesDriverBehavior: - """Test ChartScreenshot and DashboardScreenshot inherit driver behavior.""" + """Test ChartScreenshot inherits driver behavior.""" - @patch("superset.utils.screenshots.PLAYWRIGHT_AVAILABLE", True) + @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", True) @patch("superset.extensions.feature_flag_manager.is_feature_enabled") def test_chart_screenshot_uses_playwright_when_enabled(self, mock_feature_flag): """Test ChartScreenshot uses Playwright when feature enabled.""" @@ -374,27 +325,7 @@ class TestScreenshotSubclassesDriverBehavior: assert driver.__class__.__name__ == "WebDriverPlaywright" assert driver._window == chart_screenshot.window_size - @patch("superset.utils.screenshots.logger") - @patch("superset.utils.screenshots.PLAYWRIGHT_AVAILABLE", False) - @patch("superset.extensions.feature_flag_manager.is_feature_enabled") - def test_dashboard_screenshot_falls_back_to_selenium( - self, mock_feature_flag, mock_logger - ): - """Test DashboardScreenshot falls back to Selenium if no Playwright.""" - mock_feature_flag.return_value = True - - dashboard_screenshot = DashboardScreenshot( - "http://example.com/dashboard", "digest" - ) - driver = dashboard_screenshot.driver() - - assert driver.__class__.__name__ == "WebDriverSelenium" - assert driver._window == dashboard_screenshot.window_size - - # Should log the fallback message - mock_logger.info.assert_called_once() - - @patch("superset.utils.screenshots.PLAYWRIGHT_AVAILABLE", True) + @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", True) @patch("superset.extensions.feature_flag_manager.is_feature_enabled") def test_custom_window_size_passed_to_driver(self, mock_feature_flag): """Test custom window size is passed correctly to driver.""" diff --git a/tests/unit_tests/utils/test_playwright_migration_working.py b/tests/unit_tests/utils/test_playwright_migration_working.py deleted file mode 100644 index 7ab4aa4e393..00000000000 --- a/tests/unit_tests/utils/test_playwright_migration_working.py +++ /dev/null @@ -1,199 +0,0 @@ -# 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. - -""" -Working tests for Playwright migration functionality. -These tests demonstrate the core functionality works correctly. -""" - -from unittest.mock import MagicMock, patch - -from superset.utils.webdriver import ( - _PlaywrightBrowserManager, - PLAYWRIGHT_AVAILABLE, - validate_webdriver_config, -) - - -class TestPlaywrightMigrationCore: - """Core tests that demonstrate working Playwright migration functionality.""" - - def test_playwright_available_is_boolean(self): - """Test that PLAYWRIGHT_AVAILABLE is always a boolean.""" - assert isinstance(PLAYWRIGHT_AVAILABLE, bool) - - @patch("superset.extensions.feature_flag_manager.is_feature_enabled") - def test_validate_webdriver_config_structure(self, mock_feature_flag): - """Test that validate_webdriver_config returns correct structure.""" - mock_feature_flag.return_value = True - - result = validate_webdriver_config() - - # Check required keys exist - required_keys = [ - "selenium_available", - "playwright_available", - "playwright_feature_enabled", - "recommended_action", - ] - for key in required_keys: - assert key in result - - # Check data types - assert isinstance(result["selenium_available"], bool) - assert isinstance(result["playwright_available"], bool) - assert isinstance(result["playwright_feature_enabled"], bool) - assert result["recommended_action"] is None or isinstance( - result["recommended_action"], str - ) - - # Selenium should always be available - assert result["selenium_available"] is True - - @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", False) - @patch("superset.utils.webdriver.logger") - def test_webdriver_playwright_fallback_logging(self, mock_logger): - """Test that WebDriverPlaywright logs fallback correctly.""" - from superset.utils.webdriver import WebDriverPlaywright - - mock_user = MagicMock() - mock_user.username = "test_user" - - driver = WebDriverPlaywright("chrome") - result = driver.get_screenshot("http://example.com", "test-element", mock_user) - - # Should return None when unavailable - assert result is None - - # Should log the fallback message - mock_logger.info.assert_called_once() - log_call = mock_logger.info.call_args[0][0] - assert "Playwright not available" in log_call - assert "falling back to Selenium" in log_call - - def test_webdriver_classes_exist(self): - """Test that both WebDriver classes can be imported.""" - from superset.utils.webdriver import WebDriverPlaywright, WebDriverSelenium - - # Should be able to create instances without errors - playwright_driver = WebDriverPlaywright("chrome") - selenium_driver = WebDriverSelenium("chrome") - - assert playwright_driver is not None - assert selenium_driver is not None - - # Should have required attributes - assert hasattr(playwright_driver, "_driver_type") - assert hasattr(selenium_driver, "_driver_type") - - -class TestPlaywrightBrowserManager: - """Tests for the per-worker browser manager.""" - - def test_initial_state(self): - manager = _PlaywrightBrowserManager() - assert manager._playwright is None - assert manager._browser is None - - def test_get_browser_creates_browser(self): - mock_browser = MagicMock() - mock_browser.is_connected.return_value = True - - mock_pw_instance = MagicMock() - mock_pw_instance.chromium.launch.return_value = mock_browser - - mock_sync_pw = MagicMock() - mock_sync_pw.start.return_value = mock_pw_instance - - manager = _PlaywrightBrowserManager() - with patch( - "superset.utils.webdriver.sync_playwright", return_value=mock_sync_pw - ): - browser = manager.get_browser(["--headless"]) - - assert browser is mock_browser - mock_pw_instance.chromium.launch.assert_called_once_with(args=["--headless"]) - - def test_get_browser_reuses_connected_browser(self): - mock_browser = MagicMock() - mock_browser.is_connected.return_value = True - - manager = _PlaywrightBrowserManager() - manager._browser = mock_browser - manager._playwright = MagicMock() - - browser = manager.get_browser(["--headless"]) - - assert browser is mock_browser - # Should NOT launch a new browser - manager._playwright.chromium.launch.assert_not_called() - - def test_get_browser_recreates_on_disconnect(self): - stale_browser = MagicMock() - stale_browser.is_connected.return_value = False - - new_browser = MagicMock() - new_browser.is_connected.return_value = True - - mock_pw_instance = MagicMock() - mock_pw_instance.chromium.launch.return_value = new_browser - - mock_sync_pw = MagicMock() - mock_sync_pw.start.return_value = mock_pw_instance - - manager = _PlaywrightBrowserManager() - manager._browser = stale_browser - manager._playwright = MagicMock() - - with patch( - "superset.utils.webdriver.sync_playwright", return_value=mock_sync_pw - ): - browser = manager.get_browser(["--headless"]) - - assert browser is new_browser - stale_browser.close.assert_called_once() - - def test_cleanup(self): - mock_browser = MagicMock() - mock_playwright = MagicMock() - - manager = _PlaywrightBrowserManager() - manager._browser = mock_browser - manager._playwright = mock_playwright - - manager._cleanup() - - mock_browser.close.assert_called_once() - mock_playwright.stop.assert_called_once() - assert manager._browser is None - assert manager._playwright is None - - def test_cleanup_handles_exceptions(self): - mock_browser = MagicMock() - mock_browser.close.side_effect = Exception("crash") - mock_playwright = MagicMock() - mock_playwright.stop.side_effect = Exception("crash") - - manager = _PlaywrightBrowserManager() - manager._browser = mock_browser - manager._playwright = mock_playwright - - # Should not raise - manager._cleanup() - - assert manager._browser is None - assert manager._playwright is None diff --git a/tests/unit_tests/utils/webdriver_test.py b/tests/unit_tests/utils/webdriver_test.py index dbfd676935b..bc245a0ec36 100644 --- a/tests/unit_tests/utils/webdriver_test.py +++ b/tests/unit_tests/utils/webdriver_test.py @@ -15,7 +15,7 @@ # specific language governing permissions and limitations # under the License. -from unittest.mock import ANY, call, MagicMock, patch, PropertyMock +from unittest.mock import ANY, MagicMock, patch from uuid import UUID import pytest @@ -28,9 +28,7 @@ from superset.utils.webdriver import ( check_playwright_availability, PLAYWRIGHT_AVAILABLE, PLAYWRIGHT_INSTALL_MESSAGE, - validate_webdriver_config, WebDriverPlaywright, - WebDriverSelenium, ) @@ -64,391 +62,13 @@ def mock_app(): """Mock Flask app with webdriver configuration.""" app = MagicMock() app.config = { - "WEBDRIVER_TYPE": "chrome", "WEBDRIVER_OPTION_ARGS": [], - "WEBDRIVER_CONFIGURATION": {}, "SCREENSHOT_LOCATE_WAIT": 10, "SCREENSHOT_LOAD_WAIT": 10, } return app -class TestWebDriverSelenium: - """Test WebDriverSelenium timeout handling for urllib3 2.x compatibility.""" - - @patch("superset.utils.webdriver.app") - @patch("superset.utils.webdriver.firefox") - @patch("superset.utils.webdriver.chrome") - def test_timeout_conversion_to_float( - self, mock_chrome, mock_firefox, mock_app_patch, mock_app - ): - """Test that timeout values are properly converted to float.""" - # Set up app mock to be used throughout - mock_app_patch.config = { - "WEBDRIVER_TYPE": "chrome", - "WEBDRIVER_OPTION_ARGS": [], - "SCREENSHOT_LOCATE_WAIT": 10, - "SCREENSHOT_LOAD_WAIT": 10, - "WEBDRIVER_WINDOW": {}, - "WEBDRIVER_CONFIGURATION": { - "timeout": "30", - "connect_timeout": "10.5", - "socket_timeout": 20, - "read_timeout": "15.0", - "command_executor_timeout": "25", - }, - } - - mock_driver_class = MagicMock() - mock_chrome.webdriver.WebDriver = mock_driver_class - mock_chrome.service.Service = MagicMock() - mock_options = MagicMock() - mock_options.add_argument = MagicMock() - mock_chrome.options.Options = MagicMock(return_value=mock_options) - - driver = WebDriverSelenium(driver_type="chrome") - driver.create() - - # Check that the driver was called with float timeout values - mock_driver_class.assert_called_once() - call_kwargs = mock_driver_class.call_args.kwargs - assert call_kwargs["timeout"] == 30.0 - assert call_kwargs["connect_timeout"] == 10.5 - assert call_kwargs["socket_timeout"] == 20.0 - assert call_kwargs["read_timeout"] == 15.0 - assert call_kwargs["command_executor_timeout"] == 25.0 - - @patch("superset.utils.webdriver.app") - @patch("superset.utils.webdriver.chrome") - def test_timeout_none_handling(self, mock_chrome, mock_app_patch, mock_app): - """Test that None, 'None', and 'null' timeout values are set to None.""" - mock_app_patch.config = { - "WEBDRIVER_TYPE": "chrome", - "WEBDRIVER_OPTION_ARGS": [], - "SCREENSHOT_LOCATE_WAIT": 10, - "SCREENSHOT_LOAD_WAIT": 10, - "WEBDRIVER_WINDOW": {}, - "WEBDRIVER_CONFIGURATION": { - "timeout": None, - "connect_timeout": "None", - "socket_timeout": "null", - }, - } - - mock_driver_class = MagicMock() - mock_chrome.webdriver.WebDriver = mock_driver_class - mock_chrome.service.Service = MagicMock() - mock_options = MagicMock() - mock_options.add_argument = MagicMock() - mock_chrome.options.Options = MagicMock(return_value=mock_options) - - driver = WebDriverSelenium(driver_type="chrome") - driver.create() - - # Check that None values are preserved - mock_driver_class.assert_called_once() - call_kwargs = mock_driver_class.call_args.kwargs - assert call_kwargs["timeout"] is None - assert call_kwargs["connect_timeout"] is None - assert call_kwargs["socket_timeout"] is None - - @patch("superset.utils.webdriver.app") - @patch("superset.utils.webdriver.chrome") - @patch("superset.utils.webdriver.logger") - def test_invalid_timeout_warning( - self, mock_logger, mock_chrome, mock_app_patch, mock_app - ): - """Test that invalid timeout values log warnings and are set to None.""" - mock_app_patch.config = { - "WEBDRIVER_TYPE": "chrome", - "WEBDRIVER_OPTION_ARGS": [], - "SCREENSHOT_LOCATE_WAIT": 10, - "SCREENSHOT_LOAD_WAIT": 10, - "WEBDRIVER_WINDOW": {}, - "WEBDRIVER_CONFIGURATION": { - "timeout": "invalid", - "connect_timeout": "not_a_number", - "Page_Load_Timeout": "abc123", # Test case-insensitive matching - }, - } - - mock_driver_class = MagicMock() - mock_chrome.webdriver.WebDriver = mock_driver_class - mock_chrome.service.Service = MagicMock() - mock_options = MagicMock() - mock_options.add_argument = MagicMock() - mock_chrome.options.Options = MagicMock(return_value=mock_options) - - driver = WebDriverSelenium(driver_type="chrome") - driver.create() - - # Check that invalid values are set to None - mock_driver_class.assert_called_once() - call_kwargs = mock_driver_class.call_args.kwargs - assert call_kwargs["timeout"] is None - assert call_kwargs["connect_timeout"] is None - assert call_kwargs["Page_Load_Timeout"] is None - - # Check that warnings were logged with lazy logging format - assert mock_logger.warning.call_count == 3 - mock_logger.warning.assert_any_call( - "Invalid timeout value for %s: %s, setting to None", "timeout", "invalid" - ) - mock_logger.warning.assert_any_call( - "Invalid timeout value for %s: %s, setting to None", - "connect_timeout", - "not_a_number", - ) - mock_logger.warning.assert_any_call( - "Invalid timeout value for %s: %s, setting to None", - "Page_Load_Timeout", - "abc123", - ) - - @patch("superset.utils.webdriver.app") - @patch("superset.utils.webdriver.chrome") - def test_non_timeout_config_preserved(self, mock_chrome, mock_app_patch, mock_app): - """Test that non-timeout configuration values are preserved.""" - mock_app_patch.config = { - "WEBDRIVER_TYPE": "chrome", - "WEBDRIVER_OPTION_ARGS": [], - "SCREENSHOT_LOCATE_WAIT": 10, - "SCREENSHOT_LOAD_WAIT": 10, - "WEBDRIVER_WINDOW": {}, - "WEBDRIVER_CONFIGURATION": { - "timeout": "30", - "some_other_option": "value", - "another_option": 123, - "boolean_option": True, - }, - } - - mock_driver_class = MagicMock() - mock_chrome.webdriver.WebDriver = mock_driver_class - mock_chrome.service.Service = MagicMock() - mock_options = MagicMock() - mock_options.add_argument = MagicMock() - mock_chrome.options.Options = MagicMock(return_value=mock_options) - - driver = WebDriverSelenium(driver_type="chrome") - driver.create() - - # Check that all config values are passed through - mock_driver_class.assert_called_once() - call_kwargs = mock_driver_class.call_args.kwargs - assert call_kwargs["timeout"] == 30.0 - assert call_kwargs["some_other_option"] == "value" - assert call_kwargs["another_option"] == 123 - assert call_kwargs["boolean_option"] is True - - @patch("superset.utils.webdriver.app") - @patch("superset.utils.webdriver.chrome") - def test_timeout_key_case_insensitive(self, mock_chrome, mock_app_patch, mock_app): - """Test that timeout detection is case-insensitive.""" - mock_app_patch.config = { - "WEBDRIVER_TYPE": "chrome", - "WEBDRIVER_OPTION_ARGS": [], - "SCREENSHOT_LOCATE_WAIT": 10, - "SCREENSHOT_LOAD_WAIT": 10, - "WEBDRIVER_WINDOW": {}, - "WEBDRIVER_CONFIGURATION": { - "TIMEOUT": "10", - "Connect_Timeout": "20", - "SOCKET_TIMEOUT": "30", - "connection_timeout_ms": "5000", # Contains 'connection_timeout' - }, - } - - mock_driver_class = MagicMock() - mock_chrome.webdriver.WebDriver = mock_driver_class - mock_chrome.service.Service = MagicMock() - mock_options = MagicMock() - mock_options.add_argument = MagicMock() - mock_chrome.options.Options = MagicMock(return_value=mock_options) - - driver = WebDriverSelenium(driver_type="chrome") - driver.create() - - # Check that all timeout values are converted to float - mock_driver_class.assert_called_once() - call_kwargs = mock_driver_class.call_args.kwargs - assert call_kwargs["TIMEOUT"] == 10.0 - assert call_kwargs["Connect_Timeout"] == 20.0 - assert call_kwargs["SOCKET_TIMEOUT"] == 30.0 - assert call_kwargs["connection_timeout_ms"] == 5000.0 - - @patch("superset.utils.webdriver.app") - @patch("superset.utils.webdriver.chrome") - def test_empty_webdriver_config(self, mock_chrome, mock_app_patch, mock_app): - """Test handling of empty webdriver configuration.""" - mock_app_patch.config = { - "WEBDRIVER_TYPE": "chrome", - "WEBDRIVER_OPTION_ARGS": [], - "SCREENSHOT_LOCATE_WAIT": 10, - "SCREENSHOT_LOAD_WAIT": 10, - "WEBDRIVER_WINDOW": {}, - "WEBDRIVER_CONFIGURATION": {}, - } - - mock_driver_class = MagicMock() - mock_chrome.webdriver.WebDriver = mock_driver_class - mock_chrome.service.Service = MagicMock() - mock_options = MagicMock() - mock_options.add_argument = MagicMock() - mock_chrome.options.Options = MagicMock(return_value=mock_options) - - driver = WebDriverSelenium(driver_type="chrome") - driver.create() - - # Should create driver without errors - mock_driver_class.assert_called_once() - - @patch("superset.utils.webdriver.app") - def test_driver_sets_page_load_timeout(self, mock_app_patch: MagicMock) -> None: - """driver.get() must be bounded so it can't block forever (#40047).""" - mock_app_patch.config = { - "SCREENSHOT_LOCATE_WAIT": 10, - "SCREENSHOT_LOAD_WAIT": 10, - "SCREENSHOT_PAGE_LOAD_WAIT": 120, - } - mock_driver = MagicMock() - driver = WebDriverSelenium(driver_type="chrome", window=(800, 600)) - with patch.object(driver, "_create", return_value=mock_driver): - assert driver.driver is mock_driver - mock_driver.set_page_load_timeout.assert_called_once_with(120) - - @patch("superset.utils.webdriver.app") - def test_driver_skips_page_load_timeout_when_none( - self, mock_app_patch: MagicMock - ) -> None: - """Setting SCREENSHOT_PAGE_LOAD_WAIT to None disables the bound.""" - mock_app_patch.config = { - "SCREENSHOT_LOCATE_WAIT": 10, - "SCREENSHOT_LOAD_WAIT": 10, - "SCREENSHOT_PAGE_LOAD_WAIT": None, - } - mock_driver = MagicMock() - driver = WebDriverSelenium(driver_type="chrome", window=(800, 600)) - with patch.object(driver, "_create", return_value=mock_driver): - assert driver.driver is mock_driver - mock_driver.set_page_load_timeout.assert_not_called() - - @patch("superset.utils.webdriver.WebDriverWait") - @patch("superset.utils.webdriver.app") - def test_report_chart_uses_chart_readiness_not_dashboard_holders( - self, - mock_app_patch: MagicMock, - mock_wait: MagicMock, - ) -> None: - """Selenium chart reports require their chart terminal marker.""" - from selenium.common.exceptions import TimeoutException - - mock_app_patch.config = { - "SCREENSHOT_LOCATE_WAIT": 10, - "SCREENSHOT_LOAD_WAIT": 60, - "SCREENSHOT_PAGE_LOAD_WAIT": 120, - "SCREENSHOT_SELENIUM_HEADSTART": 0, - "SCREENSHOT_SELENIUM_ANIMATION_WAIT": 0, - "SCREENSHOT_REPLACE_UNEXPECTED_ERRORS": False, - } - mock_driver = MagicMock() - element = MagicMock() - mount_wait = MagicMock() - mount_wait.until.return_value = element - readiness_wait = MagicMock() - readiness_wait.until.side_effect = TimeoutException() - mock_wait.side_effect = [mount_wait, readiness_wait] - screenshot = WebDriverSelenium(driver_type="chrome") - screenshot._driver = mock_driver - - with ( - patch("superset.utils.webdriver.sleep"), - pytest.raises(TimeoutException), - ): - screenshot.get_screenshot( - "http://example.com/chart/7", - "chart-container", - report_execution_context=_report_context( - dashboard_id=None, - chart_id=7, - expected_chart_count=1, - ), - ) - - predicate = readiness_wait.until.call_args.args[0] - predicate(mock_driver) - readiness_js = mock_driver.execute_script.call_args.args[0] - assert "document.querySelector('.chart-container')" in readiness_js - assert "dashboard-component-chart-holder" not in readiness_js - assert element.screenshot_as_png.call_count == 0 - - @patch("superset.utils.webdriver.sleep") - @patch("superset.utils.webdriver.WebDriverWait") - @patch("superset.utils.webdriver.app") - def test_report_dashboard_budget_wires_selenium_timeouts_in_seconds( - self, - mock_app_patch: MagicMock, - mock_wait: MagicMock, - mock_sleep: MagicMock, - ) -> None: - """Selenium navigation, readiness, animation, and capture share one clock.""" - - mock_app_patch.config = { - "SCREENSHOT_LOCATE_WAIT": 10, - "SCREENSHOT_LOAD_WAIT": 60, - "SCREENSHOT_PAGE_LOAD_WAIT": 120, - "SCREENSHOT_SELENIUM_HEADSTART": 700, - "SCREENSHOT_SELENIUM_ANIMATION_WAIT": 700, - "SCREENSHOT_REPLACE_UNEXPECTED_ERRORS": False, - } - mock_driver = MagicMock() - mock_driver.execute_script.return_value = [ - {"chartId": "7", "state": "rendered"} - ] - element = MagicMock() - element.screenshot_as_png = b"screenshot" - mount_wait = MagicMock() - mount_wait.until.return_value = element - readiness_wait = MagicMock() - readiness_wait.until.return_value = True - mock_wait.side_effect = [mount_wait, readiness_wait] - context = ReportExecutionContext( - execution_id=UUID("084e7ee6-5557-4ecd-9632-b7f39c9ec524"), - report_schedule_id=11, - dashboard_id=805, - expected_chart_count=52, - deadline=ReportExecutionDeadline( - total_seconds=900, - started_at=0, - _clock=lambda: 100, - ), - capture_reserve_seconds=60, - delivery_reserve_seconds=120, - cleanup_reserve_seconds=30, - ) - screenshot = WebDriverSelenium(driver_type="chrome") - screenshot._driver = mock_driver - - assert ( - screenshot.get_screenshot( - "http://example.com/dashboard/805", - "standalone", - report_execution_context=context, - ) - == b"screenshot" - ) - - # 900 total - 100 elapsed - 210 reserved = 590 seconds. Selenium APIs - # take seconds (unlike Playwright's millisecond timeouts). - mock_driver.set_page_load_timeout.assert_called_once_with(590) - assert mock_wait.call_args_list == [ - call(mock_driver, 10), - call(mock_driver, 590), - ] - assert mock_sleep.call_args_list == [call(590), call(590)] - assert element.screenshot_as_png == b"screenshot" - - class TestPlaywrightAvailabilityCheck: """Test comprehensive Playwright availability checking.""" @@ -462,136 +82,14 @@ class TestPlaywrightAvailabilityCheck: @patch("superset.utils.webdriver.sync_playwright") @patch("superset.utils.webdriver.logger") - def test_check_playwright_availability_uses_lightweight_check( + def test_check_playwright_availability_returns_true_when_module_importable( self, mock_logger, mock_sync_playwright ): - """Test check_playwright_availability uses executable_path first.""" - # Setup mocks for successful executable path check - mock_playwright_instance = MagicMock() - mock_sync_playwright.return_value.__enter__.return_value = ( - mock_playwright_instance - ) - mock_playwright_instance.chromium.executable_path = "/path/to/chromium" - + """Test check_playwright_availability returns True when module is importable.""" result = check_playwright_availability() - assert result is True - # Should not launch browser if executable_path works - mock_playwright_instance.chromium.launch.assert_not_called() - - @patch("superset.utils.webdriver.sync_playwright") - @patch("superset.utils.webdriver.logger") - def test_check_playwright_availability_falls_back_to_launch( - self, mock_logger, mock_sync_playwright - ): - """Test check_playwright_availability falls back to browser launch.""" - # Setup mocks where executable_path fails but launch succeeds - mock_playwright_instance = MagicMock() - mock_browser = MagicMock() - - mock_sync_playwright.return_value.__enter__.return_value = ( - mock_playwright_instance - ) - # Make executable_path raise exception - type(mock_playwright_instance.chromium).executable_path = PropertyMock( - side_effect=Exception("executable_path failed") - ) - mock_playwright_instance.chromium.launch.return_value = mock_browser - - result = check_playwright_availability() - - assert result is True - # Should fall back to browser launch - mock_playwright_instance.chromium.launch.assert_called_once_with(headless=True) - mock_browser.close.assert_called_once() - - @patch("superset.utils.webdriver.sync_playwright") - @patch("superset.utils.webdriver.logger") - def test_check_playwright_availability_handles_browser_launch_failure( - self, mock_logger, mock_sync_playwright - ): - """Test check_playwright_availability handles browser launch failures.""" - # Setup mocks to raise exception on browser launch - mock_playwright_instance = MagicMock() - mock_sync_playwright.return_value.__enter__.return_value = ( - mock_playwright_instance - ) - # Mock executable_path to raise exception to force fallback to launch test - type(mock_playwright_instance.chromium).executable_path = PropertyMock( - side_effect=Exception("Executable path check failed") - ) - mock_playwright_instance.chromium.launch.side_effect = Exception( - "Browser binaries not installed" - ) - - result = check_playwright_availability() - - assert result is False - mock_logger.warning.assert_called_once() - warning_call = mock_logger.warning.call_args[0][0] - assert ( - "Playwright module is installed but browser launch failed" in warning_call - ) - assert "playwright install chromium" in warning_call - - @patch("superset.utils.webdriver.sync_playwright") - @patch("superset.utils.webdriver.logger") - def test_check_playwright_availability_handles_context_manager_error( - self, mock_logger, mock_sync_playwright - ): - """Test check_playwright_availability handles context manager errors.""" - # Setup mock to raise exception when entering context - mock_sync_playwright.return_value.__enter__.side_effect = Exception( - "Context error" - ) - - result = check_playwright_availability() - - assert result is False - mock_logger.warning.assert_called_once() - - -class TestPlaywrightMigrationSupport: - """Test Playwright migration and fallback functionality.""" - - @patch("superset.extensions.feature_flag_manager.is_feature_enabled") - def test_validate_webdriver_config_all_available(self, mock_feature_flag): - """Test validate_webdriver_config when all dependencies available.""" - mock_feature_flag.return_value = True - - result = validate_webdriver_config() - - assert result["selenium_available"] is True - assert isinstance(result["playwright_available"], bool) - assert isinstance(result["playwright_feature_enabled"], bool) - - if result["playwright_available"]: - assert result["recommended_action"] is None - else: - assert result["recommended_action"] == PLAYWRIGHT_INSTALL_MESSAGE - - @patch("superset.extensions.feature_flag_manager.is_feature_enabled") - def test_validate_webdriver_config_feature_flag_disabled(self, mock_feature_flag): - """Test validate_webdriver_config when feature flag is disabled.""" - mock_feature_flag.return_value = False - - result = validate_webdriver_config() - - assert result["selenium_available"] is True - assert result["playwright_feature_enabled"] is False - - @patch("superset.extensions.feature_flag_manager.is_feature_enabled") - @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", False) - def test_validate_webdriver_config_playwright_unavailable(self, mock_feature_flag): - """Test validate_webdriver_config when Playwright not available.""" - mock_feature_flag.return_value = True - - result = validate_webdriver_config() - - assert result["selenium_available"] is True - assert result["playwright_available"] is False - assert result["playwright_feature_enabled"] is True - assert result["recommended_action"] == PLAYWRIGHT_INSTALL_MESSAGE + # Only checks sync_playwright is not None — never launches browser + mock_sync_playwright.assert_not_called() class TestWebDriverPlaywrightFallback: @@ -599,24 +97,14 @@ class TestWebDriverPlaywrightFallback: @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", False) @patch("superset.utils.webdriver.logger") - def test_get_screenshot_returns_none_when_unavailable(self, mock_logger, mock_app): - """Test WebDriverPlaywright.get_screenshot returns None when unavailable.""" + def test_get_screenshot_raises_when_unavailable(self, mock_logger, mock_app): + """Test get_screenshot raises RuntimeError when Playwright is unavailable.""" mock_user = MagicMock() mock_user.username = "test_user" driver = WebDriverPlaywright("chrome") - result = driver.get_screenshot("http://example.com", "test-element", mock_user) - - assert result is None - - # Verify warning log was called with correct message - mock_logger.info.assert_called_once() - log_call = mock_logger.info.call_args[0][0] - assert "Playwright not available" in log_call - assert "falling back to Selenium" in log_call - assert "WebGL/Canvas charts may not render correctly" in log_call - # Check the substituted parameter - assert mock_logger.info.call_args[0][1] == PLAYWRIGHT_INSTALL_MESSAGE + with pytest.raises(RuntimeError, match="Playwright is required"): + driver.get_screenshot("http://example.com", "test-element", mock_user) @patch("superset.utils.webdriver.PLAYWRIGHT_AVAILABLE", True) @patch("superset.utils.webdriver._browser_manager") @@ -638,8 +126,6 @@ class TestWebDriverPlaywrightFallback: "SCREENSHOT_TILED_ENABLED": False, "SCREENSHOT_LOCATE_WAIT": 10, "SCREENSHOT_LOAD_WAIT": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE": 10, } # Setup playwright mocks @@ -688,7 +174,7 @@ class TestWebDriverPlaywrightFallback: mock_browser_manager.get_browser.return_value = mock_browser mock_browser.new_context.return_value = mock_context mock_context.new_page.return_value = mock_page - mock_page.goto.side_effect = PlaywrightTimeout() + mock_page.goto.side_effect = PlaywrightTimeout("timeout") with patch("superset.utils.webdriver.app") as mock_app: mock_app.config = { @@ -700,8 +186,6 @@ class TestWebDriverPlaywrightFallback: "SCREENSHOT_SELENIUM_ANIMATION_WAIT": 1, "SCREENSHOT_LOCATE_WAIT": 10, "SCREENSHOT_LOAD_WAIT": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE": 10, "SCREENSHOT_REPLACE_UNEXPECTED_ERRORS": True, "SCREENSHOT_TILED_ENABLED": False, } @@ -827,8 +311,6 @@ class TestWebDriverPlaywrightErrorHandling: "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_browser = MagicMock() @@ -887,8 +369,6 @@ class TestWebDriverPlaywrightErrorHandling: "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_browser = MagicMock() @@ -959,8 +439,6 @@ class TestWebDriverPlaywrightErrorHandling: "SCREENSHOT_SELENIUM_ANIMATION_WAIT": 1, "SCREENSHOT_LOCATE_WAIT": 10, "SCREENSHOT_LOAD_WAIT": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE": 10, "SCREENSHOT_REPLACE_UNEXPECTED_ERRORS": True, "SCREENSHOT_TILED_ENABLED": False, } @@ -1042,8 +520,6 @@ class TestWebDriverPlaywrightErrorHandling: "SCREENSHOT_SELENIUM_ANIMATION_WAIT": 1, "SCREENSHOT_LOCATE_WAIT": 10, "SCREENSHOT_LOAD_WAIT": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE": 10, "SCREENSHOT_REPLACE_UNEXPECTED_ERRORS": False, "SCREENSHOT_TILED_ENABLED": True, "SCREENSHOT_TILED_CHART_THRESHOLD": 20, @@ -1134,8 +610,6 @@ class TestWebDriverPlaywrightErrorHandling: "SCREENSHOT_SELENIUM_ANIMATION_WAIT": 1, "SCREENSHOT_LOCATE_WAIT": 10, "SCREENSHOT_LOAD_WAIT": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE": 10, "SCREENSHOT_REPLACE_UNEXPECTED_ERRORS": False, "SCREENSHOT_TILED_ENABLED": True, "SCREENSHOT_TILED_CHART_THRESHOLD": 20, @@ -1216,8 +690,6 @@ class TestWebDriverPlaywrightErrorHandling: "SCREENSHOT_SELENIUM_ANIMATION_WAIT": 1, "SCREENSHOT_LOCATE_WAIT": 10, "SCREENSHOT_LOAD_WAIT": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE": 10, "SCREENSHOT_REPLACE_UNEXPECTED_ERRORS": False, "SCREENSHOT_TILED_ENABLED": False, } @@ -1291,8 +763,6 @@ class TestWebDriverPlaywrightErrorHandling: "SCREENSHOT_SELENIUM_ANIMATION_WAIT": 0, "SCREENSHOT_LOCATE_WAIT": 10, "SCREENSHOT_LOAD_WAIT": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE": 10, "SCREENSHOT_REPLACE_UNEXPECTED_ERRORS": False, "SCREENSHOT_TILED_ENABLED": True, "SCREENSHOT_TILED_CHART_THRESHOLD": 20, @@ -1347,8 +817,6 @@ class TestWebDriverPlaywrightChartReadiness: "SCREENSHOT_TILED_ENABLED": False, "SCREENSHOT_LOCATE_WAIT": 10, "SCREENSHOT_LOAD_WAIT": 5, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE": 10, } def _make_pw_mocks(self, mock_browser_manager): @@ -1877,8 +1345,6 @@ class TestWebDriverPlaywrightAnimationWaitOrder: "SCREENSHOT_REPLACE_UNEXPECTED_ERRORS": False, "SCREENSHOT_LOCATE_WAIT": 10, "SCREENSHOT_LOAD_WAIT": 30, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_VISIBLE": 10, - "SCREENSHOT_WAIT_FOR_ERROR_MODAL_INVISIBLE": 10, } def _make_pw_mocks(self, mock_browser_manager):