[ES|QL] Fetch the query columns utils (#181969)

## Summary

Adds a new utility for the users who want to retrieve the columns from a
query without expressions but using the search strategy.

This is the first utility to add for fetching ES|QL data without
expressions. This is only for columns but we can extend for fetching the
entire table instead. The latter will be part of
#179641 (comment)

---------

Co-authored-by: kibanamachine <[email protected]>

packages/kbn-esql-utils/index.ts

@@ -17,6 +17,7 @@ export {
   getESQLWithSafeLimit,
   appendToESQLQuery,
   TextBasedLanguages,
+  getESQLQueryColumns,
 } from './src';
 
 export { ESQL_LATEST_VERSION } from './constants';

packages/kbn-esql-utils/src/index.ts

@@ -17,3 +17,4 @@ export {
   removeDropCommandsFromESQLQuery,
 } from './utils/query_parsing_helpers';
 export { appendToESQLQuery } from './utils/append_to_query';
+export { getESQLQueryColumns } from './utils/run_query_utils';

packages/kbn-esql-utils/src/utils/run_query_utils.ts

@@ -0,0 +1,52 @@
+/*
+ * 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 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 or the Server
+ * Side Public License, v 1.
+ */
+import type { DatatableColumn } from '@kbn/expressions-plugin/common';
+import type { ISearchStart } from '@kbn/data-plugin/public';
+import { esFieldTypeToKibanaFieldType } from '@kbn/field-types';
+import type { ESQLSearchReponse } from '@kbn/es-types';
+import { lastValueFrom } from 'rxjs';
+import { ESQL_LATEST_VERSION } from '../../constants';
+
+export async function getESQLQueryColumns({
+  esqlQuery,
+  search,
+  signal,
+}: {
+  esqlQuery: string;
+  search: ISearchStart;
+  signal?: AbortSignal;
+}): Promise<DatatableColumn[]> {
+  const response = await lastValueFrom(
+    search.search(
+      {
+        params: {
+          query: `${esqlQuery} | limit 0`,
+          version: ESQL_LATEST_VERSION,
+        },
+      },
+      {
+        abortSignal: signal,
+        strategy: 'esql_async',
+      }
+    )
+  );
+
+  const columns =
+    (response.rawResponse as unknown as ESQLSearchReponse).columns?.map(({ name, type }) => {
+      const kibanaType = esFieldTypeToKibanaFieldType(type);
+      const column = {
+        id: name,
+        name,
+        meta: { type: kibanaType, esType: type },
+      } as DatatableColumn;
+
+      return column;
+    }) ?? [];
+
+  return columns;
+}

packages/kbn-esql-utils/tsconfig.json

@@ -17,8 +17,12 @@
   ],
   "kbn_references": [
     "@kbn/data-views-plugin",
+    "@kbn/data-plugin",
     "@kbn/crypto-browser",
     "@kbn/data-view-utils",
     "@kbn/esql-ast",
+    "@kbn/expressions-plugin",
+    "@kbn/field-types",
+    "@kbn/es-types",
   ]
 }

x-pack/plugins/lens/public/app_plugin/shared/edit_on_the_fly/helpers.test.ts

@@ -5,7 +5,7 @@
  * 2.0.
  */
 import { dataViewPluginMocks } from '@kbn/data-views-plugin/public/mocks';
-import { fetchFieldsFromESQL } from '@kbn/text-based-editor';
+import { getESQLQueryColumns } from '@kbn/esql-utils';
 import type { LensPluginStartDependencies } from '../../../plugin';
 import { createMockStartDependencies } from '../../../editor_frame_service/mocks';
 import {
@@ -18,45 +18,44 @@ import { suggestionsApi } from '../../../lens_suggestions_api';
 import { getSuggestions } from './helpers';
 
 const mockSuggestionApi = suggestionsApi as jest.Mock;
-const mockFetchData = fetchFieldsFromESQL as jest.Mock;
+const mockFetchData = getESQLQueryColumns as jest.Mock;
 
 jest.mock('../../../lens_suggestions_api', () => ({
   suggestionsApi: jest.fn(() => mockAllSuggestions),
 }));
 
-jest.mock('@kbn/text-based-editor', () => ({
-  fetchFieldsFromESQL: jest.fn(() => {
-    return {
-      columns: [
-        {
-          name: '@timestamp',
-          id: '@timestamp',
-          meta: {
-            type: 'date',
-          },
+jest.mock('@kbn/esql-utils', () => {
+  return {
+    getESQLQueryColumns: jest.fn().mockResolvedValue(() => [
+      {
+        name: '@timestamp',
+        id: '@timestamp',
+        meta: {
+          type: 'date',
         },
-        {
-          name: 'bytes',
-          id: 'bytes',
-          meta: {
-            type: 'number',
-          },
+      },
+      {
+        name: 'bytes',
+        id: 'bytes',
+        meta: {
+          type: 'number',
         },
-        {
-          name: 'memory',
-          id: 'memory',
-          meta: {
-            type: 'number',
-          },
+      },
+      {
+        name: 'memory',
+        id: 'memory',
+        meta: {
+          type: 'number',
         },
-      ],
-    };
-  }),
-}));
+      },
+    ]),
+    getIndexPatternFromESQLQuery: jest.fn().mockReturnValue('index1'),
+  };
+});
 
 describe('getSuggestions', () => {
   const query = {
-    esql: 'from index1 | limit 10 | stats average = avg(bytes',
+    esql: 'from index1 | limit 10 | stats average = avg(bytes)',
   };
   const mockStartDependencies =
     createMockStartDependencies() as unknown as LensPluginStartDependencies;

x-pack/plugins/lens/public/app_plugin/shared/edit_on_the_fly/helpers.ts

@@ -4,39 +4,20 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
-import { getIndexPatternFromSQLQuery, getIndexPatternFromESQLQuery } from '@kbn/esql-utils';
+import {
+  getIndexPatternFromSQLQuery,
+  getIndexPatternFromESQLQuery,
+  getESQLAdHocDataview,
+  getESQLQueryColumns,
+} from '@kbn/esql-utils';
 import type { AggregateQuery } from '@kbn/es-query';
-import { getESQLAdHocDataview } from '@kbn/esql-utils';
 import { getLensAttributesFromSuggestion } from '@kbn/visualization-utils';
-import { fetchFieldsFromESQL } from '@kbn/text-based-editor';
 import type { DataViewSpec } from '@kbn/data-views-plugin/public';
 import type { TypedLensByValueInput } from '../../../embeddable/embeddable_component';
 import type { LensPluginStartDependencies } from '../../../plugin';
 import type { DatasourceMap, VisualizationMap } from '../../../types';
 import { suggestionsApi } from '../../../lens_suggestions_api';
 
-export const getQueryColumns = async (
-  query: AggregateQuery,
-  deps: LensPluginStartDependencies,
-  abortController?: AbortController
-) => {
-  // Fetching only columns for ES|QL for performance reasons with limit 0
-  // Important note: ES doesnt return the warnings for 0 limit,
-  // I am skipping them in favor of performance now
-  // but we should think another way to get them (from Lens embeddable or store)
-  const performantQuery = { ...query };
-  if ('esql' in performantQuery && performantQuery.esql) {
-    performantQuery.esql = `${performantQuery.esql} | limit 0`;
-  }
-  const table = await fetchFieldsFromESQL(
-    performantQuery,
-    deps.expressions,
-    undefined,
-    abortController
-  );
-  return table?.columns;
-};
-
 export const getSuggestions = async (
   query: AggregateQuery,
   deps: LensPluginStartDependencies,
@@ -65,7 +46,12 @@ export const getSuggestions = async (
     if (dataView.fields.getByName('@timestamp')?.type === 'date' && !dataViewSpec) {
       dataView.timeFieldName = '@timestamp';
     }
-    const columns = await getQueryColumns(query, deps, abortController);
+
+    const columns = await getESQLQueryColumns({
+      esqlQuery: 'esql' in query ? query.esql : '',
+      search: deps.data.search,
+      signal: abortController?.signal,
+    });
     const context = {
       dataViewSpec: dataView?.toSpec(),
       fieldName: '',

x-pack/plugins/lens/public/editor_frame_service/mocks.tsx

@@ -55,7 +55,7 @@ export function createMockSetupDependencies() {
 
 export function createMockStartDependencies() {
   return {
-    data: dataPluginMock.createSetupContract(),
+    data: dataPluginMock.createStartContract(),
     embeddable: embeddablePluginMock.createStartContract(),
     expressions: expressionsPluginMock.createStartContract(),
     charts: chartPluginMock.createStartContract(),

x-pack/plugins/lens/public/trigger_actions/open_lens_config/create_action_helpers.ts

@@ -6,13 +6,13 @@
  */
 import { createGetterSetter } from '@kbn/kibana-utils-plugin/common';
 import type { CoreStart } from '@kbn/core/public';
+import { getESQLQueryColumns } from '@kbn/esql-utils';
 import { getLensAttributesFromSuggestion } from '@kbn/visualization-utils';
 import { IncompatibleActionError } from '@kbn/ui-actions-plugin/public';
 import { PresentationContainer } from '@kbn/presentation-containers';
 import { getESQLAdHocDataview, getIndexForESQLQuery } from '@kbn/esql-utils';
 import type { Datasource, Visualization } from '../../types';
 import type { LensPluginStartDependencies } from '../../plugin';
-import { fetchDataFromAggregateQuery } from '../../datasources/text_based/fetch_data_from_aggregate_query';
 import { suggestionsApi } from '../../lens_suggestions_api';
 import { generateId } from '../../id_generator';
 import { executeEditAction } from './edit_action_helpers';
@@ -66,21 +66,17 @@ export async function executeCreateAction({
   // so we are requesting them with limit 0
   // this is much more performant than requesting
   // all the table
-  const performantQuery = {
-    esql: `from ${defaultIndex} | limit 0`,
-  };
-
-  const table = await fetchDataFromAggregateQuery(
-    performantQuery,
-    dataView,
-    deps.data,
-    deps.expressions
-  );
+  const abortController = new AbortController();
+  const columns = await getESQLQueryColumns({
+    esqlQuery: `from ${defaultIndex}`,
+    search: deps.data.search,
+    signal: abortController.signal,
+  });
 
   const context = {
     dataViewSpec: dataView.toSpec(),
     fieldName: '',
-    textBasedColumns: table?.columns,
+    textBasedColumns: columns,
     query: defaultEsqlQuery,
   };