diff --git a/x-pack/platform/plugins/shared/lens/public/editor_frame_service/editor_frame/config_panel/esql_editor.tsx b/x-pack/platform/plugins/shared/lens/public/editor_frame_service/editor_frame/config_panel/esql_editor.tsx index cb3eaf9e7eff..db8d2f10436a 100644 --- a/x-pack/platform/plugins/shared/lens/public/editor_frame_service/editor_frame/config_panel/esql_editor.tsx +++ b/x-pack/platform/plugins/shared/lens/public/editor_frame_service/editor_frame/config_panel/esql_editor.tsx @@ -30,6 +30,7 @@ import { MAX_NUM_OF_COLUMNS } from '../../../datasources/form_based/esql_layer/u import { isApiESQLVariablesCompatible } from '../../../react_embeddable/types'; import type { LayerPanelProps } from './types'; import { ESQLDataGridAccordion } from '../../../app_plugin/shared/edit_on_the_fly/esql_data_grid_accordion'; +import { useInitializeChart } from './use_initialize_chart'; export type ESQLEditorProps = Simplify< { @@ -93,6 +94,7 @@ export function ESQLEditor({ const [dataGridAttrs, setDataGridAttrs] = useState(undefined); const [isSuggestionsAccordionOpen, setIsSuggestionsAccordionOpen] = useState(false); const [isESQLResultsAccordionOpen, setIsESQLResultsAccordionOpen] = useState(false); + const [isInitialized, setIsInitialized] = useState(false); const currentAttributes = useCurrentAttributes({ textBasedMode: isTextBasedLanguage, @@ -176,29 +178,17 @@ export function ESQLEditor({ ] ); - useEffect(() => { - const abortController = new AbortController(); - const initializeChart = async () => { - if (isTextBasedLanguage && isOfAggregateQueryType(query) && !dataGridAttrs) { - try { - await runQuery(query, abortController, Boolean(attributes?.state.needsRefresh)); - } catch (e) { - setErrors([e]); - prevQuery.current = query; - } - } - }; - initializeChart(); - }, [ - adHocDataViews, - runQuery, - esqlVariables, - query, - data, - dataGridAttrs, - attributes?.state.needsRefresh, + useInitializeChart({ isTextBasedLanguage, - ]); + query, + dataGridAttrs, + isInitialized, + currentAttributes, + runQuery, + prevQueryRef: prevQuery, + setErrors, + setIsInitialized, + }); // Early exit if it's not in TextBased mode if (!isTextBasedLanguage || !canEditTextBasedQuery || !isOfAggregateQueryType(query)) { diff --git a/x-pack/platform/plugins/shared/lens/public/editor_frame_service/editor_frame/config_panel/use_initialize_chart.test.ts b/x-pack/platform/plugins/shared/lens/public/editor_frame_service/editor_frame/config_panel/use_initialize_chart.test.ts new file mode 100644 index 000000000000..750b3cc227bc --- /dev/null +++ b/x-pack/platform/plugins/shared/lens/public/editor_frame_service/editor_frame/config_panel/use_initialize_chart.test.ts @@ -0,0 +1,137 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +import type { AggregateQuery, Query } from '@kbn/es-query'; +import { + createInitializeChartFunction, + type InitializeChartLogicArgs, +} from './use_initialize_chart'; +import { TypedLensSerializedState } from '../../../react_embeddable/types'; + +describe('createInitializeChartFunction', () => { + let mockSetErrors: jest.Mock; + let mockSetIsInitialized: jest.Mock; + let mockRunQuery: jest.Mock; + let mockPrevQueryRef: { current: AggregateQuery | Query }; + let defaultArgs: Parameters[0]; + + beforeEach(() => { + jest.clearAllMocks(); + mockSetErrors = jest.fn(); + mockSetIsInitialized = jest.fn(); + mockRunQuery = jest.fn(); + mockPrevQueryRef = { current: { esql: '' } as AggregateQuery }; + + defaultArgs = { + isTextBasedLanguage: true, + query: { esql: 'FROM my_data | limit 5' } as AggregateQuery, + dataGridAttrs: undefined, + isInitialized: false, + currentAttributes: { + state: { needsRefresh: false, query: { esql: '' } }, + } as TypedLensSerializedState['attributes'], // Minimal mock + prevQueryRef: mockPrevQueryRef, + setErrors: mockSetErrors, + setIsInitialized: mockSetIsInitialized, + runQuery: mockRunQuery, + }; + }); + + it('should call runQuery and set initialized to true if all conditions are met and not initialized', async () => { + const initializeChart = createInitializeChartFunction(defaultArgs); + await initializeChart(new AbortController()); + + expect(mockRunQuery).toHaveBeenCalledTimes(1); + expect(mockRunQuery).toHaveBeenCalledWith( + defaultArgs.query, + expect.any(AbortController), + false // needsRefresh is false by default in defaultArgs + ); + expect(mockPrevQueryRef.current).toEqual(defaultArgs.query); + expect(mockSetIsInitialized).toHaveBeenCalledTimes(1); + expect(mockSetIsInitialized).toHaveBeenCalledWith(true); + }); + + it('should NOT call mockRunQuery if already initialized', async () => { + const args = { ...defaultArgs, isInitialized: true }; + const initializeChart = createInitializeChartFunction(args); + await initializeChart(new AbortController()); + + expect(mockRunQuery).not.toHaveBeenCalled(); + expect(mockSetIsInitialized).toHaveBeenCalledTimes(0); + }); + + it('should NOT call mockRunQuery if isTextBasedLanguage is false', async () => { + const args = { ...defaultArgs, isTextBasedLanguage: false }; + const initializeChart = createInitializeChartFunction(args); + await initializeChart(new AbortController()); + + expect(mockRunQuery).not.toHaveBeenCalled(); + expect(mockSetIsInitialized).toHaveBeenCalledTimes(1); + expect(mockSetIsInitialized).toHaveBeenCalledWith(true); + }); + + it('should NOT call mockRunQuery if query is not of AggregateQueryType', async () => { + const args = { ...defaultArgs, query: { query: '', language: 'kuery' } }; // KQL query + const initializeChart = createInitializeChartFunction(args); + await initializeChart(new AbortController()); + + expect(mockRunQuery).not.toHaveBeenCalled(); + expect(mockSetIsInitialized).toHaveBeenCalledTimes(1); + expect(mockSetIsInitialized).toHaveBeenCalledWith(true); + }); + + it('should NOT call runQuery if dataGridAttrs is already defined', async () => { + const args = { + ...defaultArgs, + dataGridAttrs: { columns: [], rows: [] }, + } as unknown as InitializeChartLogicArgs; + const initializeChart = createInitializeChartFunction(args); + await initializeChart(new AbortController()); + + expect(mockRunQuery).not.toHaveBeenCalled(); + expect(mockSetIsInitialized).toHaveBeenCalledTimes(1); + expect(mockSetIsInitialized).toHaveBeenCalledWith(true); + }); + + it('should set errors and update prevQueryRef if runQuery throws an error', async () => { + const simulatedError = new Error('Failed to fetch data'); + (mockRunQuery as jest.Mock).mockRejectedValue(simulatedError); + + const initializeChart = createInitializeChartFunction(defaultArgs); + await initializeChart(new AbortController()); + + expect(mockRunQuery).toHaveBeenCalledTimes(1); + expect(mockSetErrors).toHaveBeenCalledWith([simulatedError]); + expect(mockPrevQueryRef.current).toEqual(defaultArgs.query); + expect(mockSetIsInitialized).toHaveBeenCalledTimes(1); + expect(mockSetIsInitialized).toHaveBeenCalledWith(true); + }); + + it('should pass needsRefresh to runQuery if currentAttributes.state.needsRefresh is true', async () => { + const args = { + ...defaultArgs, + currentAttributes: { + state: { needsRefresh: true, query: { esql: '' } }, + } as TypedLensSerializedState['attributes'], + }; + const initializeChart = createInitializeChartFunction(args); + await initializeChart(new AbortController()); + + expect(mockRunQuery).toHaveBeenCalledTimes(1); + }); + + it('should set initialized to true even if runQuery is not called due to conditions', async () => { + // Test a case where runQuery is not called + const args = { ...defaultArgs, isTextBasedLanguage: false }; + const initializeChart = createInitializeChartFunction(args); + await initializeChart(new AbortController()); + + expect(mockRunQuery).not.toHaveBeenCalled(); + expect(mockSetIsInitialized).toHaveBeenCalledTimes(1); + expect(mockSetIsInitialized).toHaveBeenCalledWith(true); + }); +}); diff --git a/x-pack/platform/plugins/shared/lens/public/editor_frame_service/editor_frame/config_panel/use_initialize_chart.ts b/x-pack/platform/plugins/shared/lens/public/editor_frame_service/editor_frame/config_panel/use_initialize_chart.ts new file mode 100644 index 000000000000..8289e0f6605f --- /dev/null +++ b/x-pack/platform/plugins/shared/lens/public/editor_frame_service/editor_frame/config_panel/use_initialize_chart.ts @@ -0,0 +1,130 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +import { useEffect, useCallback, type MutableRefObject } from 'react'; +import { type AggregateQuery, isOfAggregateQueryType, type Query } from '@kbn/es-query'; +import { type ESQLDataGridAttrs } from '../../../app_plugin/shared/edit_on_the_fly/helpers'; +import { TypedLensSerializedState } from '../../../react_embeddable/types'; + +type LensAttributes = TypedLensSerializedState['attributes']; + +export interface InitializeChartLogicArgs { + /** + * Indicates if the query is in text-based language (ESQL). + */ + isTextBasedLanguage: boolean; + /** + * The query to be executed. + */ + query: AggregateQuery | Query; + /** + * Attributes for the ESQL data grid, if applicable. + */ + dataGridAttrs: ESQLDataGridAttrs | undefined; + /** + * Indicates if the dataGridAttrs havw been initialized. + */ + isInitialized: boolean; + /** + * Current attributes of the chart. + */ + currentAttributes: LensAttributes | undefined; + /** + * Reference to the previous query. + */ + prevQueryRef: MutableRefObject; + /** + * Function to set errors that occur during initialization. + */ + setErrors: (errors: Error[]) => void; + /** + * Function to set the initialization state. + */ + setIsInitialized: (isInitialized: boolean) => void; + /** + * Function to run the query and update the chart. + */ + runQuery: ( + q: AggregateQuery, + abortController?: AbortController, + shouldUpdateAttrs?: boolean + ) => Promise; +} + +/** + * Encapsulates the logic for initializing the chart/data grid based on ESQL query. + * + */ +export const createInitializeChartFunction = ({ + isTextBasedLanguage, + query, + dataGridAttrs, + isInitialized, + currentAttributes, + runQuery, + prevQueryRef, + setErrors, + setIsInitialized, +}: InitializeChartLogicArgs) => { + return async (abortController?: AbortController) => { + if (isInitialized) { + // If already initialized, do nothing + return; + } + if (isTextBasedLanguage && isOfAggregateQueryType(query) && !dataGridAttrs) { + try { + const shouldUpdateAttrs = Boolean(currentAttributes?.state.needsRefresh); + await runQuery(query, abortController, shouldUpdateAttrs); + } catch (e) { + setErrors([e]); + } + prevQueryRef.current = query; + } + setIsInitialized(true); + }; +}; + +export function useInitializeChart({ + isTextBasedLanguage, + query, + dataGridAttrs, + isInitialized, + currentAttributes, + runQuery, + prevQueryRef, + setErrors, + setIsInitialized, +}: InitializeChartLogicArgs) { + const initializeChartFunc = useCallback(() => { + const abortController = new AbortController(); + + const func = createInitializeChartFunction({ + isTextBasedLanguage, + query, + dataGridAttrs, + isInitialized, + currentAttributes, + runQuery, + prevQueryRef, + setErrors, + setIsInitialized, + }); + func(abortController); + }, [ + isTextBasedLanguage, + query, + dataGridAttrs, + isInitialized, + currentAttributes, + runQuery, + prevQueryRef, + setErrors, + setIsInitialized, + ]); + useEffect(() => { + initializeChartFunc(); + }, [initializeChartFunc]); +}