fix(ag-grid): jpeg export of ag-grid tables (#38781)

(cherry picked from commit 69c8eef78e)

superset-frontend/packages/superset-ui-core/src/components/ThemedAgGridReact/ThemedAgGridReact.test.tsx

@@ -158,11 +158,13 @@ test('passes all props through to AgGridReact', () => {
     />,
   );
 
+  // onGridReady and onFirstDataRendered are intercepted by the component to expose
+  // the grid API on the container element; the wrapped function is passed instead.
   expect(AgGridReact).toHaveBeenCalledWith(
     expect.objectContaining({
       rowData: mockRowData,
       columnDefs: mockColumnDefs,
-      onGridReady,
+      onGridReady: expect.any(Function),
       onCellClicked,
       pagination: true,
       paginationPageSize: 10,
@@ -171,6 +173,47 @@ test('passes all props through to AgGridReact', () => {
   );
 });
 
+test('onGridReady wrapper calls user callback and exposes api on container', () => {
+  const onGridReady = jest.fn();
+
+  render(
+    <ThemedAgGridReact
+      rowData={mockRowData}
+      columnDefs={mockColumnDefs}
+      onGridReady={onGridReady}
+    />,
+  );
+
+  // Retrieve the wrapped handler that was passed to AgGridReact
+  const lastCall = (AgGridReact as jest.Mock).mock.calls.at(-1)[0];
+  const wrappedOnGridReady = lastCall.onGridReady as Function;
+
+  const mockApi = { setGridOption: jest.fn() };
+  wrappedOnGridReady({ api: mockApi });
+
+  // The user-provided callback must be forwarded
+  expect(onGridReady).toHaveBeenCalledWith({ api: mockApi });
+});
+
+test('onFirstDataRendered wrapper calls user callback', () => {
+  const onFirstDataRendered = jest.fn();
+
+  render(
+    <ThemedAgGridReact
+      rowData={mockRowData}
+      columnDefs={mockColumnDefs}
+      onFirstDataRendered={onFirstDataRendered}
+    />,
+  );
+
+  const lastCall = (AgGridReact as jest.Mock).mock.calls.at(-1)[0];
+  const wrappedOnFirstDataRendered = lastCall.onFirstDataRendered as Function;
+
+  wrappedOnFirstDataRendered({ firstRow: 0 });
+
+  expect(onFirstDataRendered).toHaveBeenCalledWith({ firstRow: 0 });
+});
+
 test('applies custom theme colors from Superset theme', () => {
   const customTheme = {
     ...supersetTheme,

superset-frontend/packages/superset-ui-core/src/components/ThemedAgGridReact/index.tsx

@@ -16,19 +16,28 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { useMemo, forwardRef } from 'react';
+import { useMemo, useRef, useCallback, forwardRef } from 'react';
 import { css } from '@emotion/react';
 import { AgGridReact, type AgGridReactProps } from 'ag-grid-react';
 import {
   themeQuartz,
   colorSchemeDark,
   colorSchemeLight,
+  type GridApi,
+  type GridReadyEvent,
+  type FirstDataRenderedEvent,
 } from 'ag-grid-community';
 import { useTheme, useThemeMode } from '@apache-superset/core/theme';
 
 // Note: With ag-grid v34's new theming API, CSS files are injected automatically
 // Do NOT import 'ag-grid-community/styles/ag-grid.css' or theme CSS files
 
+// Extends HTMLDivElement with ag-grid state attached to the container for downloadAsImage.
+export interface AgGridContainerElement extends HTMLDivElement {
+  _agGridApi?: GridApi;
+  _agGridFirstDataRendered?: boolean;
+}
+
 export interface ThemedAgGridReactProps extends AgGridReactProps {
   /**
    * Optional theme parameter overrides to customize specific ag-grid theme values.
@@ -71,9 +80,13 @@ export interface ThemedAgGridReactProps extends AgGridReactProps {
 export const ThemedAgGridReact = forwardRef<
   AgGridReact,
   ThemedAgGridReactProps
->(function ThemedAgGridReact({ themeOverrides, ...props }, ref) {
+>(function ThemedAgGridReact(
+  { themeOverrides, onGridReady, onFirstDataRendered, ...props },
+  ref,
+) {
   const theme = useTheme();
   const isDarkMode = useThemeMode();
+  const containerRef = useRef<AgGridContainerElement>(null);
 
   // Get the appropriate ag-grid theme based on dark/light mode
   const agGridTheme = useMemo(() => {
@@ -140,8 +153,32 @@ export const ThemedAgGridReact = forwardRef<
     return baseTheme.withParams(finalParams);
   }, [theme, isDarkMode, themeOverrides]);
 
+  // Expose gridApi and first-data-rendered flag on the container for downloadAsImage.
+  const handleGridReady = useCallback(
+    (event: GridReadyEvent) => {
+      if (containerRef.current) {
+        containerRef.current._agGridFirstDataRendered = false;
+        containerRef.current._agGridApi = event.api;
+      }
+      onGridReady?.(event);
+    },
+    [onGridReady],
+  );
+
+  // Mark the container once rows are painted so downloadAsImage can gate on readiness.
+  const handleFirstDataRendered = useCallback(
+    (event: FirstDataRenderedEvent) => {
+      if (containerRef.current) {
+        containerRef.current._agGridFirstDataRendered = true;
+      }
+      onFirstDataRendered?.(event);
+    },
+    [onFirstDataRendered],
+  );
+
   return (
     <div
+      ref={containerRef}
       css={css`
         width: 100%;
         height: 100%;
@@ -151,7 +188,13 @@ export const ThemedAgGridReact = forwardRef<
       `}
       data-themed-ag-grid="true"
     >
-      <AgGridReact ref={ref} theme={agGridTheme} {...props} />
+      <AgGridReact
+        ref={ref}
+        theme={agGridTheme}
+        onGridReady={handleGridReady}
+        onFirstDataRendered={handleFirstDataRendered}
+        {...props}
+      />
     </div>
   );
 });

superset-frontend/packages/superset-ui-core/src/components/index.ts

@@ -201,6 +201,7 @@ export * from './Result';
 export {
   ThemedAgGridReact,
   type ThemedAgGridReactProps,
+  type AgGridContainerElement,
   setupAGGridModules,
   defaultModules,
 } from './ThemedAgGridReact';

superset-frontend/src/utils/downloadAsImage.test.ts

@@ -0,0 +1,685 @@
+/**
+ * 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.
+ */
+import domToImage from 'dom-to-image-more';
+import { addWarningToast } from 'src/components/MessageToasts/actions';
+import downloadAsImageOptimized, {
+  waitForStableScrollHeight,
+} from './downloadAsImage';
+
+jest.mock('dom-to-image-more', () => ({
+  __esModule: true,
+  default: { toJpeg: jest.fn() },
+}));
+
+jest.mock('src/components/MessageToasts/actions', () => ({
+  addWarningToast: jest.fn(),
+}));
+
+jest.mock('@apache-superset/core/translation', () => ({
+  t: (str: string) => str,
+}));
+
+const mockToJpeg = domToImage.toJpeg as jest.Mock;
+const mockAddWarningToast = addWarningToast as jest.Mock;
+
+// document.fonts.ready is not implemented in jsdom; provide a resolved promise
+Object.defineProperty(document, 'fonts', {
+  value: { ready: Promise.resolve() },
+  configurable: true,
+});
+
+// Build a synthetic React event that resolves `currentTarget.closest()` to a given element
+function syntheticEventFor(el: Element) {
+  return { currentTarget: { closest: () => el } } as any;
+}
+
+// Build and attach an ag-grid DOM structure; returns cleanup function
+function buildAgGridElement() {
+  const container = document.createElement('div');
+  const agContainer = document.createElement('div');
+  agContainer.setAttribute('data-themed-ag-grid', 'true');
+  const agRootWrapper = document.createElement('div');
+  agRootWrapper.className = 'ag-root-wrapper';
+  agContainer.appendChild(agRootWrapper);
+  container.appendChild(agContainer);
+  document.body.appendChild(container);
+  return {
+    container,
+    agContainer,
+    agRootWrapper,
+    cleanup: () => document.body.removeChild(container),
+  };
+}
+
+// Attach a mock GridApi and set the first-data-rendered flag on the container
+function attachMockApi(
+  agContainer: HTMLElement,
+  { firstDataRendered = true } = {},
+) {
+  const api = { setGridOption: jest.fn() };
+  (agContainer as any)._agGridApi = api;
+  (agContainer as any)._agGridFirstDataRendered = firstDataRendered;
+  return api;
+}
+
+beforeEach(() => {
+  jest.clearAllMocks();
+  mockToJpeg.mockResolvedValue('data:image/jpeg;base64,test');
+});
+
+afterEach(() => {
+  jest.useRealTimers();
+});
+
+test('waitForStableScrollHeight resolves after 2 consecutive stable scrollHeight readings', async () => {
+  jest.useFakeTimers();
+  const el = document.createElement('div');
+  Object.defineProperty(el, 'scrollHeight', {
+    get: () => 100,
+    configurable: true,
+  });
+
+  const promise = waitForStableScrollHeight(el);
+  await jest.runAllTimersAsync();
+  await expect(promise).resolves.toBeUndefined();
+
+  jest.useRealTimers();
+});
+
+test('waitForStableScrollHeight respects a custom minStablePolls', async () => {
+  jest.useFakeTimers();
+  const el = document.createElement('div');
+  Object.defineProperty(el, 'scrollHeight', {
+    get: () => 100,
+    configurable: true,
+  });
+
+  // With minStablePolls=5 the promise must not resolve after just 2 polls.
+  const promise = waitForStableScrollHeight(el, 5000, 5);
+  jest.advanceTimersByTime(300);
+  let resolved = false;
+  promise.then(() => {
+    resolved = true;
+  });
+  // Flush microtasks so the .then() above has a chance to run if resolved
+  await Promise.resolve();
+  expect(resolved).toBe(false);
+
+  // Now run the remaining polls (2 more stable polls → total 5) and confirm resolution.
+  jest.advanceTimersByTime(300);
+  await expect(promise).resolves.toBeUndefined();
+
+  jest.useRealTimers();
+});
+
+test('waitForStableScrollHeight resets stable count when height changes mid-poll', async () => {
+  jest.useFakeTimers();
+  const el = document.createElement('div');
+  let height = 100;
+  Object.defineProperty(el, 'scrollHeight', {
+    get: () => height,
+    configurable: true,
+  });
+
+  const promise = waitForStableScrollHeight(el);
+  // Poll 1: height is 100, stableCount becomes 1
+  jest.advanceTimersByTime(100);
+  // Height changes — stable counter must reset
+  height = 200;
+  // Run until new height stabilises (2 consecutive 100 ms polls)
+  jest.advanceTimersByTime(300);
+  await expect(promise).resolves.toBeUndefined();
+
+  jest.useRealTimers();
+});
+
+test('waitForStableScrollHeight resolves after maxMs even if height never stabilises', async () => {
+  jest.useFakeTimers();
+  const el = document.createElement('div');
+  let height = 0;
+  Object.defineProperty(el, 'scrollHeight', {
+    // Always increments so stableFrames never reaches 4
+    get: () => {
+      height += 1;
+      return height;
+    },
+    configurable: true,
+  });
+
+  const promise = waitForStableScrollHeight(el, 200);
+  jest.advanceTimersByTime(400); // past the 200 ms deadline
+  await expect(promise).resolves.toBeUndefined();
+
+  jest.useRealTimers();
+});
+
+test('waitForStableScrollHeight resolves if scrollHeight throws (element removed from DOM)', async () => {
+  jest.useFakeTimers();
+  const el = document.createElement('div');
+  let shouldThrow = false;
+  Object.defineProperty(el, 'scrollHeight', {
+    get: () => {
+      if (shouldThrow) throw new Error('element detached');
+      return 100;
+    },
+    configurable: true,
+  });
+
+  const promise = waitForStableScrollHeight(el);
+  jest.advanceTimersByTime(100); // poll 1: stable, stableCount = 1
+  shouldThrow = true; // simulate DOM removal
+  jest.advanceTimersByTime(100); // poll 2: throws → resolves immediately
+  await expect(promise).resolves.toBeUndefined();
+
+  jest.useRealTimers();
+});
+
+test('shows warning toast when element is not found', async () => {
+  const handler = downloadAsImageOptimized('div', 'test');
+  // closest() returning null simulates a selector that matches nothing
+  await handler({ currentTarget: { closest: () => null } } as any);
+
+  expect(mockAddWarningToast).toHaveBeenCalledWith(
+    'Image download failed, please refresh and try again.',
+  );
+  expect(mockToJpeg).not.toHaveBeenCalled();
+});
+
+test('shows "still loading" toast when grid has not yet rendered its first rows', async () => {
+  const { container, agContainer, cleanup } = buildAgGridElement();
+  attachMockApi(agContainer, { firstDataRendered: false });
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  await handler(syntheticEventFor(container));
+
+  expect(mockAddWarningToast).toHaveBeenCalledWith(
+    'The chart is still loading. Please wait a moment and try again.',
+  );
+  expect(mockToJpeg).not.toHaveBeenCalled();
+
+  cleanup();
+});
+
+test('switches to print layout, captures JPEG, and restores normal layout', async () => {
+  jest.useFakeTimers();
+  const { container, agContainer, cleanup } = buildAgGridElement();
+  const api = attachMockApi(agContainer);
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  expect(api.setGridOption).toHaveBeenCalledWith('domLayout', 'print');
+  expect(mockToJpeg).toHaveBeenCalledWith(
+    expect.any(HTMLElement),
+    expect.objectContaining({ quality: 0.95 }),
+  );
+  expect(api.setGridOption).toHaveBeenCalledWith('domLayout', 'normal');
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('restores normal layout in finally even when image capture throws', async () => {
+  jest.useFakeTimers();
+  mockToJpeg.mockRejectedValue(new Error('capture failed'));
+  const { container, agContainer, cleanup } = buildAgGridElement();
+  const api = attachMockApi(agContainer);
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  expect(api.setGridOption).toHaveBeenCalledWith('domLayout', 'normal');
+  expect(mockAddWarningToast).toHaveBeenCalledWith(
+    'Image download failed, please refresh and try again.',
+  );
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('still captures image when _agGridApi is absent (graceful degradation)', async () => {
+  jest.useFakeTimers();
+  const { container, agContainer, cleanup } = buildAgGridElement();
+  // No API — only the first-data-rendered flag
+  (agContainer as any)._agGridFirstDataRendered = true;
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  expect(mockToJpeg).toHaveBeenCalled();
+  expect(mockAddWarningToast).not.toHaveBeenCalled();
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('resolves ag-cell min-height to row pixel height when content fits within it', async () => {
+  jest.useFakeTimers();
+  const { container, agContainer, agRootWrapper, cleanup } =
+    buildAgGridElement();
+  attachMockApi(agContainer);
+
+  // Build a row with a cell inside the grid
+  const row = document.createElement('div');
+  row.className = 'ag-row';
+  Object.defineProperty(row, 'offsetHeight', {
+    get: () => 32,
+    configurable: true,
+  });
+  const cell = document.createElement('div');
+  cell.className = 'ag-cell';
+  row.appendChild(cell);
+  agRootWrapper.appendChild(row);
+
+  let capturedMinHeight = '';
+  mockToJpeg.mockImplementation(() => {
+    capturedMinHeight = cell.style.minHeight;
+    return Promise.resolve('data:image/jpeg;base64,test');
+  });
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  // Cell min-height was resolved to the row's pixel height during capture
+  expect(capturedMinHeight).toBe('32px');
+  // Cell min-height was restored after capture
+  expect(cell.style.minHeight).toBe('');
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('uses cell scrollHeight when it exceeds row offsetHeight (stale row heights for off-screen rows)', async () => {
+  jest.useFakeTimers();
+  const { container, agContainer, agRootWrapper, cleanup } =
+    buildAgGridElement();
+  attachMockApi(agContainer);
+
+  const row = document.createElement('div');
+  row.className = 'ag-row';
+  Object.defineProperty(row, 'offsetHeight', {
+    get: () => 25,
+    configurable: true,
+  }); // stale default
+  const cell = document.createElement('div');
+  cell.className = 'ag-cell';
+  Object.defineProperty(cell, 'scrollHeight', {
+    get: () => 120,
+    configurable: true,
+  }); // actual content
+  row.appendChild(cell);
+  agRootWrapper.appendChild(row);
+
+  let capturedMinHeight = '';
+  mockToJpeg.mockImplementation(() => {
+    capturedMinHeight = cell.style.minHeight;
+    return Promise.resolve('data:image/jpeg;base64,test');
+  });
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  // Uses the larger content height, not the stale row height
+  expect(capturedMinHeight).toBe('120px');
+  expect(cell.style.minHeight).toBe('');
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('derives image width from getColumnState by summing visible column pixel widths', async () => {
+  jest.useFakeTimers();
+  const { container, agContainer, cleanup } = buildAgGridElement();
+  const api = attachMockApi(agContainer);
+
+  // 3 visible columns (200 + 350 + 150 = 700 px) plus one hidden column excluded from sum
+  (api as any).getColumnState = jest.fn(() => [
+    { colId: 'col1', width: 200, hide: false },
+    { colId: 'col2', width: 350, hide: false },
+    { colId: 'col3', width: 150, hide: false },
+    { colId: 'col4', width: 999, hide: true },
+  ]);
+  (api as any).applyColumnState = jest.fn();
+
+  let capturedWidth: number | undefined;
+  mockToJpeg.mockImplementation(
+    (_el: HTMLElement, opts: { width?: number }) => {
+      capturedWidth = opts.width;
+      return Promise.resolve('data:image/jpeg;base64,test');
+    },
+  );
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  // Width passed to toJpeg is the sum of visible column widths, not agRootWrapper.offsetWidth
+  expect(capturedWidth).toBe(700);
+  expect((api as any).getColumnState).toHaveBeenCalled();
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('restores column pixel widths via applyColumnState with flex stripped after print layout', async () => {
+  jest.useFakeTimers();
+  const { container, agContainer, cleanup } = buildAgGridElement();
+  const api = attachMockApi(agContainer);
+
+  const savedState = [
+    { colId: 'col1', width: 300, flex: 1, hide: false },
+    { colId: 'col2', width: 400, flex: 1.5, hide: false },
+  ];
+  (api as any).getColumnState = jest.fn(() => savedState);
+  (api as any).applyColumnState = jest.fn();
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  // flex must be stripped (set to null) so pixel width is used, not flex ratio
+  expect((api as any).applyColumnState).toHaveBeenCalledWith({
+    state: [
+      { colId: 'col1', width: 300, flex: null },
+      { colId: 'col2', width: 400, flex: null },
+    ],
+    applyOrder: false,
+  });
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('restores original column state with flex in finally after capture', async () => {
+  jest.useFakeTimers();
+  const { container, agContainer, cleanup } = buildAgGridElement();
+  const api = attachMockApi(agContainer);
+
+  const savedState = [
+    { colId: 'col1', width: 300, flex: 1, hide: false },
+    { colId: 'col2', width: 400, flex: 1.5, hide: false },
+  ];
+  (api as any).getColumnState = jest.fn(() => savedState);
+  (api as any).applyColumnState = jest.fn();
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  // Last call must restore the original state (with flex) so the live grid is unaffected
+  expect((api as any).applyColumnState.mock.calls.at(-1)[0]).toEqual({
+    state: savedState,
+    applyOrder: false,
+  });
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('falls back to agRootWrapper.offsetWidth when getColumnState returns no visible columns', async () => {
+  jest.useFakeTimers();
+  const { container, agContainer, agRootWrapper, cleanup } =
+    buildAgGridElement();
+  const api = attachMockApi(agContainer);
+
+  // All columns hidden → visible sum is 0 → fall back to offsetWidth
+  (api as any).getColumnState = jest.fn(() => [
+    { colId: 'col1', width: 500, hide: true },
+  ]);
+  (api as any).applyColumnState = jest.fn();
+
+  Object.defineProperty(agRootWrapper, 'offsetWidth', {
+    get: () => 600,
+    configurable: true,
+  });
+
+  let capturedWidth: number | undefined;
+  mockToJpeg.mockImplementation(
+    (_el: HTMLElement, opts: { width?: number }) => {
+      capturedWidth = opts.width;
+      return Promise.resolve('data:image/jpeg;base64,test');
+    },
+  );
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  expect(capturedWidth).toBe(600);
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('restores ag-cell styles after capture even when toJpeg throws', async () => {
+  jest.useFakeTimers();
+  mockToJpeg.mockRejectedValue(new Error('capture failed'));
+  const { container, agContainer, agRootWrapper, cleanup } =
+    buildAgGridElement();
+  attachMockApi(agContainer);
+
+  const row = document.createElement('div');
+  row.className = 'ag-row';
+  Object.defineProperty(row, 'offsetHeight', {
+    get: () => 28,
+    configurable: true,
+  });
+  const cell = document.createElement('div');
+  cell.className = 'ag-cell';
+  cell.style.minHeight = '100%';
+  cell.style.overflow = 'visible';
+  row.appendChild(cell);
+  agRootWrapper.appendChild(row);
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  // Styles restored to original values despite capture error
+  expect(cell.style.minHeight).toBe('100%');
+  expect(cell.style.overflow).toBe('visible');
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('calls resetRowHeights after print layout to force ag-grid to re-measure rows with stale cached heights', async () => {
+  jest.useFakeTimers();
+  const { container, agContainer, cleanup } = buildAgGridElement();
+  const api = attachMockApi(agContainer);
+  (api as any).resetRowHeights = jest.fn();
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  expect(api.setGridOption).toHaveBeenCalledWith('domLayout', 'print');
+  expect((api as any).resetRowHeights).toHaveBeenCalled();
+  expect(api.setGridOption).toHaveBeenCalledWith('domLayout', 'normal');
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('does not throw when resetRowHeights is absent from the api', async () => {
+  jest.useFakeTimers();
+  const { container, agContainer, cleanup } = buildAgGridElement();
+  attachMockApi(agContainer); // api has only setGridOption, no resetRowHeights
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await expect(exportPromise).resolves.toBeUndefined();
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('falls through to clone path for dashboard export with a single ag-grid chart', async () => {
+  const dashboard = document.createElement('div');
+  dashboard.className = 'dashboard';
+
+  const agContainer = document.createElement('div');
+  agContainer.setAttribute('data-themed-ag-grid', 'true');
+  const agRootWrapper = document.createElement('div');
+  agRootWrapper.className = 'ag-root-wrapper';
+  agContainer.appendChild(agRootWrapper);
+  (agContainer as any)._agGridFirstDataRendered = true;
+  dashboard.appendChild(agContainer);
+  document.body.appendChild(dashboard);
+
+  const handler = downloadAsImageOptimized('.dashboard', 'My Dashboard', true);
+  await handler({ currentTarget: {} } as any);
+
+  expect(mockToJpeg).toHaveBeenCalledWith(
+    expect.any(HTMLElement),
+    expect.objectContaining({ quality: 0.95 }),
+  );
+  expect(mockAddWarningToast).not.toHaveBeenCalled();
+
+  document.body.removeChild(dashboard);
+});
+
+test('falls through to clone path for dashboard export with multiple ag-grid charts', async () => {
+  const dashboard = document.createElement('div');
+  dashboard.className = 'dashboard';
+
+  for (let i = 0; i < 2; i += 1) {
+    const agContainer = document.createElement('div');
+    agContainer.setAttribute('data-themed-ag-grid', 'true');
+    const agRootWrapper = document.createElement('div');
+    agRootWrapper.className = 'ag-root-wrapper';
+    agContainer.appendChild(agRootWrapper);
+    (agContainer as any)._agGridFirstDataRendered = true;
+    dashboard.appendChild(agContainer);
+  }
+  document.body.appendChild(dashboard);
+
+  const handler = downloadAsImageOptimized('.dashboard', 'My Dashboard', true);
+  await handler({ currentTarget: {} } as any);
+
+  expect(mockToJpeg).toHaveBeenCalledWith(
+    expect.any(HTMLElement),
+    expect.objectContaining({ quality: 0.95 }),
+  );
+  expect(mockAddWarningToast).not.toHaveBeenCalled();
+
+  document.body.removeChild(dashboard);
+});
+
+test('captures JPEG for non-ag-grid elements via the clone path', async () => {
+  const container = document.createElement('div');
+  document.body.appendChild(container);
+
+  const handler = downloadAsImageOptimized('div', 'Bar Chart');
+  await handler(syntheticEventFor(container));
+
+  expect(mockToJpeg).toHaveBeenCalledWith(
+    expect.any(HTMLElement),
+    expect.objectContaining({ quality: 0.95 }),
+  );
+  expect(mockAddWarningToast).not.toHaveBeenCalled();
+
+  document.body.removeChild(container);
+});
+
+test('shows warning toast when clone capture throws', async () => {
+  mockToJpeg.mockRejectedValue(new Error('clone capture failed'));
+  const container = document.createElement('div');
+  document.body.appendChild(container);
+
+  const handler = downloadAsImageOptimized('div', 'Bar Chart');
+  await handler(syntheticEventFor(container));
+
+  expect(mockAddWarningToast).toHaveBeenCalledWith(
+    'Image download failed, please refresh and try again.',
+  );
+
+  document.body.removeChild(container);
+});
+
+test('ag-grid path uses theme colorBgContainer as background', async () => {
+  jest.useFakeTimers();
+  const { container, agContainer, cleanup } = buildAgGridElement();
+  attachMockApi(agContainer);
+
+  const theme = { colorBgContainer: '#1a1a2e' } as any;
+  const handler = downloadAsImageOptimized('div', 'My Chart', false, theme);
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  expect(mockToJpeg).toHaveBeenCalledWith(
+    expect.any(HTMLElement),
+    expect.objectContaining({ bgcolor: '#1a1a2e' }),
+  );
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('ag-grid path falls back to white background when theme is absent', async () => {
+  jest.useFakeTimers();
+  const { container, agContainer, cleanup } = buildAgGridElement();
+  attachMockApi(agContainer);
+
+  const handler = downloadAsImageOptimized('div', 'My Chart');
+  const exportPromise = handler(syntheticEventFor(container));
+  await jest.runAllTimersAsync();
+  await exportPromise;
+
+  expect(mockToJpeg).toHaveBeenCalledWith(
+    expect.any(HTMLElement),
+    expect.objectContaining({ bgcolor: undefined }),
+  );
+
+  cleanup();
+  jest.useRealTimers();
+});
+
+test('clone path falls back to white background when theme is absent', async () => {
+  const container = document.createElement('div');
+  document.body.appendChild(container);
+
+  const handler = downloadAsImageOptimized('div', 'Bar Chart');
+  await handler(syntheticEventFor(container));
+
+  expect(mockToJpeg).toHaveBeenCalledWith(
+    expect.any(HTMLElement),
+    expect.objectContaining({ bgcolor: undefined }),
+  );
+
+  document.body.removeChild(container);
+});

superset-frontend/src/utils/downloadAsImage.tsx

@@ -19,13 +19,17 @@
 import { SyntheticEvent } from 'react';
 import domToImage from 'dom-to-image-more';
 import { kebabCase } from 'lodash';
-// eslint-disable-next-line no-restricted-imports
 import { t } from '@apache-superset/core/translation';
 import { SupersetTheme } from '@apache-superset/core/theme';
 import { addWarningToast } from 'src/components/MessageToasts/actions';
+import type { AgGridContainerElement } from '@superset-ui/core/components';
 
 const IMAGE_DOWNLOAD_QUALITY = 0.95;
 const TRANSPARENT_RGBA = 'transparent';
+const POLL_INTERVAL_MS = 100;
+
+// Tracks original cell styles to restore after capture
+type CellFixup = { el: HTMLElement; minHeight: string; overflow: string };
 
 /**
  * generate a consistent file stem from a description and date
@@ -80,6 +84,9 @@ const CRITICAL_STYLE_PROPERTIES = new Set([
   'table-layout',
   'vertical-align',
   'text-align',
+  'box-sizing',
+  'min-height',
+  'min-width',
 ]);
 
 const styleCache = new WeakMap<Element, CSSStyleDeclaration>();
@@ -262,6 +269,45 @@ const createEnhancedClone = (
   return { clone, cleanup };
 };
 
+// Polls until scrollHeight is stable for minStablePolls consecutive intervals or maxMs elapses.
+// ag-grid has no "layout settled" event, so polling is the recommended workaround.
+export const waitForStableScrollHeight = (
+  el: HTMLElement,
+  maxMs = 5000,
+  minStablePolls = 2,
+): Promise<void> =>
+  new Promise<void>(resolve => {
+    const deadline = Date.now() + maxMs;
+    let lastHeight = el.scrollHeight;
+    let stableCount = 0;
+
+    const poll = () => {
+      if (Date.now() >= deadline) {
+        resolve();
+        return;
+      }
+      try {
+        const h = el.scrollHeight;
+        if (h === lastHeight) {
+          stableCount += 1;
+          if (stableCount >= minStablePolls) {
+            resolve();
+            return;
+          }
+        } else {
+          stableCount = 0;
+          lastHeight = h;
+        }
+      } catch {
+        resolve(); // element removed from DOM
+        return;
+      }
+      setTimeout(poll, POLL_INTERVAL_MS);
+    };
+
+    setTimeout(poll, POLL_INTERVAL_MS);
+  });
+
 export default function downloadAsImageOptimized(
   selector: string,
   description: string,
@@ -280,6 +326,139 @@ export default function downloadAsImageOptimized(
       return;
     }
 
+    const filter = (node: Element) =>
+      typeof node.className === 'string'
+        ? !node.className.includes('mapboxgl-control-container') &&
+          !node.className.includes('header-controls')
+        : true;
+
+    // Only apply ag-grid path for single-chart captures.
+    // Skip entirely for dashboard-level exports (selector targets the .dashboard root).
+    const isDashboardCapture = (
+      elementToPrint as HTMLElement
+    ).classList.contains('dashboard');
+    const agContainers = isDashboardCapture
+      ? []
+      : elementToPrint.querySelectorAll('[data-themed-ag-grid]');
+    const agContainer =
+      agContainers.length === 1
+        ? (agContainers[0] as AgGridContainerElement)
+        : null;
+    const agRootWrapper = agContainer
+      ? (agContainer.querySelector('.ag-root-wrapper') as HTMLElement | null)
+      : null;
+
+    if (agContainer && agRootWrapper) {
+      const api = agContainer._agGridApi;
+      const isFirstDataRendered =
+        agContainer._agGridFirstDataRendered === true;
+
+      if (!isFirstDataRendered) {
+        addWarningToast(
+          t('The chart is still loading. Please wait a moment and try again.'),
+        );
+        return;
+      }
+
+      // Capture resolved pixel widths before print layout can re-trigger sizeColumnsToFit.
+      // sizeColumnsToFit() sets flex (not pixel widths), so after print layout expands the
+      // container it recalculates column widths wider. We restore with flex: null to force
+      // pixel widths when calling applyColumnState after the layout switch.
+      const savedColumnState = api?.getColumnState?.();
+      const visibleColumnState =
+        savedColumnState?.filter(col => !col.hide) ?? [];
+      const originalWidth =
+        visibleColumnState.reduce((sum, col) => sum + (col.width ?? 0), 0) ||
+        agRootWrapper.offsetWidth;
+
+      // Chrome SVG foreignObject bug: % min-height resolves against canvas height,
+      // causing cells to expand to full image height and overlap adjacent rows.
+      const cellFixups: CellFixup[] = [];
+
+      try {
+        await document.fonts.ready;
+
+        if (api) {
+          api.setGridOption('domLayout', 'print');
+
+          // Wait for ResizeObserver + any triggered sizeColumnsToFit() to settle,
+          // then restore column widths before measurement.
+          await new Promise<void>(resolve =>
+            requestAnimationFrame(() => requestAnimationFrame(() => resolve())),
+          );
+
+          if (visibleColumnState.length > 0) {
+            api.applyColumnState?.({
+              state: visibleColumnState.map(col => ({
+                colId: col.colId,
+                width: col.width,
+                flex: null,
+              })),
+              applyOrder: false,
+            });
+          }
+
+          // Rows never scrolled into view have stale cached heights; remeasure all.
+          api.resetRowHeights?.();
+
+          // 5 polls × POLL_INTERVAL_MS = 500 ms; autoHeight rows batch-measure slowly.
+          await waitForStableScrollHeight(agRootWrapper, 5000, 5);
+        }
+
+        agRootWrapper.querySelectorAll('.ag-cell').forEach(cell => {
+          const el = cell as HTMLElement;
+          const rowHeight =
+            (el.parentElement as HTMLElement)?.offsetHeight ?? 0;
+          // scrollHeight catches any cells where resetRowHeights lagged behind.
+          const minH = Math.max(rowHeight, el.scrollHeight);
+          cellFixups.push({
+            el,
+            minHeight: el.style.minHeight,
+            overflow: el.style.overflow,
+          });
+          el.style.minHeight = minH > 0 ? `${minH}px` : '0px';
+          el.style.overflow = 'hidden';
+        });
+
+        const imageHeight = agRootWrapper.scrollHeight;
+
+        const dataUrl = await domToImage.toJpeg(agRootWrapper, {
+          bgcolor: theme?.colorBgContainer,
+          filter,
+          quality: IMAGE_DOWNLOAD_QUALITY,
+          height: imageHeight,
+          width: originalWidth,
+          cacheBust: true,
+        });
+
+        const link = document.createElement('a');
+        link.download = `${generateFileStem(description)}.jpg`;
+        link.href = dataUrl;
+        link.click();
+      } catch (error) {
+        console.error('Creating image failed', error);
+        addWarningToast(
+          t('Image download failed, please refresh and try again.'),
+        );
+      } finally {
+        cellFixups.forEach(({ el, minHeight, overflow }) => {
+          el.style.minHeight = minHeight;
+          el.style.overflow = overflow;
+        });
+        if (api) {
+          api.setGridOption('domLayout', 'normal');
+          if (savedColumnState) {
+            api.applyColumnState?.({
+              state: savedColumnState,
+              applyOrder: false,
+            });
+          }
+        }
+      }
+      return;
+    }
+
+    // All other chart types: use the clone-based approach
     let cleanup: (() => void) | null = null;
 
     try {
@@ -289,12 +468,6 @@ export default function downloadAsImageOptimized(
       );
       cleanup = cleanupFn;
 
-      const filter = (node: Element) =>
-        typeof node.className === 'string'
-          ? !node.className.includes('mapboxgl-control-container') &&
-            !node.className.includes('header-controls')
-          : true;
-
       const dataUrl = await domToImage.toJpeg(clone, {
         bgcolor: theme?.colorBgContainer,
         filter,