How to use the Dashboard API from a Splunk visualization (e.g. to reload the datasource)

DashboardApiContext.jsx

import React, { useState } from 'react';
import PropTypes from 'prop-types';

const DashboardApiContext = React.createContext();

const DashboardApiProvider = ({ children }) => {
    const [api, setApi] = useState();
    return (
        <DashboardApiContext.Provider
            value={{
                api,
                setApi
            }}
        >
            {children}
        </DashboardApiContext.Provider>
    );
};

DashboardApiProvider.propTypes = {
    children: PropTypes.node.isRequired,
};

const useDashboardApi = () => {
    const context = React.useContext(DashboardApiContext);
    if (context === undefined) {
        throw new Error('useDashboardApi must be used within a DashboardApiProvider');
    }
    return context;
};

export { DashboardApiProvider, useDashboardApi };

RefreshButton.jsx

import React, { useEffect } from 'react';
import { ActionButton } from '@splunk/dashboard-action-buttons';
import Refresh from '@splunk/dashboard-icons/Refresh';
import PropTypes from 'prop-types';
import { useDashboardApi } from './DashboardApiContext';

const RefreshButton = ({ itemId, dashboardApi, screenReaderText }) => {
    const { setApi } = useDashboardApi();
    useEffect(() => {
        setApi(dashboardApi);
    }, [dashboardApi, setApi]);

    const handleOnClick = () => {
        if (!dashboardApi || !itemId) {
            return;
        }
        dashboardApi.refreshVisualization(itemId);
    };

    return (
        <ActionButton
            data-test="RefreshButton" // for testing
            onClick={handleOnClick}
            icon={<Refresh screenReaderText={screenReaderText} />} // icon for this button
        />
    );
};

RefreshButton.propTypes = {
    itemId: PropTypes.string,
    dashboardApi: PropTypes.object,
    screenReaderText: PropTypes.string,
};

export default RefreshButton;

Usage.jsx

// 1. Add `RefreshButton` and `DashboardApiProvider` to your dashboard JSX file

<SplunkThemeProvider {...themeToVariant.prisma}>
    <DashboardContextProvider>
        <DashboardApiProvider>
            <DashboardCore
                width="100%"
                height="100%"
                preset={customPreset}
                definition={definition}
                actionMenus={[<RefreshButton key="refresh" />]}
            />
        </DashboardApiProvider>
    </DashboardContextProvider>
</SplunkThemeProvider>

// 2. Then you can use the Dashboard API from your visualization, e.g.

import React from 'react';
import Button from '@splunk/react-ui/Button';
import Table from '@splunk/visualizations/Table';
import { useDashboardApi } from './DashboardApiContext';

const Test = ({ id, dataSources }) => {
    const { api } = useDashboardApi();

    const handleOnClick = async () => {
        api.refreshVisualization(id);
    };

    return (
        <div>
            <Button onClick={handleOnClick}>Refresh</Button>
            <Table dataSources={dataSources} />
        </div>
    );
};