[ES|QL] More usages of getESQLQueryColumns (#182691)

## Summary This PR is an attempt to use the `getESQLQueryColumns` util everywhere we are trying to fetch the columns (using the limit 0). --------- Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
elastic · May 10, 2024 · 7054f58 · 7054f58
1 parent 37412c2
commit 7054f58
Show file tree

Hide file tree

Showing 17 changed files with 142 additions and 120 deletions.
diff --git a/.i18nrc.json b/.i18nrc.json
@@ -96,6 +96,7 @@
     "randomSampling": "x-pack/packages/kbn-random-sampling",
     "reactPackages": "packages/react",
     "textBasedEditor": "packages/kbn-text-based-editor",
+    "esqlUtils": "packages/kbn-esql-utils",
     "reporting": "packages/kbn-reporting",
     "savedObjects": "src/plugins/saved_objects",
     "savedObjectsFinder": "src/plugins/saved_objects_finder",

diff --git a/packages/kbn-esql-utils/index.ts b/packages/kbn-esql-utils/index.ts
@@ -17,6 +17,7 @@ export {
   appendToESQLQuery,
   appendWhereClauseToESQLQuery,
   getESQLQueryColumns,
+  getESQLQueryColumnsRaw,
   TextBasedLanguages,
 } from './src';
 

diff --git a/packages/kbn-esql-utils/src/index.ts b/packages/kbn-esql-utils/src/index.ts
@@ -16,4 +16,4 @@ export {
   removeDropCommandsFromESQLQuery,
 } from './utils/query_parsing_helpers';
 export { appendToESQLQuery, appendWhereClauseToESQLQuery } from './utils/append_to_query';
-export { getESQLQueryColumns } from './utils/run_query';
+export { getESQLQueryColumns, getESQLQueryColumnsRaw } from './utils/run_query';
diff --git a/packages/kbn-esql-utils/src/utils/run_query.ts b/packages/kbn-esql-utils/src/utils/run_query.ts
@@ -5,10 +5,11 @@
  * in compliance with, at your election, the Elastic License 2.0 or the Server
  * Side Public License, v 1.
  */
+import { i18n } from '@kbn/i18n';
 import type { DatatableColumn } from '@kbn/expressions-plugin/common';
 import type { ISearchGeneric } from '@kbn/search-types';
 import { esFieldTypeToKibanaFieldType } from '@kbn/field-types';
-import type { ESQLSearchReponse } from '@kbn/es-types';
+import type { ESQLColumn, ESQLSearchReponse } from '@kbn/es-types';
 import { lastValueFrom } from 'rxjs';
 import { ESQL_LATEST_VERSION } from '../../constants';
 
@@ -21,32 +22,77 @@ export async function getESQLQueryColumns({
   search: ISearchGeneric;
   signal?: AbortSignal;
 }): Promise<DatatableColumn[]> {
-  const response = await lastValueFrom(
-    search(
-      {
-        params: {
-          query: `${esqlQuery} | limit 0`,
-          version: ESQL_LATEST_VERSION,
+  try {
+    const response = await lastValueFrom(
+      search(
+        {
+          params: {
+            query: `${esqlQuery} | limit 0`,
+            version: ESQL_LATEST_VERSION,
+          },
         },
-      },
-      {
-        abortSignal: signal,
-        strategy: 'esql_async',
-      }
-    )
-  );
+        {
+          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;
+    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 column;
+      }) ?? [];
 
-  return columns;
+    return columns;
+  } catch (error) {
+    throw new Error(
+      i18n.translate('esqlUtils.columnsErrorMsg', {
+        defaultMessage: 'Unable to load columns. {errorMessage}',
+        values: { errorMessage: error.message },
+      })
+    );
+  }
+}
+
+export async function getESQLQueryColumnsRaw({
+  esqlQuery,
+  search,
+  signal,
+}: {
+  esqlQuery: string;
+  search: ISearchGeneric;
+  signal?: AbortSignal;
+}): Promise<ESQLColumn[]> {
+  try {
+    const response = await lastValueFrom(
+      search(
+        {
+          params: {
+            query: `${esqlQuery} | limit 0`,
+            version: ESQL_LATEST_VERSION,
+          },
+        },
+        {
+          abortSignal: signal,
+          strategy: 'esql_async',
+        }
+      )
+    );
+
+    return (response.rawResponse as unknown as ESQLSearchReponse).columns ?? [];
+  } catch (error) {
+    throw new Error(
+      i18n.translate('esqlUtils.columnsErrorMsg', {
+        defaultMessage: 'Unable to load columns. {errorMessage}',
+        values: { errorMessage: error.message },
+      })
+    );
+  }
 }
diff --git a/packages/kbn-esql-utils/tsconfig.json b/packages/kbn-esql-utils/tsconfig.json
@@ -23,6 +23,7 @@
     "@kbn/search-types",
     "@kbn/expressions-plugin",
     "@kbn/field-types",
-    "@kbn/es-types"
+    "@kbn/es-types",
+    "@kbn/i18n"
   ]
 }
diff --git a/x-pack/plugins/maps/public/classes/sources/esql_source/esql_utils.ts b/x-pack/plugins/maps/public/classes/sources/esql_source/esql_utils.ts
@@ -6,14 +6,13 @@
  */
 
 import { i18n } from '@kbn/i18n';
-import { lastValueFrom } from 'rxjs';
 import type { DataView } from '@kbn/data-plugin/common';
 import {
-  ESQL_LATEST_VERSION,
   getESQLAdHocDataview,
   getIndexPatternFromESQLQuery,
+  getESQLQueryColumnsRaw,
 } from '@kbn/esql-utils';
-import type { ESQLColumn, ESQLSearchReponse } from '@kbn/es-types';
+import type { ESQLColumn } from '@kbn/es-types';
 import { ES_GEO_FIELD_TYPE } from '../../../../common/constants';
 import { getData, getIndexPatternService } from '../../../kibana_services';
 
@@ -59,7 +58,10 @@ export async function getESQLMeta(esql: string) {
     getIndexPatternService()
   );
   return {
-    columns: await getColumns(esql),
+    columns: await getESQLQueryColumnsRaw({
+      esqlQuery: esql,
+      search: getData().search.search,
+    }),
     adhocDataViewId: adhocDataView.id!,
     ...getFields(adhocDataView),
   };
@@ -87,34 +89,6 @@ export function getFieldType(column: ESQLColumn) {
   }
 }
 
-async function getColumns(esql: string) {
-  const params = {
-    query: esql + ' | limit 0',
-    version: ESQL_LATEST_VERSION,
-  };
-
-  try {
-    const resp = await lastValueFrom(
-      getData().search.search(
-        { params },
-        {
-          strategy: 'esql',
-        }
-      )
-    );
-
-    const searchResponse = resp.rawResponse as unknown as ESQLSearchReponse;
-    return searchResponse.all_columns ? searchResponse.all_columns : searchResponse.columns;
-  } catch (error) {
-    throw new Error(
-      i18n.translate('xpack.maps.source.esql.getColumnsErrorMsg', {
-        defaultMessage: 'Unable to load columns. {errorMessage}',
-        values: { errorMessage: error.message },
-      })
-    );
-  }
-}
-
 export function getFields(dataView: DataView) {
   const dateFields: string[] = [];
   const geoFields: string[] = [];

diff --git a/...k/plugins/security_solution/public/detection_engine/rule_creation/logic/esql_validator.ts b/...k/plugins/security_solution/public/detection_engine/rule_creation/logic/esql_validator.ts
@@ -71,16 +71,16 @@ export const esqlValidator = async (
       };
     }
 
-    const data = await securitySolutionQueryClient.fetchQuery(
-      getEsqlQueryConfig({ esqlQuery: query, expressions: services.expressions })
+    const columns = await securitySolutionQueryClient.fetchQuery(
+      getEsqlQueryConfig({ esqlQuery: query, search: services.data.search.search })
     );
 
-    if (data && 'error' in data) {
-      return constructValidationError(data.error);
+    if (columns && 'error' in columns) {
+      return constructValidationError(columns.error);
     }
 
     // check whether _id field is present in response
-    const isIdFieldPresent = (data?.columns ?? []).find(({ id }) => '_id' === id);
+    const isIdFieldPresent = (columns ?? []).find(({ id }) => '_id' === id);
     // for non-aggregating query, we want to disable queries w/o _id property returned in response
     if (!isEsqlQueryAggregating && !isIdFieldPresent) {
       return {

diff --git a/...ns/security_solution/public/detection_engine/rule_creation/logic/get_esql_query_config.ts b/...ns/security_solution/public/detection_engine/rule_creation/logic/get_esql_query_config.ts
@@ -5,29 +5,31 @@
  * 2.0.
  */
 
-import { fetchFieldsFromESQL } from '@kbn/text-based-editor';
-import type { ExpressionsStart } from '@kbn/expressions-plugin/public';
+import { getESQLQueryColumns } from '@kbn/esql-utils';
+import type { ISearchGeneric } from '@kbn/search-types';
 
 /**
  * react-query configuration to be used to fetch ES|QL fields
  * it sets limit in query to 0, so we don't fetch unnecessary results, only fields
  */
 export const getEsqlQueryConfig = ({
   esqlQuery,
-  expressions,
+  search,
 }: {
   esqlQuery: string | undefined;
-  expressions: ExpressionsStart;
+  search: ISearchGeneric;
 }) => {
-  const emptyResultsEsqlQuery = `${esqlQuery} | limit 0`;
   return {
     queryKey: [(esqlQuery ?? '').trim()],
     queryFn: async () => {
       if (!esqlQuery) {
         return null;
       }
       try {
-        const res = await fetchFieldsFromESQL({ esql: emptyResultsEsqlQuery }, expressions);
+        const res = await getESQLQueryColumns({
+          esqlQuery,
+          search,
+        });
         return res;
       } catch (e) {
         return { error: e };

diff --git a/...tion_engine/rule_creation_ui/components/esql_autocomplete/use_esql_fields_options.test.ts b/...tion_engine/rule_creation_ui/components/esql_autocomplete/use_esql_fields_options.test.ts
@@ -19,17 +19,13 @@ describe('esqlToOptions', () => {
   });
   it('should transform all columns if fieldTYpe is not passed', () => {
     expect(
-      esqlToOptions({
-        type: 'datatable',
-        rows: [],
-        columns: [
-          { name: '@timestamp', id: '@timestamp', meta: { type: 'date' } },
-          { name: 'agent.build.original', id: 'agent.build.original', meta: { type: 'string' } },
-          { name: 'amqp.app-id', id: 'amqp.app-id', meta: { type: 'string' } },
-          { name: 'amqp.auto-delete', id: 'amqp.auto-delete', meta: { type: 'number' } },
-          { name: 'amqp.class-id', id: 'amqp.class-id', meta: { type: 'boolean' } },
-        ],
-      })
+      esqlToOptions([
+        { name: '@timestamp', id: '@timestamp', meta: { type: 'date' } },
+        { name: 'agent.build.original', id: 'agent.build.original', meta: { type: 'string' } },
+        { name: 'amqp.app-id', id: 'amqp.app-id', meta: { type: 'string' } },
+        { name: 'amqp.auto-delete', id: 'amqp.auto-delete', meta: { type: 'number' } },
+        { name: 'amqp.class-id', id: 'amqp.class-id', meta: { type: 'boolean' } },
+      ])
     ).toEqual([
       { label: '@timestamp' },
       { label: 'agent.build.original' },
@@ -41,17 +37,13 @@ describe('esqlToOptions', () => {
   it('should transform only columns of exact fieldType', () => {
     expect(
       esqlToOptions(
-        {
-          type: 'datatable',
-          rows: [],
-          columns: [
-            { name: '@timestamp', id: '@timestamp', meta: { type: 'date' } },
-            { name: 'agent.build.original', id: 'agent.build.original', meta: { type: 'string' } },
-            { name: 'amqp.app-id', id: 'amqp.app-id', meta: { type: 'string' } },
-            { name: 'amqp.auto-delete', id: 'amqp.auto-delete', meta: { type: 'number' } },
-            { name: 'amqp.class-id', id: 'amqp.class-id', meta: { type: 'boolean' } },
-          ],
-        },
+        [
+          { name: '@timestamp', id: '@timestamp', meta: { type: 'date' } },
+          { name: 'agent.build.original', id: 'agent.build.original', meta: { type: 'string' } },
+          { name: 'amqp.app-id', id: 'amqp.app-id', meta: { type: 'string' } },
+          { name: 'amqp.auto-delete', id: 'amqp.auto-delete', meta: { type: 'number' } },
+          { name: 'amqp.class-id', id: 'amqp.class-id', meta: { type: 'boolean' } },
+        ],
         'string'
       )
     ).toEqual([{ label: 'agent.build.original' }, { label: 'amqp.app-id' }]);

diff --git a/...detection_engine/rule_creation_ui/components/esql_autocomplete/use_esql_fields_options.ts b/...detection_engine/rule_creation_ui/components/esql_autocomplete/use_esql_fields_options.ts
@@ -6,7 +6,8 @@
  */
 import { useMemo } from 'react';
 import type { EuiComboBoxOptionOption } from '@elastic/eui';
-import type { Datatable, ExpressionsStart } from '@kbn/expressions-plugin/public';
+import type { DatatableColumn } from '@kbn/expressions-plugin/public';
+import type { DataPublicPluginStart } from '@kbn/data-plugin/public';
 
 import { useQuery } from '@tanstack/react-query';
 
@@ -16,14 +17,14 @@ import { getEsqlQueryConfig } from '../../../rule_creation/logic/get_esql_query_
 import type { FieldType } from '../../../rule_creation/logic/esql_validator';
 
 export const esqlToOptions = (
-  data: { error: unknown } | Datatable | undefined | null,
+  columns: { error: unknown } | DatatableColumn[] | undefined | null,
   fieldType?: FieldType
 ) => {
-  if (data && 'error' in data) {
+  if (columns && 'error' in columns) {
     return [];
   }
 
-  const options = (data?.columns ?? []).reduce<Array<{ label: string }>>((acc, { id, meta }) => {
+  const options = (columns ?? []).reduce<Array<{ label: string }>>((acc, { id, meta }) => {
     // if fieldType absent, we do not filter columns by type
     if (!fieldType || fieldType === meta.type) {
       acc.push({ label: id });
@@ -46,11 +47,11 @@ type UseEsqlFieldOptions = (
  * fetches ES|QL fields and convert them to Combobox options
  */
 export const useEsqlFieldOptions: UseEsqlFieldOptions = (esqlQuery, fieldType) => {
-  const kibana = useKibana<{ expressions: ExpressionsStart }>();
+  const kibana = useKibana<{ data: DataPublicPluginStart }>();
 
-  const { expressions } = kibana.services;
+  const { data: dataService } = kibana.services;
 
-  const queryConfig = getEsqlQueryConfig({ esqlQuery, expressions });
+  const queryConfig = getEsqlQueryConfig({ esqlQuery, search: dataService.search.search });
   const { data, isLoading } = useQuery(queryConfig);
 
   const options = useMemo(() => {

diff --git a/..._solution/public/detection_engine/rule_creation_ui/hooks/use_investigation_fields.test.ts b/..._solution/public/detection_engine/rule_creation_ui/hooks/use_investigation_fields.test.ts
@@ -13,18 +13,21 @@ import { useInvestigationFields } from './use_investigation_fields';
 import { createQueryWrapperMock } from '../../../common/__mocks__/query_wrapper';
 
 import { computeIsESQLQueryAggregating } from '@kbn/securitysolution-utils';
-import { fetchFieldsFromESQL } from '@kbn/text-based-editor';
+import { getESQLQueryColumns } from '@kbn/esql-utils';
 
 jest.mock('@kbn/securitysolution-utils', () => ({
   computeIsESQLQueryAggregating: jest.fn(),
 }));
 
-jest.mock('@kbn/text-based-editor', () => ({
-  fetchFieldsFromESQL: jest.fn(),
-}));
+jest.mock('@kbn/esql-utils', () => {
+  return {
+    getESQLQueryColumns: jest.fn(),
+    getIndexPatternFromESQLQuery: jest.fn().mockReturnValue('auditbeat*'),
+  };
+});
 
 const computeIsESQLQueryAggregatingMock = computeIsESQLQueryAggregating as jest.Mock;
-const fetchFieldsFromESQLMock = fetchFieldsFromESQL as jest.Mock;
+const getESQLQueryColumnsMock = getESQLQueryColumns as jest.Mock;
 
 const { wrapper } = createQueryWrapperMock();
 
@@ -48,7 +51,7 @@ const mockEsqlDatatable = {
 describe('useInvestigationFields', () => {
   beforeEach(() => {
     jest.clearAllMocks();
-    fetchFieldsFromESQLMock.mockResolvedValue(mockEsqlDatatable);
+    getESQLQueryColumnsMock.mockResolvedValue(mockEsqlDatatable.columns);
   });
 
   it('should return loading true when esql fields still loading', () => {