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

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';