fix(reports): positive readiness check for non-tiled screenshots

superset/config.py

@@ -1278,8 +1278,7 @@ def sync_theme_logo_href(
 # Time before selenium times out after trying to locate an element on the page and wait
 # for that element to load for a screenshot.
 SCREENSHOT_LOCATE_WAIT = int(timedelta(seconds=10).total_seconds())
-# Time before selenium times out after waiting for all DOM class elements named
-# "loading" are gone.
+# 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

superset/utils/screenshot_utils.py

@@ -22,13 +22,71 @@
 import time
 from typing import TYPE_CHECKING
 
+from celery import current_task
 from PIL import Image
 
 logger = logging.getLogger(__name__)
 
 # Time to wait after scrolling for content to settle and load (in milliseconds)
 SCROLL_SETTLE_TIMEOUT_MS = 1000
 
+# Runtime task-budget policy shared with the approach introduced in #42118.
+# Celery exposes the effective per-task hard/soft limits only on the running
+# task, so a static Superset timeout cannot reliably stay below task-level
+# overrides. Reserve at most 20% (capped at five minutes) for browser cleanup,
+# cache error transition, and the remaining report pipeline.
+SCREENSHOT_TASK_BUDGET_MARGIN_FRACTION = 0.2
+SCREENSHOT_TASK_BUDGET_MAX_MARGIN_SECONDS = 300
+
+
+def resolve_screenshot_task_budget_seconds(
+    log_context: str | None = None,
+) -> float | None:
+    """
+    Return the safe screenshot budget derived from the active Celery task.
+
+    Celery exposes ``request.timelimit`` as ``(hard, soft)``. Prefer the soft
+    limit because cleanup must finish before Celery raises it, falling back to
+    the hard limit when no soft limit is configured. Outside Celery, or when
+    the metadata is absent or malformed, return ``None`` so callers preserve
+    their configured standalone timeout.
+    """
+    context_suffix = f" [{log_context}]" if log_context else ""
+    try:
+        if not current_task:
+            return None
+        timelimit = current_task.request.timelimit
+        if not isinstance(timelimit, (tuple, list)) or len(timelimit) != 2:
+            return None
+        hard_limit, soft_limit = timelimit
+        limit = soft_limit or hard_limit
+        if isinstance(limit, bool) or not isinstance(limit, (int, float)) or limit <= 0:
+            return None
+        margin = min(
+            SCREENSHOT_TASK_BUDGET_MAX_MARGIN_SECONDS,
+            limit * SCREENSHOT_TASK_BUDGET_MARGIN_FRACTION,
+        )
+        budget = max(0.0, float(limit) - margin)
+        logger.debug(
+            "Screenshot budget derived from Celery task %s=%.1fs: %.1fs "
+            "(cleanup margin=%.1fs)%s",
+            "soft_time_limit" if soft_limit else "time_limit",
+            limit,
+            budget,
+            margin,
+            context_suffix,
+        )
+        return budget
+    except Exception:
+        logger.debug(
+            "Failed to derive screenshot budget from Celery task context; "
+            "using the configured screenshot timeout%s",
+            context_suffix,
+            exc_info=True,
+        )
+        return None
+
+
 try:
     from playwright.sync_api import TimeoutError as PlaywrightTimeout
 except ImportError:
@@ -40,79 +98,113 @@
     except ImportError:
         Page = None
 
-# Selectors used to build a positive per-tile readiness check. A chart holder
-# is only "ready" once it shows a terminal state (a rendered chart or an
-# error/empty state) -- the mere absence of a `.loading` element is not
-# sufficient, since a chart holder that intersects the viewport but hasn't
-# mounted anything yet (e.g. its IntersectionObserver callback hasn't fired)
-# would otherwise pass vacuously.
-# See superset-frontend/src/dashboard/components/gridComponents/ChartHolder/
-# ChartHolder.tsx for `data-test="dashboard-component-chart-holder"`,
-# superset-frontend/src/components/Chart/Chart.tsx for `.slice_container`
-# (rendered chart container, `data-test="slice-container"`) and `.loading`
-# (spinner, via the shared Loading component), and
-# superset-frontend/packages/superset-ui-core/src/components/EmptyState for
-# `.ant-empty` (e.g. "no results"/"add required control values" states).
-#
-# For diagnostics, each unready holder is additionally classified by *why*
-# it isn't ready, distinguishing a slow query from the virtualization race:
-#   - "waiting_on_database": `.loading` present with no `.slice_container`
-#     -- Chart.tsx's `renderSpinner()` replaces the whole container while
-#     the initial query is in flight (`chartStatus === 'loading'`).
-#   - "spinner_mounted": `.loading` present *inside* `.slice_container`
-#     -- the chart's query finished, but it isn't in the virtualization
-#     viewport yet, so `renderChartContainer()` shows a bare spinner instead
-#     of the chart.
-#   - "nothing_mounted": neither `.loading` nor any ready marker present --
-#     the vacuous-pass race this check exists to close.
-_UNREADY_CHART_HOLDERS_JS_BODY = """
-    const holders = document.querySelectorAll(
-        '[data-test="dashboard-component-chart-holder"]'
-    );
+# Production dashboard builds run ``babel-plugin-jsx-remove-data-test-id``
+# under the production BABEL_ENV (including Docker builds), so readiness must
+# never depend on ``data-test`` attributes. These runtime classes are the
+# production contract shared by readiness polling and diagnostics.
+CHART_HOLDER_SELECTOR = (
+    r'.dashboard-component-chart-holder[class*="dashboard-chart-id-"]'
+)
+SLICE_CONTAINER_SELECTOR = r".slice_container"
+LOADING_SELECTOR = r".loading"
+ALERT_SELECTOR = r'[role="alert"]'
+EMPTY_SELECTOR = r".ant-empty"
+MISSING_CHART_SELECTOR = r".missing-chart-container"
+TERMINAL_MARKER_SELECTOR = (
+    f"{SLICE_CONTAINER_SELECTOR}, {ALERT_SELECTOR}, {EMPTY_SELECTOR}, "
+    f"{MISSING_CHART_SELECTOR}"
+)
+CHART_ID_CLASS_PATTERN = r"\bdashboard-chart-id-(\d+)\b"
+
+# Shared body for holder readiness and timeout diagnostics. A holder is ready
+# only after a terminal marker appears and its loading marker disappears.
+UNREADY_CHART_HOLDERS_JS_BODY = f"""
+    const holders = document.querySelectorAll('{CHART_HOLDER_SELECTOR}');
     const unready = [];
-    for (const holder of holders) {
+    for (const holder of holders) {{
         const r = holder.getBoundingClientRect();
-        if (!(r.top < window.innerHeight && r.bottom > 0)) {
+        if (!(r.top < window.innerHeight && r.bottom > 0)) {{
             continue;
-        }
+        }}
         const hasSliceContainer = holder.querySelector(
-            '[data-test="slice-container"]'
+            '{SLICE_CONTAINER_SELECTOR}'
         ) !== null;
-        const stillLoading = holder.querySelector('.loading') !== null;
-        const isReady = hasSliceContainer || holder.querySelector(
-            '[role="alert"], .ant-empty, .missing-chart-container'
-        ) !== null;
-        if (stillLoading || !isReady) {
-            const chartIdEl = holder.querySelector('[data-test-chart-id]');
+        const stillLoading = holder.querySelector('{LOADING_SELECTOR}') !== null;
+        const isReady = holder.querySelector('{TERMINAL_MARKER_SELECTOR}') !== null;
+        if (stillLoading || !isReady) {{
+            const chartIdMatch = holder.className.match(/{CHART_ID_CLASS_PATTERN}/);
+            const chartId = chartIdMatch ? chartIdMatch[1] : null;
             let state;
-            if (stillLoading && hasSliceContainer) {
+            if (stillLoading && hasSliceContainer) {{
                 state = 'spinner_mounted';
-            } else if (stillLoading) {
+            }} else if (stillLoading) {{
                 state = 'waiting_on_database';
-            } else {
+            }} else {{
                 state = 'nothing_mounted';
-            }
-            unready.push({
-                chartId: chartIdEl
-                    ? chartIdEl.getAttribute('data-test-chart-id')
-                    : 'unknown',
+            }}
+            unready.push({{
+                chartId: chartId,
                 state: state,
-            });
-        }
-    }
+            }});
+        }}
+    }}
 """
 
-# Predicate for page.wait_for_function: true once every viewport-visible chart
-# holder has reached a terminal state.
-_TILE_READY_CHECK_JS = (
-    f"() => {{ {_UNREADY_CHART_HOLDERS_JS_BODY} return unready.length === 0; }}"
-)
+# Diagnostic query for every chart holder, including terminal and virtualized
+# states. It interpolates the same selector constants as the predicates.
+FIND_CHART_HOLDER_STATES_JS = f"""
+() => {{
+    const holders = document.querySelectorAll('{CHART_HOLDER_SELECTOR}');
+    return Array.from(holders).map(holder => {{
+        const chartIdMatch = holder.className.match(/{CHART_ID_CLASS_PATTERN}/);
+        const chartId = chartIdMatch ? chartIdMatch[1] : null;
+        const r = holder.getBoundingClientRect();
+        if (!(r.top < window.innerHeight && r.bottom > 0)) {{
+            return {{ chartId, state: 'virtualized' }};
+        }}
+        const hasSliceContainer = holder.querySelector(
+            '{SLICE_CONTAINER_SELECTOR}'
+        ) !== null;
+        const stillLoading = holder.querySelector('{LOADING_SELECTOR}') !== null;
+        if (stillLoading && hasSliceContainer) {{
+            return {{ chartId, state: 'spinner_mounted' }};
+        }}
+        if (stillLoading) {{
+            return {{ chartId, state: 'waiting_on_database' }};
+        }}
+        if (holder.querySelector('{ALERT_SELECTOR}') !== null) {{
+            return {{ chartId, state: 'error' }};
+        }}
+        if (holder.querySelector(
+            '{EMPTY_SELECTOR}, {MISSING_CHART_SELECTOR}'
+        ) !== null) {{
+            return {{ chartId, state: 'empty' }};
+        }}
+        if (hasSliceContainer) {{
+            return {{ chartId, state: 'rendered' }};
+        }}
+        return {{ chartId, state: 'nothing_mounted' }};
+    }});
+}}
+"""
 
-# Diagnostic query for page.evaluate: chart id + state of holders still not
-# ready, used to build the timeout log message.
-_FIND_UNREADY_CHART_HOLDERS_JS = (
-    f"() => {{ {_UNREADY_CHART_HOLDERS_JS_BODY} return unready; }}"
+CHART_HOLDERS_READY_JS = (
+    f"() => {{ {UNREADY_CHART_HOLDERS_JS_BODY} return unready.length === 0; }}"
 )
+FIND_UNREADY_CHART_HOLDERS_JS = (
+    f"() => {{ {UNREADY_CHART_HOLDERS_JS_BODY} return unready; }}"
+)
+
+# A chart capture has one target rather than dashboard holders, but needs the
+# same positive terminal-state guarantee and loading exclusion.
+CHART_CONTAINER_READY_JS = f"""
+() => {{
+    const chart = document.querySelector('.chart-container');
+    return chart !== null
+        && chart.querySelector('{LOADING_SELECTOR}') === null
+        && chart.querySelector('{TERMINAL_MARKER_SELECTOR}') !== null;
+}}
+"""
 
 
 def combine_screenshot_tiles(screenshot_tiles: list[bytes]) -> bytes:
@@ -249,12 +341,12 @@ def take_tiled_screenshot(
             tile_wait_start = time.monotonic()
             try:
                 page.wait_for_function(
-                    _TILE_READY_CHECK_JS,
+                    CHART_HOLDERS_READY_JS,
                     timeout=load_wait * 1000,
                 )
             except PlaywrightTimeout:
                 elapsed = time.monotonic() - tile_wait_start
-                unready_chart_holders = page.evaluate(_FIND_UNREADY_CHART_HOLDERS_JS)
+                unready_chart_holders = page.evaluate(FIND_UNREADY_CHART_HOLDERS_JS)
                 # A chart failing to load in time is a customer chart-loading
                 # issue (slow query, error state, etc.), not a Superset system
                 # fault, so this stays at WARNING -- the report still fails