[ES|QL Controls] Refresh "Values from a query" options on dashboard reload

src/platform/plugins/shared/controls/public/controls/esql_control/esql_control_selections.ts

@@ -7,10 +7,17 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 import deepEqual from 'react-fast-compare';
-import { BehaviorSubject, combineLatest, map, merge } from 'rxjs';
-import type { ESQLControlVariable, ESQLControlState, EsqlControlType } from '@kbn/esql-types';
-import { ESQLVariableType } from '@kbn/esql-types';
+import { BehaviorSubject, combineLatest, filter, map, merge, switchMap } from 'rxjs';
+import {
+  ESQLControlVariable,
+  ESQLControlState,
+  EsqlControlType,
+  ESQLVariableType,
+} from '@kbn/esql-types';
 import { PublishingSubject, StateComparators } from '@kbn/presentation-publishing';
+import { dataService } from '../../services/kibana_services';
+import { ControlGroupApi } from '../../control_group/types';
+import { getESQLSingleColumnValues } from './utils/get_esql_single_column_values';
 
 function selectedOptionsComparatorFunction(a?: string[], b?: string[]) {
   return deepEqual(a ?? [], b ?? []);
@@ -37,7 +44,10 @@ export const selectionComparators: StateComparators<
   title: 'referenceEquality',
 };
 
-export function initializeESQLControlSelections(initialState: ESQLControlState) {
+export function initializeESQLControlSelections(
+  initialState: ESQLControlState,
+  controlFetch$: ReturnType<ControlGroupApi['controlFetch$']>
+) {
   const availableOptions$ = new BehaviorSubject<string[]>(initialState.availableOptions ?? []);
   const selectedOptions$ = new BehaviorSubject<string[]>(initialState.selectedOptions ?? []);
   const hasSelections$ = new BehaviorSubject<boolean>(false); // hardcoded to false to prevent clear action from appearing.
@@ -55,6 +65,25 @@ export function initializeESQLControlSelections(initialState: ESQLControlState)
     }
   }
 
+  // For Values From Query controls, update values on dashboard load/reload
+  const fetchSubscription = controlFetch$
+    .pipe(
+      filter(() => controlType$.getValue() === EsqlControlType.VALUES_FROM_QUERY),
+      switchMap(
+        async ({ timeRange }) =>
+          await getESQLSingleColumnValues({
+            query: esqlQuery$.getValue(),
+            search: dataService.search.search,
+            timeRange,
+          })
+      )
+    )
+    .subscribe((result) => {
+      if (getESQLSingleColumnValues.isSuccess(result)) {
+        availableOptions$.next(result.values);
+      }
+    });
+
   // derive ESQL control variable from state.
   const getEsqlVariable = () => ({
     key: variableName$.value,
@@ -64,12 +93,18 @@ export function initializeESQLControlSelections(initialState: ESQLControlState)
     type: variableType$.value,
   });
   const esqlVariable$ = new BehaviorSubject<ESQLControlVariable>(getEsqlVariable());
-  const subscriptions = combineLatest([variableName$, variableType$, selectedOptions$]).subscribe(
-    () => esqlVariable$.next(getEsqlVariable())
-  );
+  const variableSubscriptions = combineLatest([
+    variableName$,
+    variableType$,
+    selectedOptions$,
+    availableOptions$,
+  ]).subscribe(() => esqlVariable$.next(getEsqlVariable()));
 
   return {
-    cleanup: () => subscriptions.unsubscribe(),
+    cleanup: () => {
+      variableSubscriptions.unsubscribe();
+      fetchSubscription.unsubscribe();
+    },
     api: {
       hasSelections$: hasSelections$ as PublishingSubject<boolean | undefined>,
       esqlVariable$: esqlVariable$ as PublishingSubject<ESQLControlVariable>,

src/platform/plugins/shared/controls/public/controls/esql_control/get_esql_control_factory.test.tsx

@@ -9,19 +9,36 @@
 
 import React from 'react';
 import { fireEvent, render, waitFor } from '@testing-library/react';
-import type { ESQLControlState } from '@kbn/esql-types';
+import { EsqlControlType, type ESQLControlState } from '@kbn/esql-types';
 import { getMockedControlGroupApi, getMockedFinalizeApi } from '../mocks/control_mocks';
 import { getESQLControlFactory } from './get_esql_control_factory';
+import { BehaviorSubject } from 'rxjs';
+import { ControlFetchContext } from '../../control_group/control_fetch';
+
+const mockGetESQLSingleColumnValues = jest.fn(() => ({ options: ['option1', 'option2'] }));
+const mockIsSuccess = jest.fn(() => true);
+
+jest.mock('./utils/get_esql_single_column_values', () => {
+  const getESQLSingleColumnValues = () => mockGetESQLSingleColumnValues();
+  getESQLSingleColumnValues.isSuccess = () => mockIsSuccess();
+  return {
+    getESQLSingleColumnValues,
+  };
+});
 
 describe('ESQLControlApi', () => {
+  beforeEach(() => {
+    jest.resetAllMocks();
+  });
+
   const uuid = 'myESQLControl';
 
   const dashboardApi = {};
   const controlGroupApi = getMockedControlGroupApi(dashboardApi);
   const factory = getESQLControlFactory();
   const finalizeApi = getMockedFinalizeApi(uuid, factory, controlGroupApi);
 
-  test('Should publish ES|QL variable', async () => {
+  test('should publish ES|QL variable', async () => {
     const initialState = {
       selectedOptions: ['option1'],
       availableOptions: ['option1', 'option2'],
@@ -43,7 +60,7 @@ describe('ESQLControlApi', () => {
     });
   });
 
-  test('Should serialize state', async () => {
+  test('should serialize state', async () => {
     const initialState = {
       selectedOptions: ['option1'],
       availableOptions: ['option1', 'option2'],
@@ -74,37 +91,70 @@ describe('ESQLControlApi', () => {
     });
   });
 
-  test('changing the dropdown should publish new ES|QL variable', async () => {
-    const initialState = {
-      selectedOptions: ['option1'],
-      availableOptions: ['option1', 'option2'],
-      variableName: 'variable1',
-      variableType: 'values',
-      esqlQuery: 'FROM foo | WHERE column = ?variable1',
-      controlType: 'STATIC_VALUES',
-    } as ESQLControlState;
-    const { Component, api } = await factory.buildControl({
-      initialState,
-      finalizeApi,
-      uuid,
-      controlGroupApi,
-    });
-
-    expect(api.esqlVariable$.value).toStrictEqual({
-      key: 'variable1',
-      type: 'values',
-      value: 'option1',
+  describe('values from query', () => {
+    test('should update on load and fetch', async () => {
+      const initialState = {
+        selectedOptions: ['option1'],
+        availableOptions: ['option1', 'option2'],
+        variableName: 'variable1',
+        variableType: 'values',
+        esqlQuery: 'FROM foo | STATS BY column',
+        controlType: EsqlControlType.VALUES_FROM_QUERY,
+      } as ESQLControlState;
+      await factory.buildControl({
+        initialState,
+        finalizeApi,
+        uuid,
+        controlGroupApi,
+      });
+      await waitFor(() => {
+        expect(mockGetESQLSingleColumnValues).toHaveBeenCalledTimes(1);
+        expect(mockIsSuccess).toHaveBeenCalledTimes(1);
+      });
+      const controlFetch$ = controlGroupApi.controlFetch$(
+        uuid
+      ) as BehaviorSubject<ControlFetchContext>;
+      controlFetch$.next({});
+      await waitFor(() => {
+        expect(mockGetESQLSingleColumnValues).toHaveBeenCalledTimes(2);
+        expect(mockIsSuccess).toHaveBeenCalledTimes(2);
+      });
     });
+  });
 
-    const { findByTestId, findByTitle } = render(<Component className="" />);
-    fireEvent.click(await findByTestId('comboBoxSearchInput'));
-    fireEvent.click(await findByTitle('option2'));
+  describe('changing the dropdown', () => {
+    test('should publish new ES|QL variable', async () => {
+      const initialState = {
+        selectedOptions: ['option1'],
+        availableOptions: ['option1', 'option2'],
+        variableName: 'variable1',
+        variableType: 'values',
+        esqlQuery: 'FROM foo | WHERE column = ?variable1',
+        controlType: 'STATIC_VALUES',
+      } as ESQLControlState;
+      const { Component, api } = await factory.buildControl({
+        initialState,
+        finalizeApi,
+        uuid,
+        controlGroupApi,
+      });
 
-    await waitFor(() => {
       expect(api.esqlVariable$.value).toStrictEqual({
         key: 'variable1',
         type: 'values',
-        value: 'option2',
+        value: 'option1',
+      });
+
+      const { findByTestId, findByTitle } = render(<Component className="" />);
+      fireEvent.click(await findByTestId('comboBoxSearchInput'));
+      fireEvent.click(await findByTitle('option2'));
+
+      await waitFor(() => {
+        expect(api.esqlVariable$.value).toStrictEqual({
+          key: 'variable1',
+          type: 'values',
+          value: 'option2',
+        });
       });
     });
   });

src/platform/plugins/shared/controls/public/controls/esql_control/get_esql_control_factory.tsx

@@ -37,7 +37,10 @@ export const getESQLControlFactory = (): ControlFactory<ESQLControlState, ESQLCo
     getDisplayName: () => displayName,
     buildControl: async ({ initialState, finalizeApi, uuid, controlGroupApi }) => {
       const defaultControlManager = initializeDefaultControlManager(initialState);
-      const selections = initializeESQLControlSelections(initialState);
+      const selections = initializeESQLControlSelections(
+        initialState,
+        controlGroupApi.controlFetch$(uuid)
+      );
 
       const closeOverlay = () => {
         if (apiHasParentApi(controlGroupApi) && tracksOverlays(controlGroupApi.parentApi)) {

src/platform/plugins/shared/controls/public/controls/esql_control/utils/get_esql_single_column_values.test.ts

@@ -0,0 +1,99 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+import { ISearchGeneric } from '@kbn/search-types';
+import {
+  getESQLSingleColumnValues,
+  GetESQLSingleColumnValuesSuccess,
+  GetESQLSingleColumnValuesFailure,
+} from './get_esql_single_column_values';
+
+const mockGetESQLResults = jest.fn();
+jest.mock('@kbn/esql-utils', () => ({
+  getESQLResults: (...args: any[]) => mockGetESQLResults(...args),
+}));
+
+const searchMock = {} as ISearchGeneric;
+
+describe('getESQLSingleColumnValues', () => {
+  beforeEach(() => {
+    jest.resetAllMocks();
+  });
+  it('returns only options on success', async () => {
+    mockGetESQLResults.mockResolvedValueOnce({
+      response: {
+        columns: [{ name: 'column1' }],
+        values: [['option1'], ['option2']],
+      },
+    });
+    const result = (await getESQLSingleColumnValues({
+      query: 'FROM index | STATS BY column',
+      search: searchMock,
+    })) as GetESQLSingleColumnValuesSuccess;
+    expect(getESQLSingleColumnValues.isSuccess(result)).toBe(true);
+    expect(result).toMatchInlineSnapshot(`
+      Object {
+        "values": Array [
+          "option1",
+          "option2",
+        ],
+      }
+    `);
+  });
+  it('returns an error when query returns multiple columns', async () => {
+    mockGetESQLResults.mockResolvedValueOnce({
+      response: {
+        columns: [{ name: 'column1' }, { name: 'column2' }],
+        values: [['option1'], ['option2']],
+      },
+    });
+    const result = (await getESQLSingleColumnValues({
+      query: 'FROM index',
+      search: searchMock,
+    })) as GetESQLSingleColumnValuesFailure;
+    expect(getESQLSingleColumnValues.isSuccess(result)).toBe(false);
+    expect('values' in result).toBe(false);
+    expect(result).toMatchInlineSnapshot(`
+      Object {
+        "errors": Array [
+          [Error: Query must return a single column],
+        ],
+      }
+    `);
+  });
+  it('returns an error on a failed query', async () => {
+    mockGetESQLResults.mockRejectedValueOnce('Invalid ES|QL query');
+    const result = (await getESQLSingleColumnValues({
+      query: 'FROM index | EVAL',
+      search: searchMock,
+    })) as GetESQLSingleColumnValuesFailure;
+    expect(getESQLSingleColumnValues.isSuccess(result)).toBe(false);
+    expect('values' in result).toBe(false);
+    expect(result).toMatchInlineSnapshot(`
+      Object {
+        "errors": Array [
+          "Invalid ES|QL query",
+        ],
+      }
+    `);
+  });
+
+  it('passes timeRange successfully', async () => {
+    const timeRange = { from: 'now-10m', to: 'now' };
+    await getESQLSingleColumnValues({
+      query: 'FROM index | STATS BY column',
+      search: searchMock,
+      timeRange,
+    });
+    expect(mockGetESQLResults).toHaveBeenCalledWith(
+      expect.objectContaining({
+        timeRange,
+      })
+    );
+  });
+});