[Security Solution][Detection Engine] improves ES|QL investigation fields for detection rules

x-pack/plugins/security_solution/public/common/__mocks__/query_wrapper.tsx

@@ -0,0 +1,34 @@
+/*
+ * 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 React from 'react';
+import { QueryClient, QueryClientProvider } from '@tanstack/react-query';
+
+export const createQueryWrapperMock = (): {
+  queryClient: QueryClient;
+  wrapper: React.FC;
+} => {
+  const queryClient = new QueryClient({
+    defaultOptions: {
+      queries: {
+        retry: false,
+      },
+    },
+    logger: {
+      error: () => undefined,
+      log: () => undefined,
+      warn: () => undefined,
+    },
+  });
+
+  return {
+    queryClient,
+    wrapper: ({ children }) => (
+      <QueryClientProvider client={queryClient}>{children}</QueryClientProvider>
+    ),
+  };
+};

x-pack/plugins/security_solution/public/detection_engine/rule_creation/components/esql_autocomplete/use_esql_fields_options.ts → x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/esql_autocomplete/use_esql_fields_options.ts

@@ -12,8 +12,8 @@ import { useQuery } from '@tanstack/react-query';
 
 import { useKibana } from '@kbn/kibana-react-plugin/public';
 
-import { getEsqlQueryConfig } from '../../logic/get_esql_query_config';
-import type { FieldType } from '../../logic/esql_validator';
+import { getEsqlQueryConfig } from '../../../rule_creation/logic/get_esql_query_config';
+import type { FieldType } from '../../../rule_creation/logic/esql_validator';
 
 export const esqlToOptions = (
   data: { error: unknown } | Datatable | undefined | null,

x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_about_rule/index.tsx

@@ -37,8 +37,9 @@ import { useFetchIndex } from '../../../../common/containers/source';
 import { DEFAULT_INDICATOR_SOURCE_PATH } from '../../../../../common/constants';
 import { useKibana } from '../../../../common/lib/kibana';
 import { useRuleIndices } from '../../../rule_management/logic/use_rule_indices';
-import { EsqlAutocomplete } from '../../../rule_creation/components/esql_autocomplete';
+import { EsqlAutocomplete } from '../esql_autocomplete';
 import { MultiSelectFieldsAutocomplete } from '../multi_select_fields';
+import { useInvestigationFields } from '../../hooks/use_investigation_fields';
 
 const CommonUseField = getUseField({ component: Field });
 
@@ -128,6 +129,11 @@ const StepAboutRuleComponent: FC<StepAboutRuleProps> = ({
     [getFields]
   );
 
+  const { investigationFields, isLoading: isInvestigationFieldsLoading } = useInvestigationFields({
+    esqlQuery: isEsqlRuleValue ? esqlQuery : undefined,
+    indexPatternsFields: indexPattern.fields,
+  });
+
   return (
     <>
       <StepContentWrapper addPadding={!isUpdateView}>
@@ -240,8 +246,8 @@ const StepAboutRuleComponent: FC<StepAboutRuleProps> = ({
               path="investigationFields"
               component={MultiSelectFieldsAutocomplete}
               componentProps={{
-                browserFields: indexPattern.fields,
-                isDisabled: isLoading || indexPatternLoading,
+                browserFields: investigationFields,
+                isDisabled: isLoading || indexPatternLoading || isInvestigationFieldsLoading,
                 fullWidth: true,
                 dataTestSubj: 'detectionEngineStepAboutRuleInvestigationFields',
               }}

x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/hooks/use_investigation_fields.test.ts

@@ -0,0 +1,135 @@
+/*
+ * 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 { renderHook } from '@testing-library/react-hooks';
+import type { DataViewFieldBase } from '@kbn/es-query';
+
+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';
+
+jest.mock('@kbn/securitysolution-utils', () => ({
+  computeIsESQLQueryAggregating: jest.fn(),
+}));
+
+jest.mock('@kbn/text-based-editor', () => ({
+  fetchFieldsFromESQL: jest.fn(),
+}));
+
+const computeIsESQLQueryAggregatingMock = computeIsESQLQueryAggregating as jest.Mock;
+const fetchFieldsFromESQLMock = fetchFieldsFromESQL as jest.Mock;
+
+const { wrapper } = createQueryWrapperMock();
+
+const mockEsqlQuery = 'from auditbeat* [metadata _id]';
+const mockIndexPatternFields: DataViewFieldBase[] = [
+  {
+    name: 'agent.name',
+    type: 'string',
+  },
+  {
+    name: 'agent.type',
+    type: 'string',
+  },
+];
+const mockEsqlDatatable = {
+  type: 'datatable',
+  rows: [],
+  columns: [{ id: '_custom_field', name: '_custom_field', meta: { type: 'string' } }],
+};
+
+describe('useInvestigationFields', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    fetchFieldsFromESQLMock.mockResolvedValue(mockEsqlDatatable);
+  });
+
+  it('should return loading true when esql fields still loading', () => {
+    const { result } = renderHook(
+      () =>
+        useInvestigationFields({
+          esqlQuery: mockEsqlQuery,
+          indexPatternsFields: mockIndexPatternFields,
+        }),
+      { wrapper }
+    );
+
+    expect(result.current.isLoading).toBe(true);
+  });
+
+  it('should return only index pattern fields when ES|QL query is empty', async () => {
+    const { result, waitForNextUpdate } = renderHook(
+      () =>
+        useInvestigationFields({
+          esqlQuery: '',
+          indexPatternsFields: mockIndexPatternFields,
+        }),
+      { wrapper }
+    );
+
+    await waitForNextUpdate();
+
+    expect(result.current.investigationFields).toEqual(mockIndexPatternFields);
+  });
+
+  it('should return only index pattern fields when ES|QL query is undefined', async () => {
+    const { result } = renderHook(
+      () =>
+        useInvestigationFields({
+          esqlQuery: undefined,
+          indexPatternsFields: mockIndexPatternFields,
+        }),
+      { wrapper }
+    );
+
+    expect(result.current.investigationFields).toEqual(mockIndexPatternFields);
+  });
+
+  it('should return index pattern fields concatenated with ES|QL fields when ES|QL query is non-aggregating', async () => {
+    computeIsESQLQueryAggregatingMock.mockReturnValue(false);
+
+    const { result } = renderHook(
+      () =>
+        useInvestigationFields({
+          esqlQuery: mockEsqlQuery,
+          indexPatternsFields: mockIndexPatternFields,
+        }),
+      { wrapper }
+    );
+
+    expect(result.current.investigationFields).toEqual([
+      {
+        name: '_custom_field',
+        type: 'string',
+      },
+      ...mockIndexPatternFields,
+    ]);
+  });
+
+  it('should return only ES|QL fields when ES|QL query is aggregating', async () => {
+    computeIsESQLQueryAggregatingMock.mockReturnValue(true);
+
+    const { result } = renderHook(
+      () =>
+        useInvestigationFields({
+          esqlQuery: mockEsqlQuery,
+          indexPatternsFields: mockIndexPatternFields,
+        }),
+      { wrapper }
+    );
+
+    expect(result.current.investigationFields).toEqual([
+      {
+        name: '_custom_field',
+        type: 'string',
+      },
+    ]);
+  });
+});

x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/hooks/use_investigation_fields.ts

@@ -0,0 +1,91 @@
+/*
+ * 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 { useMemo } from 'react';
+import type { Datatable, ExpressionsStart } from '@kbn/expressions-plugin/public';
+import type { DataViewFieldBase } from '@kbn/es-query';
+import { computeIsESQLQueryAggregating } from '@kbn/securitysolution-utils';
+
+import { useQuery } from '@tanstack/react-query';
+
+import { useKibana } from '@kbn/kibana-react-plugin/public';
+
+import { getEsqlQueryConfig } from '../../rule_creation/logic/get_esql_query_config';
+
+const esqlToFields = (
+  data: { error: unknown } | Datatable | undefined | null
+): DataViewFieldBase[] => {
+  if (data && 'error' in data) {
+    return [];
+  }
+
+  const fields = (data?.columns ?? []).map(({ id, meta }) => {
+    return {
+      name: id,
+      type: meta.type,
+    };
+  });
+
+  return fields;
+};
+
+type UseEsqlFields = (esqlQuery: string | undefined) => {
+  isLoading: boolean;
+  fields: DataViewFieldBase[];
+};
+
+/**
+ * fetches ES|QL fields and convert them to DataViewBase fields
+ */
+const useEsqlFields: UseEsqlFields = (esqlQuery) => {
+  const kibana = useKibana<{ expressions: ExpressionsStart }>();
+
+  const { expressions } = kibana.services;
+
+  const queryConfig = getEsqlQueryConfig({ esqlQuery, expressions });
+  const { data, isLoading } = useQuery(queryConfig);
+
+  const fields = useMemo(() => {
+    return esqlToFields(data);
+  }, [data]);
+
+  return {
+    fields,
+    isLoading,
+  };
+};
+
+type UseInvestigationFields = (params: {
+  esqlQuery: string | undefined;
+  indexPatternsFields: DataViewFieldBase[];
+}) => {
+  isLoading: boolean;
+  investigationFields: DataViewFieldBase[];
+};
+
+export const useInvestigationFields: UseInvestigationFields = ({
+  esqlQuery,
+  indexPatternsFields,
+}) => {
+  const { fields: esqlFields, isLoading } = useEsqlFields(esqlQuery);
+
+  const investigationFields = useMemo(() => {
+    if (!esqlQuery) {
+      return indexPatternsFields;
+    }
+
+    // alerts generated from non-aggregating queries are enriched with source document
+    // so, index patterns fields should be included in the list of investigation fields
+    const isEsqlQueryAggregating = computeIsESQLQueryAggregating(esqlQuery);
+
+    return isEsqlQueryAggregating ? esqlFields : [...esqlFields, ...indexPatternsFields];
+  }, [esqlFields, esqlQuery, indexPatternsFields]);
+
+  return {
+    investigationFields,
+    isLoading,
+  };
+};

x-pack/test/security_solution_cypress/cypress/e2e/detection_response/detection_engine/rule_creation/esql_rule_ess.cy.ts

@@ -7,7 +7,11 @@
 
 import { getEsqlRule } from '../../../../objects/rule';
 
-import { RULES_MANAGEMENT_TABLE, RULE_NAME } from '../../../../screens/alerts_detection_rules';
+import {
+  RULES_MANAGEMENT_TABLE,
+  RULE_NAME,
+  INVESTIGATION_FIELDS_VALUE_ITEM,
+} from '../../../../screens/alerts_detection_rules';
 import {
   RULE_NAME_HEADER,
   RULE_TYPE_DETAILS,
@@ -29,6 +33,11 @@ import {
   fillEsqlQueryBar,
   fillAboutSpecificEsqlRuleAndContinue,
   createRuleWithoutEnabling,
+  expandAdvancedSettings,
+  fillCustomInvestigationFields,
+  fillRuleName,
+  fillDescription,
+  getAboutContinueButton,
 } from '../../../../tasks/create_new_rule';
 import { login } from '../../../../tasks/login';
 import { visit } from '../../../../tasks/navigation';
@@ -176,4 +185,38 @@ describe('Detection ES|QL rules, creation', { tags: ['@ess'] }, () => {
       cy.get(ESQL_QUERY_BAR).contains('Error validating ES|QL');
     });
   });
+
+  describe('ES|QL investigation fields', () => {
+    beforeEach(() => {
+      login();
+      visit(CREATE_RULE_URL);
+    });
+    it('shows custom ES|QL field in investigation fields autocomplete and saves it in rule', function () {
+      const CUSTOM_ESQL_FIELD = '_custom_agent_name';
+      const queryWithCustomFields = [
+        `from auditbeat* [metadata _id, _version, _index]`,
+        `eval ${CUSTOM_ESQL_FIELD} = agent.name`,
+        `keep _id, _custom_agent_name`,
+        `limit 5`,
+      ].join(' | ');
+
+      workaroundForResizeObserver();
+
+      selectEsqlRuleType();
+      expandEsqlQueryBar();
+      fillEsqlQueryBar(queryWithCustomFields);
+      getDefineContinueButton().click();
+
+      expandAdvancedSettings();
+      fillRuleName();
+      fillDescription();
+      fillCustomInvestigationFields([CUSTOM_ESQL_FIELD]);
+      getAboutContinueButton().click();
+
+      fillScheduleRuleAndContinue(rule);
+      createRuleWithoutEnabling();
+
+      cy.get(INVESTIGATION_FIELDS_VALUE_ITEM).should('have.text', CUSTOM_ESQL_FIELD);
+    });
+  });
 });

x-pack/test/security_solution_cypress/cypress/tasks/create_new_rule.ts

@@ -684,7 +684,7 @@ export const getIndicatorAtLeastOneInvalidationText = () => cy.contains(AT_LEAST
 export const getIndexPatternInvalidationText = () => cy.contains(AT_LEAST_ONE_INDEX_PATTERN);
 
 /** Returns the continue button on the step of about */
-const getAboutContinueButton = () => cy.get(ABOUT_CONTINUE_BTN);
+export const getAboutContinueButton = () => cy.get(ABOUT_CONTINUE_BTN);
 
 /** Returns the continue button on the step of define */
 export const getDefineContinueButton = () => cy.get(DEFINE_CONTINUE_BUTTON);