diff --git a/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx b/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx index 1d8fc029f003e..bda5dc608ebbd 100644 --- a/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx +++ b/src/platform/packages/private/kbn-esql-editor/src/esql_editor.tsx @@ -493,6 +493,7 @@ const ESQLEditorInternal = function ESQLEditor({ name: c.name, type: c.meta.esType as FieldType, hasConflict: c.meta.type === KBN_FIELD_TYPES.CONFLICT, + userDefined: false, }; }) || []; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/__tests__/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/__tests__/autocomplete.ts index bd7ca560affbe..76eb9a680c7a2 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/__tests__/autocomplete.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/__tests__/autocomplete.ts @@ -15,11 +15,10 @@ import { uniq } from 'lodash'; import type { LicenseType } from '@kbn/licensing-types'; import type { - ESQLUserDefinedColumn, - ESQLFieldWithMetadata, ICommandCallbacks, ISuggestionItem, Location, + ESQLColumnData, } from '../commands_registry/types'; import { getLocationFromCommandOrOptionName } from '../commands_registry/types'; import { aggFunctionDefinitions } from '../definitions/generated/aggregation_functions'; @@ -50,8 +49,7 @@ export const suggest = ( arg1: ESQLCommand, arg2: ICommandCallbacks, arg3: { - userDefinedColumns: Map; - fields: Map; + columns: Map; }, arg4?: number ) => Promise, @@ -80,8 +78,7 @@ export const expectSuggestions = async ( arg1: ESQLCommand, arg2: ICommandCallbacks, arg3: { - userDefinedColumns: Map; - fields: Map; + columns: Map; }, arg4?: number ) => Promise, @@ -105,12 +102,10 @@ export function getFieldNamesByType( _requestedType: Readonly>, excludeUserDefined: boolean = false ) { - const fieldsMap = mockContext.fields; - const userDefinedColumnsMap = mockContext.userDefinedColumns; - const fields = Array.from(fieldsMap.values()); - const userDefinedColumns = Array.from(userDefinedColumnsMap.values()).flat(); + const columnMap = mockContext.columns; + const columns = Array.from(columnMap.values()); const requestedType = Array.isArray(_requestedType) ? _requestedType : [_requestedType]; - const finalArray = excludeUserDefined ? fields : [...fields, ...userDefinedColumns]; + const finalArray = excludeUserDefined ? columns.filter((col) => !col.userDefined) : columns; return finalArray .filter( ({ type }) => diff --git a/src/platform/packages/shared/kbn-esql-ast/src/__tests__/context_fixtures.ts b/src/platform/packages/shared/kbn-esql-ast/src/__tests__/context_fixtures.ts index 8522f18e26891..4106ddce9a729 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/__tests__/context_fixtures.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/__tests__/context_fixtures.ts @@ -8,9 +8,8 @@ */ import type { IndexAutocompleteItem, InferenceEndpointAutocompleteItem } from '@kbn/esql-types'; import type { - ESQLFieldWithMetadata, + ESQLColumnData, ESQLPolicy, - ESQLUserDefinedColumn, ICommandCallbacks, ICommandContext, } from '../commands_registry/types'; @@ -115,84 +114,79 @@ export const editorExtensions = { }; export const mockContext: ICommandContext = { - userDefinedColumns: new Map([ + columns: new Map([ [ 'var0', - [ - { - name: 'var0', - type: 'double', - location: { min: 0, max: 10 }, - }, - ], + { + name: 'var0', + type: 'double', + location: { min: 0, max: 10 }, + userDefined: true, + }, ], [ 'col0', - [ - { - name: 'col0', - type: 'double', - location: { min: 0, max: 10 }, - }, - ], + { + name: 'col0', + type: 'double', + location: { min: 0, max: 10 }, + userDefined: true, + }, ], [ 'prompt', - [ - { - name: 'prompt', - type: 'keyword', - location: { min: 0, max: 10 }, - }, - ], + { + name: 'prompt', + type: 'keyword', + location: { min: 0, max: 10 }, + userDefined: true, + }, ], [ 'integerPrompt', - [ - { - name: 'integerPrompt', - type: 'integer', - location: { min: 0, max: 10 }, - }, - ], + { + name: 'integerPrompt', + type: 'integer', + location: { min: 0, max: 10 }, + userDefined: true, + }, ], [ 'ipPrompt', - [ - { - name: 'ipPrompt', - type: 'ip', - location: { min: 0, max: 10 }, - }, - ], + { + name: 'ipPrompt', + type: 'ip', + location: { min: 0, max: 10 }, + userDefined: true, + }, ], [ 'renamedField', - [ - { - name: 'renamedField', - type: 'keyword', - location: { min: 0, max: 10 }, - }, - ], + { + name: 'renamedField', + type: 'keyword', + location: { min: 0, max: 10 }, + userDefined: true, + }, ], - ]), - fields: new Map([ - ['keywordField', { name: 'keywordField', type: 'keyword' }], - ['any#Char$Field', { name: 'any#Char$Field', type: 'keyword' }], - ['textField', { name: 'textField', type: 'text' }], - ['doubleField', { name: 'doubleField', type: 'double' }], - ['integerField', { name: 'integerField', type: 'integer' }], - ['counterIntegerField', { name: 'counterIntegerField', type: 'counter_integer' }], - ['dateField', { name: 'dateField', type: 'date' }], - ['dateNanosField', { name: 'dateNanosField', type: 'date_nanos' }], - ['@timestamp', { name: '@timestamp', type: 'date' }], - ['ipField', { name: 'ipField', type: 'ip' }], - ['booleanField', { name: 'booleanField', type: 'boolean' }], - ['geoPointField', { name: 'geoPointField', type: 'geo_point' }], - ['geoShapeField', { name: 'geoShapeField', type: 'geo_shape' }], - ['versionField', { name: 'versionField', type: 'version' }], - ['longField', { name: 'longField', type: 'long' }], + ['keywordField', { name: 'keywordField', type: 'keyword', userDefined: false }], + ['any#Char$Field', { name: 'any#Char$Field', type: 'keyword', userDefined: false }], + ['textField', { name: 'textField', type: 'text', userDefined: false }], + ['doubleField', { name: 'doubleField', type: 'double', userDefined: false }], + ['integerField', { name: 'integerField', type: 'integer', userDefined: false }], + [ + 'counterIntegerField', + { name: 'counterIntegerField', type: 'counter_integer', userDefined: false }, + ], + ['dateField', { name: 'dateField', type: 'date', userDefined: false }], + ['dateNanosField', { name: 'dateNanosField', type: 'date_nanos', userDefined: false }], + ['@timestamp', { name: '@timestamp', type: 'date', userDefined: false }], + ['ipField', { name: 'ipField', type: 'ip', userDefined: false }], + ['booleanField', { name: 'booleanField', type: 'boolean', userDefined: false }], + ['geoPointField', { name: 'geoPointField', type: 'geo_point', userDefined: false }], + ['geoShapeField', { name: 'geoShapeField', type: 'geo_shape', userDefined: false }], + ['versionField', { name: 'versionField', type: 'version', userDefined: false }], + ['longField', { name: 'longField', type: 'long', userDefined: false }], ]), policies: new Map(policies.map((policy) => [policy.name, policy])), sources: indexes.map((name) => ({ diff --git a/src/platform/packages/shared/kbn-esql-ast/src/__tests__/validation.ts b/src/platform/packages/shared/kbn-esql-ast/src/__tests__/validation.ts index 859db40e90303..3652598d184d0 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/__tests__/validation.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/__tests__/validation.ts @@ -7,7 +7,7 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ -import type { ESQLUserDefinedColumn, ESQLFieldWithMetadata } from '../commands_registry/types'; +import type { ESQLColumnData } from '../commands_registry/types'; import { Parser } from '../parser'; import type { ESQLCommand, ESQLMessage } from '../types'; import { mockContext } from './context_fixtures'; @@ -28,8 +28,7 @@ export const expectErrors = ( arg0: ESQLCommand, arg1: ESQLCommand[], arg2: { - userDefinedColumns: Map; - fields: Map; + columns: Map; } ) => any ) => { diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts index beb59f176b7ea..56a086e22c668 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.test.ts @@ -7,54 +7,48 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import { synth } from '../../../..'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; +import type { ESQLColumnData } from '../../types'; import { columnsAfter } from './columns_after'; describe('CHANGE_POINT > columnsAfter', () => { - const context = { - userDefinedColumns: new Map([]), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }; it('adds "type" and "pvalue" fields, when AS option not specified', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'count', type: 'double' }, - ] as ESQLFieldWithMetadata[]; + const previousCommandFields: ESQLColumnData[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'count', type: 'double', userDefined: false }, + ]; - const result = columnsAfter( - synth.cmd`CHANGE_POINT count ON field1`, - previousCommandFields, - context - ); + const result = columnsAfter(synth.cmd`CHANGE_POINT count ON field1`, previousCommandFields, ''); - expect(result).toEqual([ - { name: 'field1', type: 'keyword' }, - { name: 'count', type: 'double' }, - { name: 'type', type: 'keyword' }, - { name: 'pvalue', type: 'double' }, + expect(result).toEqual([ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'count', type: 'double', userDefined: false }, + { name: 'type', type: 'keyword', userDefined: false }, + { name: 'pvalue', type: 'double', userDefined: false }, ]); }); it('adds the given names as fields, when AS option is specified', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'count', type: 'double' }, - ] as ESQLFieldWithMetadata[]; + const previousCommandFields: ESQLColumnData[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'count', type: 'double', userDefined: false }, + ]; const result = columnsAfter( synth.cmd`CHANGE_POINT count ON field1 AS changePointType, pValue`, previousCommandFields, - context + '' ); - expect(result).toEqual([ - { name: 'field1', type: 'keyword' }, - { name: 'count', type: 'double' }, - { name: 'changePointType', type: 'keyword' }, - { name: 'pValue', type: 'double' }, + expect(result).toEqual([ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'count', type: 'double', userDefined: false }, + { + name: 'changePointType', + type: 'keyword', + userDefined: true, + location: { min: 0, max: 0 }, + }, + { name: 'pValue', type: 'double', userDefined: true, location: { min: 0, max: 0 } }, ]); }); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts index 00680ddee9de7..45f7cbedbfd0f 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/columns_after.ts @@ -7,30 +7,49 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import uniqBy from 'lodash/uniqBy'; -import { type ESQLCommand, type ESQLAstChangePointCommand } from '../../../types'; import { LeafPrinter } from '../../../pretty_print/leaf_printer'; -import type { ESQLFieldWithMetadata } from '../../types'; -import type { ICommandContext } from '../../types'; +import { type ESQLAstChangePointCommand, type ESQLCommand } from '../../../types'; +import type { ESQLColumnData } from '../../types'; export const columnsAfter = ( command: ESQLCommand, - previousColumns: ESQLFieldWithMetadata[], - context?: ICommandContext + previousColumns: ESQLColumnData[], + query: string ) => { const { target } = command as ESQLAstChangePointCommand; - return uniqBy( - [ - ...previousColumns, - { - name: target ? LeafPrinter.column(target.type) : 'type', - type: 'keyword' as const, - }, - { - name: target ? LeafPrinter.column(target.pvalue) : 'pvalue', - type: 'double' as const, - }, - ], - 'name' - ); + let typeField: ESQLColumnData; + let pvalueField: ESQLColumnData; + + if (target?.type) { + typeField = { + name: LeafPrinter.column(target.type), + type: 'keyword' as const, + userDefined: true, + location: target.type.location, + }; + } else { + typeField = { + name: 'type', + type: 'keyword' as const, + userDefined: false, + }; + } + + if (target?.pvalue) { + pvalueField = { + name: LeafPrinter.column(target.pvalue), + type: 'double' as const, + userDefined: true, + location: target.pvalue.location, + }; + } else { + pvalueField = { + name: 'pvalue', + type: 'double' as const, + userDefined: false, + }; + } + + return uniqBy([...previousColumns, typeField, pvalueField], 'name'); }; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.test.ts index af5d8857495a0..fecfc2760b627 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.test.ts @@ -50,11 +50,11 @@ describe('CHANGE_POINT Validation', () => { test('raises error when the default @timestamp field is missing', () => { // make sure that @timestamp field is not present - const newFields = new Map(mockContext.fields); - newFields.delete('@timestamp'); + const newColumns = new Map(mockContext.columns); + newColumns.delete('@timestamp'); const context = { ...mockContext, - fields: newFields, + columns: newColumns, }; changePointExpectErrors( 'FROM a_index | CHANGE_POINT doubleField', diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts index 94e56d0a09d16..3d378c7ff780d 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/change_point/validate.ts @@ -27,14 +27,10 @@ export const validate = ( const valueArg = command.args[0]; if (isColumn(valueArg)) { const columnName = valueArg.name; - // look up for columns in userDefinedColumns and existing fields let valueColumnType: SupportedDataType | 'unknown' | undefined; - const userDefinedColumnRef = context?.userDefinedColumns.get(columnName); - if (userDefinedColumnRef) { - valueColumnType = userDefinedColumnRef.find((v) => v.name === columnName)?.type; - } else { - const fieldRef = context?.fields.get(columnName); - valueColumnType = fieldRef?.type; + + if (context?.columns.has(columnName)) { + valueColumnType = context?.columns.get(columnName)?.type; } if (valueColumnType && !isNumericType(valueColumnType)) { @@ -45,7 +41,7 @@ export const validate = ( // validate ON column const defaultOnColumnName = '@timestamp'; const onColumn = command.args.find((arg) => isOptionNode(arg) && arg.name === 'on'); - const hasDefaultOnColumn = context?.fields.has(defaultOnColumnName); + const hasDefaultOnColumn = context?.columns.has(defaultOnColumnName); if (!onColumn && !hasDefaultOnColumn) { messages.push({ location: command.location, @@ -58,20 +54,15 @@ export const validate = ( }); } - // validate AS - const asArg = command.args.find((arg) => isOptionNode(arg) && arg.name === 'as'); - if (asArg && isOptionNode(asArg)) { - // populate userDefinedColumns references to prevent the common check from failing with unknown column - asArg.args.forEach((arg, index) => { - if (isColumn(arg)) { - context?.userDefinedColumns.set(arg.name, [ - { name: arg.name, location: arg.location, type: index === 0 ? 'keyword' : 'long' }, - ]); - } - }); - } - - messages.push(...validateCommandArguments(command, ast, context, callbacks)); + messages.push( + ...validateCommandArguments( + // exclude AS option from generic validation + { ...command, args: command.args.filter((arg) => !(isOptionNode(arg) && arg.name === 'as')) }, + ast, + context, + callbacks + ) + ); return messages; }; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts index cf96358dd49c3..a10f5f8100ad3 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/autocomplete.ts @@ -62,11 +62,7 @@ function getPosition( } const expressionRoot = prompt?.text !== EDITOR_MARKER ? prompt : undefined; - const expressionType = getExpressionType( - expressionRoot, - context?.fields, - context?.userDefinedColumns - ); + const expressionType = getExpressionType(expressionRoot, context?.columns); if (isExpressionComplete(expressionType, query)) { return CompletionPosition.AFTER_PROMPT; @@ -161,8 +157,7 @@ export async function autocomplete( callbacks?.getByType, { functions: true, - fields: true, - userDefinedColumns: context?.userDefinedColumns, + columns: true, }, {}, callbacks?.hasMinimumLicenseRequired, diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts index 4020af765f063..c08c11e2d2498 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.test.ts @@ -7,52 +7,45 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import { synth } from '../../../..'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; +import type { ESQLColumnData } from '../../types'; import { columnsAfter } from './columns_after'; describe('COMPLETION', () => { - const context = { - userDefinedColumns: new Map([]), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }; it('adds "completion" field, when targetField is not specified', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'count', type: 'double' }, - ] as ESQLFieldWithMetadata[]; + const previousCommandColumns: ESQLColumnData[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'count', type: 'double', userDefined: false }, + ]; const result = columnsAfter( synth.cmd`COMPLETION "prompt" WITH {"inference_id": "my-inference-id"}`, - previousCommandFields, - context + previousCommandColumns, + '' ); expect(result).toEqual([ - { name: 'field1', type: 'keyword' }, - { name: 'count', type: 'double' }, - { name: 'completion', type: 'keyword' }, + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'count', type: 'double', userDefined: false }, + { name: 'completion', type: 'keyword', userDefined: false }, ]); }); it('adds the given targetField as field, when targetField is specified', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'count', type: 'double' }, - ] as ESQLFieldWithMetadata[]; + const previousCommandColumns: ESQLColumnData[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'count', type: 'double', userDefined: false }, + ]; const result = columnsAfter( synth.cmd`COMPLETION customField = "prompt" WITH {"inference_id": "my-inference-id"}`, - previousCommandFields, - context + previousCommandColumns, + '' ); expect(result).toEqual([ - { name: 'field1', type: 'keyword' }, - { name: 'count', type: 'double' }, - { name: 'customField', type: 'keyword' }, + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'count', type: 'double', userDefined: false }, + { name: 'customField', type: 'keyword', userDefined: true, location: { min: 0, max: 0 } }, ]); }); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts index 82e3c43d951be..e128a20492fca 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/columns_after.ts @@ -7,25 +7,32 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import uniqBy from 'lodash/uniqBy'; -import { type ESQLCommand, type ESQLAstCompletionCommand } from '../../../types'; import { LeafPrinter } from '../../../pretty_print/leaf_printer'; -import type { ESQLFieldWithMetadata } from '../../types'; -import type { ICommandContext } from '../../types'; +import { type ESQLAstCompletionCommand, type ESQLCommand } from '../../../types'; +import type { ESQLColumnData, ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; export const columnsAfter = ( command: ESQLCommand, - previousColumns: ESQLFieldWithMetadata[], - context?: ICommandContext + previousColumns: ESQLColumnData[], + query: string ) => { const { targetField } = command as ESQLAstCompletionCommand; return uniqBy( [ ...previousColumns, - { - name: targetField ? LeafPrinter.column(targetField) : 'completion', - type: 'keyword' as const, - }, + targetField + ? ({ + name: LeafPrinter.column(targetField), + type: 'keyword' as const, + userDefined: true, + location: targetField.location, + } as ESQLUserDefinedColumn) + : ({ + name: 'completion', + type: 'keyword' as const, + userDefined: false, + } as ESQLFieldWithMetadata), ], 'name' ); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts index ffd90ad3370a5..2daf08a996565 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/completion/validate.ts @@ -22,13 +22,9 @@ export const validate = ( ): ESQLMessage[] => { const messages: ESQLMessage[] = []; - const { prompt, location, targetField, inferenceId } = command as ESQLAstCompletionCommand; + const { prompt, location, inferenceId } = command as ESQLAstCompletionCommand; - const promptExpressionType = getExpressionType( - prompt, - context?.fields, - context?.userDefinedColumns - ); + const promptExpressionType = getExpressionType(prompt, context?.columns); if (!supportedPromptTypes.includes(promptExpressionType)) { messages.push( @@ -43,17 +39,6 @@ export const validate = ( messages.push(errors.byId('inferenceIdRequired', command.location, { command: 'COMPLETION' })); } - const targetName = targetField?.name || 'completion'; - - // Sets the target field so the column is recognized after the command is applied - context?.userDefinedColumns.set(targetName, [ - { - name: targetName, - location: targetField?.location || command.location, - type: 'keyword', - }, - ]); - messages.push(...validateCommandArguments(command, ast, context, callbacks)); return messages; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts index c154efe00f4bf..18d155383b8ba 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.test.ts @@ -7,7 +7,7 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import { synth } from '../../../..'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; +import type { ESQLColumnData } from '../../types'; import { columnsAfter, extractDissectColumnNames } from './columns_after'; describe('DISSECT', () => { @@ -68,29 +68,18 @@ describe('DISSECT', () => { }); }); describe('columnsAfter', () => { - const context = { - userDefinedColumns: new Map([]), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }; it('adds the DISSECT pattern columns as fields', () => { - const previousColumns = [ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - ] as ESQLFieldWithMetadata[]; + const previousColumns: ESQLColumnData[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; - const result = columnsAfter( - synth.cmd`DISSECT agent "%{firstWord}"`, - previousColumns, - context - ); + const result = columnsAfter(synth.cmd`DISSECT agent "%{firstWord}"`, previousColumns, ''); - expect(result).toEqual([ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - { name: 'firstWord', type: 'keyword' }, + expect(result).toEqual([ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + { name: 'firstWord', type: 'keyword', userDefined: false }, ]); }); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts index c8a929322fcf3..c8d7cbdf3c38c 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/columns_after.ts @@ -8,8 +8,7 @@ */ import { walk } from '../../../walker'; import { type ESQLCommand } from '../../../types'; -import type { ESQLFieldWithMetadata } from '../../types'; -import type { ICommandContext } from '../../types'; +import type { ESQLColumnData } from '../../types'; function unquoteTemplate(inputString: string): string { if (inputString.startsWith('"') && inputString.endsWith('"') && inputString.length >= 2) { @@ -35,8 +34,8 @@ export function extractDissectColumnNames(pattern: string): string[] { export const columnsAfter = ( command: ESQLCommand, - previousColumns: ESQLFieldWithMetadata[], - context?: ICommandContext + previousColumns: ESQLColumnData[], + query: string ) => { const columns: string[] = []; @@ -49,6 +48,13 @@ export const columnsAfter = ( return [ ...previousColumns, - ...columns.map((column) => ({ name: column, type: 'keyword' as const })), + ...columns.map((column) => { + const newColumn: ESQLColumnData = { + name: column, + type: 'keyword' as const, + userDefined: false, + }; + return newColumn; + }), ]; }; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/validate.ts index d38a39beb44cf..ae689cbad4195 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/validate.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/dissect/validate.ts @@ -22,9 +22,9 @@ import { validateCommandArguments } from '../../../definitions/utils/validation' const validateColumnForGrokDissect = (command: ESQLCommand, context?: ICommandContext) => { const acceptedColumnTypes: FieldType[] = ['keyword', 'text']; const astCol = command.args[0] as ESQLColumn; - const columnRef = context?.fields.get(astCol.name); + const columnRef = context?.columns.get(astCol.name); - if (columnRef && !acceptedColumnTypes.includes(columnRef.type)) { + if (columnRef && !acceptedColumnTypes.includes(columnRef.type as FieldType)) { return [ getMessageFromId({ messageId: 'unsupportedColumnTypeForCommand', diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts index d0428da653a62..32bcb28fb34c1 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.test.ts @@ -7,25 +7,18 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import { synth } from '../../../..'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; +import type { ESQLColumnData } from '../../types'; import { columnsAfter } from './columns_after'; describe('DROP', () => { - const context = { - userDefinedColumns: new Map([]), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }; it('removes the columns defined in the command', () => { - const previousColumns = [ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - ] as ESQLFieldWithMetadata[]; + const previousColumns: ESQLColumnData[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; - const result = columnsAfter(synth.cmd`DROP field1`, previousColumns, context); + const result = columnsAfter(synth.cmd`DROP field1`, previousColumns, ''); - expect(result).toEqual([{ name: 'field2', type: 'double' }]); + expect(result).toEqual([{ name: 'field2', type: 'double', userDefined: false }]); }); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts index 8033d2e754b97..491095041f95c 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/columns_after.ts @@ -8,13 +8,12 @@ */ import { walk } from '../../../walker'; import { type ESQLCommand } from '../../../types'; -import type { ESQLFieldWithMetadata } from '../../types'; -import type { ICommandContext } from '../../types'; +import type { ESQLColumnData } from '../../types'; export const columnsAfter = ( command: ESQLCommand, - previousColumns: ESQLFieldWithMetadata[], - context?: ICommandContext + previousColumns: ESQLColumnData[], + query: string ) => { const columnsToDrop: string[] = []; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/validate.test.ts index 34e3d6fa91cc1..bab0ef07c51bd 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/validate.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/drop/validate.test.ts @@ -21,24 +21,22 @@ describe('DROP Validation', () => { test('validates the most basic query', () => { // make sure that @timestamp field is not present - const newUserDefinedColumns = new Map(mockContext.userDefinedColumns); - newUserDefinedColumns.set('MIN(doubleField * 10)', [ - { - name: 'MIN(doubleField * 10)', - type: 'double', - location: { min: 0, max: 10 }, - }, - ]); - newUserDefinedColumns.set('COUNT(*)', [ - { - name: 'COUNT(*)', - type: 'integer', - location: { min: 0, max: 10 }, - }, - ]); + const newColumns = new Map(mockContext.columns); + newColumns.set('MIN(doubleField * 10)', { + name: 'MIN(doubleField * 10)', + type: 'double', + location: { min: 0, max: 10 }, + userDefined: true, + }); + newColumns.set('COUNT(*)', { + name: 'COUNT(*)', + type: 'integer', + location: { min: 0, max: 10 }, + userDefined: true, + }); const context = { ...mockContext, - userDefinedColumns: newUserDefinedColumns, + columns: newColumns, }; dropExpectErrors('from index | drop textField, doubleField, dateField', []); dropExpectErrors('from index | drop `any#Char$Field`', []); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/autocomplete.ts index c9d0326cee9ff..1f1c937c3c774 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/autocomplete.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/autocomplete.ts @@ -6,30 +6,30 @@ * 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 { findFinalWord, findPreviousWord } from '../../../definitions/utils/autocomplete/helpers'; +import { buildFieldsDefinitions } from '../../../definitions/utils/functions'; +import { getOperatorSuggestions } from '../../../definitions/utils/operators'; +import { unescapeColumnName } from '../../../definitions/utils/shared'; import type { ESQLCommand, ESQLSource } from '../../../types'; import { - pipeCompleteItem, commaCompleteItem, getNewUserDefinedColumnSuggestion, + pipeCompleteItem, } from '../../complete_items'; -import { findFinalWord, findPreviousWord } from '../../../definitions/utils/autocomplete/helpers'; -import { unescapeColumnName } from '../../../definitions/utils/shared'; -import type { ESQLPolicy, ICommandCallbacks } from '../../types'; -import { type ISuggestionItem, type ICommandContext, Location } from '../../types'; +import { TRIGGER_SUGGESTION_COMMAND } from '../../constants'; +import type { ESQLColumnData, ESQLPolicy, ICommandCallbacks } from '../../types'; +import { Location, type ICommandContext, type ISuggestionItem } from '../../types'; import { Position, buildMatchingFieldsDefinition, + buildPoliciesDefinitions, + getPolicyMetadata, getPosition, modeSuggestions, noPoliciesAvailableSuggestion, onSuggestion, withSuggestion, - getPolicyMetadata, - buildPoliciesDefinitions, } from './util'; -import { TRIGGER_SUGGESTION_COMMAND } from '../../constants'; -import { getOperatorSuggestions } from '../../../definitions/utils/operators'; -import { buildFieldsDefinitions } from '../../../definitions/utils/functions'; export async function autocomplete( query: string, @@ -41,8 +41,14 @@ export async function autocomplete( const innerText = query.substring(0, cursorPosition); const pos = getPosition(innerText, command); const policies = context?.policies ?? new Map(); - const fieldsMap = context?.fields ?? new Map(); - const allColumnNames = Array.from(fieldsMap.keys()); + const columnMap = context?.columns ?? new Map(); + + const fieldNames: string[] = []; + for (const name of columnMap.keys()) { + const col = columnMap.get(name); + if (col && !col.userDefined) fieldNames.push(name); + } + const policyName = ( command.args.find((arg) => !Array.isArray(arg) && arg.type === 'source') as | ESQLSource @@ -107,7 +113,7 @@ export async function autocomplete( return []; } - return buildMatchingFieldsDefinition(policyMetadata.matchField, allColumnNames); + return buildMatchingFieldsDefinition(policyMetadata.matchField, fieldNames); } case Position.AFTER_ON_CLAUSE: diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts new file mode 100644 index 0000000000000..0838236692969 --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.test.ts @@ -0,0 +1,114 @@ +/* + * 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 { synth } from '../../../..'; +import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types'; +import { columnsAfter } from './columns_after'; + +describe('ENRICH columnsAfter', () => { + it('returns previousColumns when no enrich columns', async () => { + const previousColumns: ESQLColumnData[] = [ + { name: 'fieldA', type: 'keyword', userDefined: false }, + { name: 'fieldB', type: 'long', userDefined: false }, + ]; + const result = await columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', { + fromJoin: () => Promise.resolve([]), + fromEnrich: () => Promise.resolve([]), + fromFrom: () => Promise.resolve([]), + }); + expect(result).toEqual(previousColumns); + }); + + it('adds all enrich columns to when no WITH clause', async () => { + const previousColumns: ESQLColumnData[] = [ + { name: 'fieldA', type: 'keyword', userDefined: false }, + { name: 'fieldB', type: 'long', userDefined: false }, + ]; + const enrichColumns: ESQLFieldWithMetadata[] = [ + { name: 'enrichField1', type: 'keyword', userDefined: false }, + { name: 'enrichField2', type: 'double', userDefined: false }, + ]; + const result = await columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', { + fromJoin: () => Promise.resolve([]), + fromEnrich: () => Promise.resolve(enrichColumns), + fromFrom: () => Promise.resolve([]), + }); + expect(result).toEqual([...enrichColumns, ...previousColumns]); + }); + + it('adds only declared columns when WITH clause is present', async () => { + const previousColumns: ESQLColumnData[] = [ + { name: 'fieldA', type: 'keyword', userDefined: false }, + { name: 'fieldB', type: 'long', userDefined: false }, + ]; + const enrichColumns: ESQLFieldWithMetadata[] = [ + { name: 'enrichField1', type: 'keyword', userDefined: false }, + { name: 'enrichField2', type: 'double', userDefined: false }, + ]; + const result = await columnsAfter( + synth.cmd`ENRICH policy ON matchfield WITH enrichField2`, + previousColumns, + '', + { + fromJoin: () => Promise.resolve([]), + fromEnrich: () => Promise.resolve(enrichColumns), + fromFrom: () => Promise.resolve([]), + } + ); + expect(result).toEqual([enrichColumns[1], ...previousColumns]); + }); + + it('renames enrichment fields using WITH', async () => { + const previousColumns: ESQLColumnData[] = [ + { name: 'fieldA', type: 'keyword', userDefined: false }, + { name: 'fieldB', type: 'long', userDefined: false }, + ]; + const enrichColumns: ESQLFieldWithMetadata[] = [ + { name: 'enrichField1', type: 'keyword', userDefined: false }, + { name: 'enrichField2', type: 'double', userDefined: false }, + ]; + const result = await columnsAfter( + synth.cmd`ENRICH policy ON matchfield WITH foo = enrichField1, bar = enrichField2`, + previousColumns, + '', + { + fromEnrich: () => Promise.resolve(enrichColumns), + fromJoin: () => Promise.resolve([]), + fromFrom: () => Promise.resolve([]), + } + ); + const expected = [ + { name: 'foo', type: 'keyword', userDefined: false }, + { name: 'bar', type: 'double', userDefined: false }, + ...previousColumns, + ]; + expect(result).toEqual(expected); + }); + + it('overwrites previous columns with the same name', async () => { + const previousColumns: ESQLColumnData[] = [ + { name: 'fieldA', type: 'keyword', userDefined: false }, + { name: 'fieldB', type: 'long', userDefined: false }, + ]; + const enrichFields: ESQLFieldWithMetadata[] = [ + { name: 'fieldA', type: 'text', userDefined: false }, + { name: 'fieldC', type: 'double', userDefined: false }, + ]; + const result = await columnsAfter(synth.cmd`ENRICH policy ON matchfield`, previousColumns, '', { + fromEnrich: () => Promise.resolve(enrichFields), + fromJoin: () => Promise.resolve([]), + fromFrom: () => Promise.resolve([]), + }); + expect(result).toEqual([ + { name: 'fieldA', type: 'text', userDefined: false }, + { name: 'fieldC', type: 'double', userDefined: false }, + { name: 'fieldB', type: 'long', userDefined: false }, + ]); + }); +}); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts new file mode 100644 index 0000000000000..9a8898a16369b --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/columns_after.ts @@ -0,0 +1,56 @@ +/* + * 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 { uniqBy } from 'lodash'; +import { isAssignment, isColumn, isOptionNode } from '../../../ast/is'; +import type { ESQLCommandOption } from '../../../types'; +import { type ESQLCommand } from '../../../types'; +import type { ESQLColumnData } from '../../types'; +import type { IAdditionalFields } from '../../registry'; + +export const columnsAfter = async ( + command: ESQLCommand, + previousColumns: ESQLColumnData[], + query: string, + additionalFields: IAdditionalFields +) => { + const enrichFields = await additionalFields.fromEnrich(command); + let fieldsToAdd = enrichFields; + + // the with option scopes down the fields that are added + // and potentially renames things + const withOption = command.args.find((arg) => isOptionNode(arg) && arg.name === 'with') as + | ESQLCommandOption + | undefined; + if (withOption) { + const declaredFieldEntries = withOption.args + .map((arg) => { + if ( + isAssignment(arg) && + isColumn(arg.args[0]) && + Array.isArray(arg.args[1]) && + isColumn(arg.args[1][0]) + ) { + return [arg.args[1][0].parts.join('.'), arg.args[0].parts.join('.')]; + } + return undefined; + }) + .filter(Boolean) as [string, string][]; + + const declaredFields = new Map(declaredFieldEntries); + + fieldsToAdd = enrichFields + .filter((field) => declaredFields.has(field.name)) + .map((field) => { + const newName = declaredFields.get(field.name)!; + return { ...field, name: newName }; + }); + } + + return uniqBy([...fieldsToAdd, ...previousColumns], 'name'); +}; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/index.ts index 79303aa63ee77..230d81d28f831 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/index.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/index.ts @@ -10,11 +10,13 @@ import { i18n } from '@kbn/i18n'; import type { ICommandMethods } from '../../registry'; import { autocomplete } from './autocomplete'; import { validate } from './validate'; +import { columnsAfter } from './columns_after'; import type { ICommandContext } from '../../types'; const enrichCommandMethods: ICommandMethods = { validate, autocomplete, + columnsAfter, }; export const enrichCommand = { diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.test.ts index c402abea064f7..cdf34554d19f6 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.test.ts @@ -6,6 +6,7 @@ * 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 { mockContext } from '../../../__tests__/context_fixtures'; import { validate } from './validate'; @@ -22,79 +23,67 @@ describe('ENRICH Validation', () => { }); test('validates the most basic query', () => { - enrichExpectErrors(`from a_index | enrich policy `, []); - enrichExpectErrors(`from a_index | enrich policy on textField with var0 = doubleField `, []); - enrichExpectErrors( - `from a_index | enrich policy on textField with var0 = doubleField, textField `, - [] - ); - enrichExpectErrors( - `from a_index | enrich policy on textField with col0 = doubleField, var0 = textField`, - [] - ); - enrichExpectErrors(`from a_index | enrich policy with doubleField`, []); - enrichExpectErrors(`from a_index | enrich policy | eval doubleField`, []); - enrichExpectErrors(`from a_index | enrich policy with col0 = doubleField | eval col0`, []); + enrichExpectErrors(`FROM a_index | ENRICH policy `, []); }); + test('validates the coordinators', () => { for (const value of ['any', 'coordinator', 'remote']) { - enrichExpectErrors(`from a_index | enrich _${value}:policy `, []); - enrichExpectErrors(`from a_index | enrich _${value} : policy `, [ + enrichExpectErrors(`FROM a_index | enrich _${value}:policy `, []); + enrichExpectErrors(`FROM a_index | enrich _${value} : policy `, [ `Unknown policy "_${value}"`, ]); - enrichExpectErrors(`from a_index | enrich _${value}: policy `, [ + enrichExpectErrors(`FROM a_index | enrich _${value}: policy `, [ `Unknown policy "_${value}"`, ]); - enrichExpectErrors(`from a_index | enrich _${camelCase(value)}:policy `, []); - enrichExpectErrors(`from a_index | enrich _${value.toUpperCase()}:policy `, []); + enrichExpectErrors(`FROM a_index | enrich _${camelCase(value)}:policy `, []); + enrichExpectErrors(`FROM a_index | enrich _${value.toUpperCase()}:policy `, []); } - enrichExpectErrors(`from a_index | enrich _unknown:policy`, [ + enrichExpectErrors(`FROM a_index | enrich _unknown:policy`, [ 'Unrecognized value "_unknown" for ENRICH, mode needs to be one of [_any, _coordinator, _remote]', ]); - enrichExpectErrors(`from a_index | enrich any:policy`, [ + enrichExpectErrors(`FROM a_index | enrich any:policy`, [ 'Unrecognized value "any" for ENRICH, mode needs to be one of [_any, _coordinator, _remote]', ]); }); test('raises error on unknown policy', () => { - enrichExpectErrors(`from a_index | enrich _`, ['Unknown policy "_"']); - enrichExpectErrors(`from a_index | enrich _:`, ['Unknown policy "_"']); - enrichExpectErrors(`from a_index | enrich any:`, ['Unknown policy "any"']); - enrichExpectErrors(`from a_index | enrich _:policy`, [ + enrichExpectErrors(`FROM a_index | enrich _`, ['Unknown policy "_"']); + enrichExpectErrors(`FROM a_index | enrich _:`, ['Unknown policy "_"']); + enrichExpectErrors(`FROM a_index | enrich any:`, ['Unknown policy "any"']); + enrichExpectErrors(`FROM a_index | enrich _:policy`, [ 'Unrecognized value "_" for ENRICH, mode needs to be one of [_any, _coordinator, _remote]', ]); - enrichExpectErrors(`from a_index | enrich _any:`, ['Unknown policy "_any"']); - enrichExpectErrors('from a_index | enrich `this``is fine`', ['Unknown policy "`this``is"']); - enrichExpectErrors('from a_index | enrich this is fine', ['Unknown policy "this"']); - enrichExpectErrors(`from a_index |enrich missing-policy `, ['Unknown policy "missing-policy"']); + enrichExpectErrors(`FROM a_index | enrich _any:`, ['Unknown policy "_any"']); + enrichExpectErrors('FROM a_index | enrich `this``is fine`', ['Unknown policy "`this``is"']); + enrichExpectErrors('FROM a_index | enrich this is fine', ['Unknown policy "this"']); + enrichExpectErrors(`FROM a_index |enrich missing-policy `, ['Unknown policy "missing-policy"']); }); - test('validates the columns', () => { - enrichExpectErrors(`from a_index | enrich policy on b `, ['Unknown column "b"']); + test('validates match field', () => { + enrichExpectErrors(`FROM a_index | ENRICH policy ON b `, ['Unknown column "b"']); - enrichExpectErrors('from a_index | enrich policy on `this``is fine`', [ + enrichExpectErrors('FROM a_index | ENRICH policy ON `this``is fine`', [ 'Unknown column "this`is fine"', ]); - enrichExpectErrors('from a_index | enrich policy on this is fine', ['Unknown column "this"']); - enrichExpectErrors(`from a_index | enrich policy on textField with col1 `, [ - 'Unknown column "col1"', - ]); - enrichExpectErrors(`from a_index |enrich policy on doubleField with col1 = `, [ - 'Unknown column "col1"', - ]); - enrichExpectErrors(`from a_index | enrich policy on textField with col1 = c `, [ - 'Unknown column "col1"', - 'Unknown column "c"', - ]); - enrichExpectErrors(`from a_index |enrich policy on doubleField with col1 = , `, [ - 'Unknown column "col1"', - ]); - enrichExpectErrors(`from a_index | enrich policy on textField with var0 = doubleField, col1 `, [ - 'Unknown column "col1"', - ]); + enrichExpectErrors('FROM a_index | ENRICH policy ON this is fine', ['Unknown column "this"']); + }); + + test('validates enrich fields against policy', () => { + // otherField and yetAnotherField are enrich fields in the policy + enrichExpectErrors(`FROM a_index | ENRICH policy ON textField WITH otherField `, []); enrichExpectErrors( - `from a_index |enrich policy on doubleField with col0 = doubleField, col1 = `, - ['Unknown column "col1"'] + `FROM a_index | ENRICH policy ON textField WITH otherField, newname = yetAnotherField `, + [] ); + + // keywordField is a field in the index, but not an enrich field + enrichExpectErrors(`FROM a_index | ENRICH policy ON textField WITH keywordField `, [ + 'Unknown column "keywordField"', + ]); + + // col1 shouldn't be validated, only foo + enrichExpectErrors(`FROM a_index |ENRICH policy ON doubleField WITH col1 = foo`, [ + 'Unknown column "foo"', + ]); }); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.ts index 2fc375fe8bdb1..a0fe9df7c9d2c 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/enrich/validate.ts @@ -6,11 +6,18 @@ * 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 { getMessageFromId } from '../../../definitions/utils/errors'; -import type { ESQLSource, ESQLCommand, ESQLMessage, ESQLAst } from '../../../types'; -import { ENRICH_MODES } from './util'; -import type { ESQLPolicy, ICommandContext, ICommandCallbacks } from '../../types'; +import { isAssignment, isColumn, isOptionNode } from '../../../ast/is'; +import { errors, getMessageFromId } from '../../../definitions/utils/errors'; import { validateCommandArguments } from '../../../definitions/utils/validation'; +import type { + ESQLAst, + ESQLCommand, + ESQLCommandOption, + ESQLMessage, + ESQLSource, +} from '../../../types'; +import type { ESQLPolicy, ICommandCallbacks, ICommandContext } from '../../types'; +import { ENRICH_MODES } from './util'; export const validate = ( command: ESQLCommand, @@ -53,7 +60,34 @@ export const validate = ( } } - messages.push(...validateCommandArguments(command, ast, context, callbacks)); + const policy = index && policies.get(index.valueUnquoted); + const withOption = command.args.find( + (arg) => isOptionNode(arg) && arg.name === 'with' + ) as ESQLCommandOption; + + if (withOption && policy) { + withOption.args.forEach((arg) => { + if (isAssignment(arg) && Array.isArray(arg.args[1]) && isColumn(arg.args[1][0])) { + const column = arg.args[1][0]; + if (!policy.enrichFields.includes(column.parts.join('.'))) { + messages.push(errors.unknownColumn(column)); + } + } + }); + } + + messages.push( + ...validateCommandArguments( + { + ...command, + // exclude WITH from generic validation since it shouldn't be compared against the generic column list + args: command.args.filter((arg) => arg !== withOption), + }, + ast, + context, + callbacks + ) + ); return messages; }; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/autocomplete.ts index 4bd50f4cac152..ceae006efafa4 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/autocomplete.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/autocomplete.ts @@ -78,10 +78,7 @@ export async function autocomplete( if ( // don't suggest finishing characters if incomplete expression - isExpressionComplete( - getExpressionType(expressionRoot, context?.fields, context?.userDefinedColumns), - innerText - ) && + isExpressionComplete(getExpressionType(expressionRoot, context?.columns), innerText) && // don't suggest finishing characters if the expression is a column // because "EVAL columnName" is a useless expression expressionRoot && diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts new file mode 100644 index 0000000000000..0126c48f87211 --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.test.ts @@ -0,0 +1,136 @@ +/* + * 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 { Parser, synth } from '../../../..'; +import type { ESQLColumnData } from '../../types'; +import { columnsAfter } from './columns_after'; + +describe('EVAL > columnsAfter', () => { + const baseColumns: ESQLColumnData[] = [ + { name: 'foo', type: 'integer', userDefined: false }, + { name: 'bar', type: 'keyword', userDefined: false }, + ]; + + it('adds a new column for a simple assignment', () => { + const command = synth.cmd`EVAL baz = foo + 1`; + const result = columnsAfter(command, baseColumns, ''); + + expect(result).toEqual([ + { + name: 'baz', + type: 'integer', + location: { min: 0, max: 0 }, + userDefined: true, + }, + ...baseColumns, + ]); + }); + + it('adds multiple new columns for multiple assignments', () => { + const command = synth.cmd`EVAL baz = foo + 1, qux = bar`; + const result = columnsAfter(command, baseColumns, ''); + + expect(result).toEqual([ + { + name: 'baz', + type: 'integer', + location: { min: 0, max: 0 }, + userDefined: true, + }, + { + name: 'qux', + type: 'keyword', + location: { min: 0, max: 0 }, + userDefined: true, + }, + ...baseColumns, + ]); + }); + + it('adds a column for a single expression (not assignment)', () => { + const queryString = `FROM index | EVAL foo + 1`; + + // Can't use synth because it steps on the location information + // which is used to determine the name of the new column + const { + root: { + commands: [, command], + }, + } = Parser.parseQuery(queryString); + + const result = columnsAfter(command, baseColumns, queryString); + + expect(result).toEqual([ + { + name: 'foo + 1', + type: 'integer', + location: { min: 18, max: 24 }, + userDefined: true, + }, + ...baseColumns, + ]); + }); + + it('handles mix of assignments and expressions', () => { + const queryString = `FROM index | EVAL baz = foo + 1, TRIM(bar)`; + + // Can't use synth because it steps on the location information + // which is used to determine the name of the new column + const { + root: { + commands: [, command], + }, + } = Parser.parseQuery(queryString); + + const result = columnsAfter(command, baseColumns, queryString); + + expect(result).toEqual([ + { + name: 'baz', + type: 'integer', + location: { min: 18, max: 20 }, + userDefined: true, + }, + { + name: 'TRIM(bar)', + type: 'keyword', + location: { min: 33, max: 41 }, + userDefined: true, + }, + ...baseColumns, + ]); + }); + + it('returns previous columns if no args', () => { + const command = { args: [] } as any; + const result = columnsAfter(command, baseColumns, ''); + + expect(result).toEqual(baseColumns); + }); + + it('handles overwriting columns', () => { + const command = synth.cmd`EVAL foo = "", bar = 23`; + const result = columnsAfter(command, baseColumns, ''); + + expect(result).toEqual([ + { + name: 'foo', + type: 'keyword', // originally integer + location: { min: 0, max: 0 }, + userDefined: true, + }, + { + name: 'bar', + type: 'integer', // originally keyword + location: { min: 0, max: 0 }, + userDefined: true, + }, + ]); + }); +}); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts new file mode 100644 index 0000000000000..0e500606a63f6 --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/columns_after.ts @@ -0,0 +1,50 @@ +/* + * 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 { uniqBy } from 'lodash'; +import { isAssignment, isColumn } from '../../../ast/is'; +import { getExpressionType } from '../../../definitions/utils'; +import type { ESQLAstItem, ESQLCommand } from '../../../types'; +import type { ESQLColumnData, ESQLUserDefinedColumn } from '../../types'; + +export const columnsAfter = ( + command: ESQLCommand, + previousColumns: ESQLColumnData[], + query: string +) => { + const columnMap = new Map(); + previousColumns.forEach((col) => columnMap.set(col.name, col)); // TODO make this more efficient + + const typeOf = (thing: ESQLAstItem) => getExpressionType(thing, columnMap); + + const newColumns = []; + + for (const expression of command.args) { + if (isAssignment(expression) && isColumn(expression.args[0])) { + const name = expression.args[0].parts.join('.'); + const newColumn: ESQLUserDefinedColumn = { + name, + type: typeOf(expression.args[1]), + location: expression.args[0].location, + userDefined: true, + }; + newColumns.push(newColumn); + } else if (!Array.isArray(expression)) { + const newColumn: ESQLUserDefinedColumn = { + name: query.substring(expression.location.min, expression.location.max + 1), + type: typeOf(expression), + location: expression.location, + userDefined: true, + }; + newColumns.push(newColumn); + } + } + + return uniqBy([...newColumns, ...previousColumns], 'name'); +}; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/index.ts index 7804a7069183a..18c649ab6e348 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/index.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/eval/index.ts @@ -10,11 +10,13 @@ import { i18n } from '@kbn/i18n'; import type { ICommandMethods } from '../../registry'; import { autocomplete } from './autocomplete'; import { validate } from './validate'; +import { columnsAfter } from './columns_after'; import type { ICommandContext } from '../../types'; const evalCommandMethods: ICommandMethods = { autocomplete, validate, + columnsAfter, }; export const evalCommand = { diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts index 850321b897fa1..5d0c8803d1ed8 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.test.ts @@ -7,36 +7,231 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import { synth } from '../../../..'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; +import type { ESQLColumnData } from '../../types'; import { columnsAfter } from './columns_after'; describe('FORK', () => { - const context = { - userDefinedColumns: new Map([]), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }; - it('adds the _fork in the list of fields', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - ] as ESQLFieldWithMetadata[]; - - const result = columnsAfter( + it('adds the _fork in the list of fields', async () => { + const previousCommandFields: ESQLColumnData[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; + + const result = await columnsAfter( synth.cmd`FORK (LIMIT 10 ) (LIMIT 1000 ) `, previousCommandFields, - context + '', + { + fromEnrich: () => Promise.resolve([]), + fromJoin: () => Promise.resolve([]), + fromFrom: () => Promise.resolve([]), + } + ); + + expect(result).toEqual([ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + { + name: '_fork', + type: 'keyword', + userDefined: false, + }, + ]); + }); + + it('collects columns from branches', async () => { + const previousCommandFields: ESQLColumnData[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; + + const result = await columnsAfter( + synth.cmd`FORK (EVAL foo = 1 | RENAME foo AS bar) (EVAL lolz = 2 + 3 | EVAL field1 = 2.) `, + previousCommandFields, + '', + { + fromEnrich: () => Promise.resolve([]), + fromJoin: () => Promise.resolve([]), + fromFrom: () => Promise.resolve([]), + } + ); + + expect(result).toEqual([ + { name: 'bar', type: 'integer', userDefined: true, location: { min: 0, max: 0 } }, + { name: 'field1', type: 'double', userDefined: true, location: { min: 0, max: 0 } }, + { name: 'field2', type: 'double', userDefined: false }, + { name: 'lolz', type: 'integer', userDefined: true, location: { min: 0, max: 0 } }, + { + name: '_fork', + type: 'keyword', + userDefined: false, + }, + ]); + }); + + it('supports JOIN and ENRICH', async () => { + const previousCommandFields: ESQLColumnData[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; + + const result = await columnsAfter( + synth.cmd`FORK (LOOKUP JOIN lookup-index ON joinField) (ENRICH policy ON matchField)`, + previousCommandFields, + '', + { + fromEnrich: () => + Promise.resolve([ + { + name: 'from-enrich', + type: 'keyword', + userDefined: false, + }, + ]), + fromJoin: () => + Promise.resolve([ + { + name: 'from-join', + type: 'keyword', + userDefined: false, + }, + ]), + fromFrom: () => + Promise.resolve([ + { + name: 'from-from', + type: 'keyword', + userDefined: false, + }, + ]), + } + ); + + expect(result).toEqual([ + { name: 'from-join', type: 'keyword', userDefined: false }, + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + { name: 'from-enrich', type: 'keyword', userDefined: false }, + { + name: '_fork', + type: 'keyword', + userDefined: false, + }, + ]); + }); + + it('prefers userDefined columns over fields with the same name', async () => { + const previousCommandFields: ESQLColumnData[] = [ + { name: 'foo', type: 'keyword', userDefined: false }, + { name: 'bar', type: 'double', userDefined: false }, + ]; + + // Branch 1: foo is userDefined, bar is not + // Branch 2: foo is not userDefined, bar is userDefined + const result = await columnsAfter( + synth.cmd`FORK (EVAL foo = 1) (EVAL bar = 2)`, + previousCommandFields, + '', + { + fromEnrich: () => Promise.resolve([]), + fromJoin: () => Promise.resolve([]), + fromFrom: () => Promise.resolve([]), + } ); - expect(result).toEqual([ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, + // foo from branch 1 is userDefined, bar from branch 2 is userDefined + expect(result).toEqual([ + { name: 'foo', type: 'integer', userDefined: true, location: { min: 0, max: 0 } }, + { name: 'bar', type: 'integer', userDefined: true, location: { min: 0, max: 0 } }, { name: '_fork', type: 'keyword', + userDefined: false, }, ]); }); + + describe('conflicts between branches', () => { + it('keeps the first userDefined column if both branches define userDefined columns with the same name', async () => { + const previousCommandFields: ESQLColumnData[] = [ + { name: 'foo', type: 'keyword', userDefined: false }, + ]; + + // Both branches define foo as userDefined, but with different types + const result = await columnsAfter( + synth.cmd`FORK (EVAL foo = 1) (EVAL foo = 2.5)`, + previousCommandFields, + '', + { + fromEnrich: () => Promise.resolve([]), + fromJoin: () => Promise.resolve([]), + fromFrom: () => Promise.resolve([]), + } + ); + + // The first userDefined column wins (type: integer) + expect(result).toEqual([ + { name: 'foo', type: 'integer', userDefined: true, location: { min: 0, max: 0 } }, + { + name: '_fork', + type: 'keyword', + userDefined: false, + }, + ]); + }); + + it("doesn't duplicate fields from branches", async () => { + const previousCommandFields: ESQLColumnData[] = [ + { name: 'foo', type: 'keyword', userDefined: false }, + ]; + + // All branches keep foo as a field, but with different types + const result = await columnsAfter( + synth.cmd`FORK (KEEP foo) (KEEP foo) (KEEP foo)`, + previousCommandFields, + '', + { + fromEnrich: () => Promise.resolve([]), + fromJoin: () => Promise.resolve([]), + fromFrom: () => Promise.resolve([]), + } + ); + + expect(result).toEqual([ + { name: 'foo', type: 'keyword', userDefined: false }, + { + name: '_fork', + type: 'keyword', + userDefined: false, + }, + ]); + }); + + it('prefers userDefined column if one branch overwrites a field', async () => { + const previousCommandFields: ESQLColumnData[] = [ + { name: 'foo', type: 'keyword', userDefined: false }, + ]; + + // Branch 1: foo is userDefined, Branch 2: foo is not userDefined + const result = await columnsAfter( + synth.cmd`FORK (EVAL foo = 1) (LIMIT 1)`, + previousCommandFields, + '', + { + fromEnrich: () => Promise.resolve([]), + fromJoin: () => Promise.resolve([]), + fromFrom: () => Promise.resolve([]), + } + ); + + expect(result).toEqual([ + { name: 'foo', type: 'integer', userDefined: true, location: { min: 0, max: 0 } }, + { + name: '_fork', + type: 'keyword', + userDefined: false, + }, + ]); + }); + }); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts index a44beb912b818..36cecba91e721 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/columns_after.ts @@ -7,22 +7,87 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import { uniqBy } from 'lodash'; +import { esqlCommandRegistry } from '../../../..'; +import type { ESQLAstQueryExpression } from '../../../types'; import { type ESQLCommand } from '../../../types'; -import type { ESQLFieldWithMetadata } from '../../types'; -import type { ICommandContext } from '../../types'; +import type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types'; +import type { IAdditionalFields } from '../../registry'; -export const columnsAfter = ( +export const columnsAfter = async ( command: ESQLCommand, - previousColumns: ESQLFieldWithMetadata[], - context?: ICommandContext + previousColumns: ESQLColumnData[], + query: string, + additionalFields: IAdditionalFields ) => { + const branches = command.args as ESQLAstQueryExpression[]; + + const columnsFromBranches = []; + + for (const branch of branches) { + // start with columns from before FORK + let columnsFromBranch = [...previousColumns]; + for (const branchCommand of branch.commands) { + const commandDef = esqlCommandRegistry.getCommandByName(branchCommand.name); + if (commandDef?.methods?.columnsAfter) { + columnsFromBranch = await commandDef.methods?.columnsAfter?.( + branchCommand, + columnsFromBranch, + query, + additionalFields + ); + } + } + + columnsFromBranches.push(columnsFromBranch); + } + + const maps = columnsFromBranches.map((cols) => new Map(cols.map((_col) => [_col.name, _col]))); + + const merged = new Map(); + + // O(b * n), where b is the branches and n is the max number of columns in a branch + for (const map of maps) { + for (const [name, colData] of map) { + /** + * Check for conflicts... + * + * The branches often produce large numbers of columns with the same names + * because they will often include most or all fields in their results. + * + * Sometimes, user-defined columns have been created that overwrite + * fields with the same name, in which case user-defined columns should win. + * + * In other cases, the branches may define columns of the same name but with different + * data types. That is not allowed but, to keep things simple, we don't touch that case... + * we leave it to Elasticsearch validation. + */ + if (merged.has(name)) { + const existingColData = merged.get(name) as ESQLColumnData; + if (existingColData.userDefined) { + // If the existing column is user-defined and the new one is not, keep the existing one + continue; + } else if (!existingColData.userDefined && colData.userDefined) { + // If the existing column is not user-defined and the new one is, use the new one + merged.set(name, colData); + continue; + } else { + // If both columns are user-defined or both are not, keep the existing one + continue; + } + } + + merged.set(name, colData); + } + } + return uniqBy( [ - ...previousColumns, + ...merged.values(), { name: '_fork', type: 'keyword' as const, - }, + userDefined: false, + } as ESQLFieldWithMetadata, ], 'name' ); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts index f070903d6aa37..7f26d94fc916b 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fork/validate.ts @@ -68,10 +68,5 @@ export const validate = ( messages.push(errors.tooManyForks(forks[1])); } - context?.fields.set('_fork', { - name: '_fork', - type: 'keyword', - }); - return messages; }; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/autocomplete.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/autocomplete.test.ts index a1f708fc05d87..ff0b42f940e13 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/autocomplete.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/autocomplete.test.ts @@ -6,14 +6,14 @@ * 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 { mockContext, integrations } from '../../../__tests__/context_fixtures'; -import { autocomplete } from './autocomplete'; import { expectSuggestions, getFieldNamesByType } from '../../../__tests__/autocomplete'; -import type { ICommandCallbacks } from '../../types'; +import { integrations, mockContext } from '../../../__tests__/context_fixtures'; import { correctQuerySyntax, findAstPosition } from '../../../definitions/utils/ast'; import { parse } from '../../../parser'; -import { getRecommendedQueriesTemplates } from '../../options/recommended_queries'; import { METADATA_FIELDS } from '../../options/metadata'; +import { getRecommendedQueriesTemplates } from '../../options/recommended_queries'; +import type { ICommandCallbacks } from '../../types'; +import { autocomplete } from './autocomplete'; const metadataFields = [...METADATA_FIELDS].sort(); @@ -57,7 +57,7 @@ describe('FROM Autocomplete', () => { getByType: jest.fn(), }; - const expectedFields = getFieldNamesByType('any'); + const expectedFields = getFieldNamesByType('any', true); (mockCallbacks.getByType as jest.Mock).mockResolvedValue( expectedFields.map((name) => ({ label: name, text: name })) ); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.test.ts new file mode 100644 index 0000000000000..6cda1b953ea20 --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.test.ts @@ -0,0 +1,28 @@ +/* + * 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 type { ESQLFieldWithMetadata } from '../../types'; +import { columnsAfter } from './columns_after'; + +describe('FROM columnsAfter', () => { + it('returns fields from the source', async () => { + const sourceFields: ESQLFieldWithMetadata[] = [ + { name: 'field1', type: 'double', userDefined: false }, + { name: 'field2', type: 'keyword', userDefined: false }, + ]; + + const result = await columnsAfter({} as any, [], '', { + fromFrom: () => Promise.resolve(sourceFields), + fromJoin: () => Promise.resolve([]), + fromEnrich: () => Promise.resolve([]), + }); + + expect(result).toEqual(sourceFields); + }); +}); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.ts new file mode 100644 index 0000000000000..6e5e8502964fa --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/columns_after.ts @@ -0,0 +1,20 @@ +/* + * 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 { type ESQLCommand } from '../../../types'; +import type { ESQLColumnData } from '../../types'; +import type { IAdditionalFields } from '../../registry'; + +export const columnsAfter = ( + command: ESQLCommand, + _previousColumns: ESQLColumnData[], // will always be empty for FROM + _query: string, + additionalFields: IAdditionalFields +) => { + return additionalFields.fromFrom(command); +}; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/index.ts index b5551e073892d..16fd887e4a34c 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/index.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/from/index.ts @@ -11,10 +11,12 @@ import type { ICommandMethods } from '../../registry'; import { autocomplete } from './autocomplete'; import { validate } from './validate'; import type { ICommandContext } from '../../types'; +import { columnsAfter } from './columns_after'; const fromCommandMethods: ICommandMethods = { autocomplete, validate, + columnsAfter, }; export const fromCommand = { diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.test.ts index 0480c05cb0aa5..d2905ab15923a 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.test.ts @@ -6,10 +6,10 @@ * 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 { METADATA_FIELDS } from '../../../..'; import { mockContext } from '../../../__tests__/context_fixtures'; -import { validate } from './validate'; import { expectErrors } from '../../../__tests__/validation'; -import { METADATA_FIELDS } from '../../../..'; +import { validate } from './validate'; const fuseExpectErrors = (query: string, expectedErrors: string[], context = mockContext) => { return expectErrors(query, expectedErrors, context, 'fuse', validate); @@ -22,13 +22,13 @@ describe('FUSE Validation', () => { describe('FUSE', () => { test('no errors for valid command', () => { - const newFields = new Map(mockContext.fields); + const newColumns = new Map(mockContext.columns); METADATA_FIELDS.forEach((fieldName) => { - newFields.set(fieldName, { name: fieldName, type: 'keyword' }); + newColumns.set(fieldName, { name: fieldName, type: 'keyword', userDefined: false }); }); const context = { ...mockContext, - fields: newFields, + columns: newColumns, }; fuseExpectErrors( `FROM index METADATA _id, _score, _index diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.ts index c580ccb6699d9..7a8a5464e1d85 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/fuse/validate.ts @@ -17,15 +17,15 @@ export const validate = ( ): ESQLMessage[] => { const messages: ESQLMessage[] = []; - if (!context?.fields.get('_id')) { + if (!context?.columns.get('_id')) { messages.push(buildMissingMetadataMessage(command, '_id')); } - if (!context?.fields.get('_index')) { + if (!context?.columns.get('_index')) { messages.push(buildMissingMetadataMessage(command, '_index')); } - if (!context?.fields.get('_score')) { + if (!context?.columns.get('_score')) { messages.push(buildMissingMetadataMessage(command, '_score')); } diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts index ffa3b48d1f53d..1770473508517 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.test.ts @@ -7,7 +7,7 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import { synth } from '../../../..'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; +import type { ESQLFieldWithMetadata } from '../../types'; import { columnsAfter, extractSemanticsFromGrok } from './columns_after'; describe('GROK', () => { @@ -40,29 +40,22 @@ describe('GROK', () => { }); }); describe('columnsAfter', () => { - const context = { - userDefinedColumns: new Map([]), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }; it('adds the GROK columns from the pattern in the list', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - ] as ESQLFieldWithMetadata[]; + const previousCommandFields: ESQLFieldWithMetadata[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; const result = columnsAfter( synth.cmd`GROK agent "%{WORD:firstWord}"`, previousCommandFields, - context + '' ); expect(result).toEqual([ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - { name: 'firstWord', type: 'keyword' }, + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + { name: 'firstWord', type: 'keyword', userDefined: false }, ]); }); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts index a2803afe17154..70cd7d18af47e 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/columns_after.ts @@ -6,10 +6,9 @@ * 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 { walk } from '../../../walker'; import { type ESQLCommand } from '../../../types'; -import type { ESQLFieldWithMetadata } from '../../types'; -import type { ICommandContext } from '../../types'; +import { walk } from '../../../walker'; +import type { ESQLColumnData } from '../../types'; function unquoteTemplate(inputString: string): string { if (inputString.startsWith('"') && inputString.endsWith('"') && inputString.length >= 2) { @@ -50,8 +49,8 @@ export function extractSemanticsFromGrok(pattern: string): string[] { export const columnsAfter = ( command: ESQLCommand, - previousColumns: ESQLFieldWithMetadata[], - context?: ICommandContext + previousColumns: ESQLColumnData[], + query: string ) => { const columns: string[] = []; @@ -64,6 +63,8 @@ export const columnsAfter = ( return [ ...previousColumns, - ...columns.map((column) => ({ name: column, type: 'keyword' as const })), + ...columns.map( + (column) => ({ name: column, type: 'keyword' as const, userDefined: false } as ESQLColumnData) + ), ]; }; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/validate.ts index 8b8100afd4545..d522b97c5ae11 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/validate.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/grok/validate.ts @@ -21,9 +21,9 @@ export const validate = ( const messages: ESQLMessage[] = []; const acceptedColumnTypes: FieldType[] = ['keyword', 'text']; const astCol = command.args[0] as ESQLColumn; - const columnRef = context?.fields.get(astCol.name); + const columnRef = context?.columns.get(astCol.name); - if (columnRef && !acceptedColumnTypes.includes(columnRef.type)) { + if (columnRef && !acceptedColumnTypes.includes(columnRef.type as FieldType)) { messages.push( getMessageFromId({ messageId: 'unsupportedColumnTypeForCommand', diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts index cb15a668d4a37..875695d5cd378 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.test.ts @@ -6,137 +6,35 @@ * 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 { synth } from '../../../..'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; +import { Parser } from '../../../..'; +import type { ESQLColumnData } from '../../types'; +import { type ESQLFieldWithMetadata } from '../../types'; import { columnsAfter } from './columns_after'; -// Test data fixtures -const createPreviousFields = (fields: Array<[string, string]>): ESQLFieldWithMetadata[] => - fields.map(([name, type]) => ({ name, type } as ESQLFieldWithMetadata)); - describe('INLINESTATS', () => { - const createUserDefinedColumn = ( - name: string, - type: ESQLFieldWithMetadata['type'], - location = { min: 0, max: 10 } - ): ESQLUserDefinedColumn => ({ name, type, location }); - - const createContext = (userDefinedColumns: Array<[string, ESQLUserDefinedColumn[]]> = []) => ({ - userDefinedColumns: new Map(userDefinedColumns), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }); - - const expectColumnsAfter = ( - command: string, - previousFields: ESQLFieldWithMetadata[], - userColumns: Array<[string, string]>, - expectedResult: Array<[string, string]> - ) => { - const context = createContext( - userColumns.map(([name, type]) => [name, [createUserDefinedColumn(name, type as 'keyword')]]) - ); - const result = columnsAfter(synth.cmd(command), previousFields, context); - const expected = createPreviousFields(expectedResult); - expect(result).toEqual(expected); - }; - it('preserves all previous columns and adds the user defined column, when no grouping is given', () => { - const previousFields = createPreviousFields([ - ['field1', 'keyword'], - ['field2', 'double'], + it('gets the columns after the query', () => { + const previousCommandFields: ESQLFieldWithMetadata[] = [ + { name: 'field1', type: 'double', userDefined: false }, + { name: 'buckets', type: 'double', userDefined: false }, // should be overwritten + { name: '@timestamp', type: 'date', userDefined: false }, + ]; + + const queryString = `FROM a | INLINESTATS AVG(field1) BY buckets=BUCKET(@timestamp,50,?_tstart,?_tend)`; + + // Can't use synth because it steps on the location information + // which is used to determine the name of the new column + const { + root: { + commands: [, command], + }, + } = Parser.parseQuery(queryString); + + const result = columnsAfter(command, previousCommandFields, queryString); + + expect(result).toEqual([ + { name: 'AVG(field1)', type: 'double', userDefined: true, location: { min: 21, max: 31 } }, + { name: 'buckets', type: 'date', userDefined: true, location: { min: 36, max: 42 } }, + ...previousCommandFields.filter(({ name }) => name !== 'buckets'), ]); - - expectColumnsAfter( - 'INLINESTATS var0=AVG(field2)', - previousFields, - [['var0', 'double']], - [ - ['field1', 'keyword'], - ['field2', 'double'], - ['var0', 'double'], - ] - ); - }); - - it('preserves all previous columns and adds the escaped column, when no grouping is given', () => { - const previousFields = createPreviousFields([ - ['field1', 'keyword'], - ['field2', 'double'], - ]); - - expectColumnsAfter( - 'INLINESTATS AVG(field2)', - previousFields, - [['AVG(field2)', 'double']], - [ - ['field1', 'keyword'], - ['field2', 'double'], - ['AVG(field2)', 'double'], - ] - ); - }); - - it('preserves all previous columns and adds the escaped column, with grouping', () => { - const previousFields = createPreviousFields([ - ['field1', 'keyword'], - ['field2', 'double'], - ]); - - // Note: Unlike STATS, INLINESTATS doesn't care about BY clause for column preservation - expectColumnsAfter( - 'INLINESTATS AVG(field2) BY field1', - previousFields, - [['AVG(field2)', 'double']], - [ - ['field1', 'keyword'], - ['field2', 'double'], - ['AVG(field2)', 'double'], - ] - ); - }); - - it('preserves all previous columns and adds user defined and grouping columns', () => { - const previousFields = createPreviousFields([ - ['field1', 'keyword'], - ['field2', 'double'], - ['@timestamp', 'date'], - ]); - - expectColumnsAfter( - 'INLINESTATS AVG(field2) BY buckets=BUCKET(@timestamp,50,?_tstart,?_tend)', - previousFields, - [ - ['AVG(field2)', 'double'], - ['buckets', 'unknown'], - ], - [ - ['field1', 'keyword'], - ['field2', 'double'], - ['@timestamp', 'date'], - ['AVG(field2)', 'double'], - ['buckets', 'unknown'], - ] - ); - }); - - it('handles duplicate column names by keeping the original column type', () => { - const previousFields = createPreviousFields([ - ['field1', 'keyword'], - ['field2', 'double'], - ['avg_field', 'integer'], // This will be preserved since it comes first - ]); - - expectColumnsAfter( - 'INLINESTATS avg_field=AVG(field2)', - previousFields, - [['avg_field', 'double']], // This will be ignored due to duplicate name - [ - ['field1', 'keyword'], - ['field2', 'double'], - ['avg_field', 'integer'], - ] - ); }); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts index 0d3fd9d787926..b599f132cbb68 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/columns_after.ts @@ -8,39 +8,15 @@ */ import uniqBy from 'lodash/uniqBy'; import type { ESQLCommand } from '../../../types'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; -import type { ICommandContext } from '../../types'; -import type { FieldType } from '../../../definitions/types'; - -function transformMapToESQLFields( - inputMap: Map -): ESQLFieldWithMetadata[] { - const esqlFields: ESQLFieldWithMetadata[] = []; - - for (const [, userDefinedColumns] of inputMap) { - for (const userDefinedColumn of userDefinedColumns) { - if (userDefinedColumn.type) { - esqlFields.push({ - name: userDefinedColumn.name, - type: userDefinedColumn.type as FieldType, - }); - } - } - } - - return esqlFields; -} +import { columnsAfter as columnsAfterStats } from '../stats/columns_after'; +import type { ESQLColumnData } from '../../types'; export const columnsAfter = ( - _command: ESQLCommand, - previousColumns: ESQLFieldWithMetadata[], - context?: ICommandContext + command: ESQLCommand, + previousColumns: ESQLColumnData[], + query: string ) => { - const userDefinedColumns = - context?.userDefinedColumns ?? new Map(); - - const arrayOfUserDefinedColumns: ESQLFieldWithMetadata[] = - transformMapToESQLFields(userDefinedColumns); + const newColumns = columnsAfterStats(command, previousColumns, query); - return uniqBy([...previousColumns, ...arrayOfUserDefinedColumns], 'name'); + return uniqBy([...newColumns, ...previousColumns], 'name'); }; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/validate.test.ts index 4b47b7a2f19e1..de9cf10e3d76c 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/validate.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/inlinestats/validate.test.ts @@ -26,21 +26,19 @@ describe('INLINESTATS Validation', () => { }); describe('INLINESTATS [ BY ]', () => { - const newUserDefinedColumns = new Map(mockContext.userDefinedColumns); - newUserDefinedColumns.set('doubleField * 3.281', [ - { - name: 'doubleField * 3.281', - type: 'double', - location: { min: 0, max: 10 }, - }, - ]); - newUserDefinedColumns.set('avg_doubleField', [ - { - name: 'avg_doubleField', - type: 'double', - location: { min: 0, max: 10 }, - }, - ]); + const newColumns = new Map(mockContext.columns); + newColumns.set('doubleField * 3.281', { + name: 'doubleField * 3.281', + type: 'double', + location: { min: 0, max: 10 }, + userDefined: true, + }); + newColumns.set('avg_doubleField', { + name: 'avg_doubleField', + type: 'double', + location: { min: 0, max: 10 }, + userDefined: true, + }); describe('... ...', () => { test('no errors on correct usage', () => { diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts new file mode 100644 index 0000000000000..3b2d34c78b92c --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.test.ts @@ -0,0 +1,70 @@ +/* + * 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 type { ESQLColumnData, ESQLFieldWithMetadata } from '../../types'; +import { columnsAfter } from './columns_after'; + +describe('JOIN columnsAfter', () => { + it('returns previousColumns when no join columns', async () => { + const previousColumns: ESQLColumnData[] = [ + { name: 'fieldA', type: 'keyword', userDefined: false }, + { name: 'fieldB', type: 'long', userDefined: false }, + ]; + + const result = await columnsAfter({} as any, previousColumns, '', { + fromJoin: () => Promise.resolve([]), + fromEnrich: () => Promise.resolve([]), + fromFrom: () => Promise.resolve([]), + }); + + expect(result).toEqual(previousColumns); + }); + + it('prepends joinColumns to previousColumns when joinColumns is provided', async () => { + const previousColumns: ESQLColumnData[] = [ + { name: 'fieldA', type: 'keyword', userDefined: false }, + { name: 'fieldB', type: 'long', userDefined: false }, + ]; + const joinColumns: ESQLFieldWithMetadata[] = [ + { name: 'joinField1', type: 'keyword', userDefined: false }, + { name: 'joinField2', type: 'double', userDefined: false }, + ]; + + const result = await columnsAfter({} as any, previousColumns, '', { + fromJoin: () => Promise.resolve(joinColumns), + fromEnrich: () => Promise.resolve([]), + fromFrom: () => Promise.resolve([]), + }); + + expect(result).toEqual([...joinColumns, ...previousColumns]); + }); + + it('overwrites previous columns with the same name', async () => { + const previousColumns: ESQLColumnData[] = [ + { name: 'fieldA', type: 'keyword', userDefined: false }, + { name: 'fieldB', type: 'long', userDefined: false }, + ]; + const joinColumns: ESQLFieldWithMetadata[] = [ + { name: 'fieldA', type: 'text', userDefined: false }, + { name: 'fieldC', type: 'double', userDefined: false }, + ]; + + const result = await columnsAfter({} as any, previousColumns, '', { + fromJoin: () => Promise.resolve(joinColumns), + fromEnrich: () => Promise.resolve([]), + fromFrom: () => Promise.resolve([]), + }); + + expect(result).toEqual([ + { name: 'fieldA', type: 'text', userDefined: false }, + { name: 'fieldC', type: 'double', userDefined: false }, + { name: 'fieldB', type: 'long', userDefined: false }, + ]); + }); +}); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts new file mode 100644 index 0000000000000..b437293906aae --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/columns_after.ts @@ -0,0 +1,22 @@ +/* + * 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 { uniqBy } from 'lodash'; +import { type ESQLCommand } from '../../../types'; +import type { ESQLColumnData } from '../../types'; +import type { IAdditionalFields } from '../../registry'; + +export const columnsAfter = async ( + command: ESQLCommand, + previousColumns: ESQLColumnData[], + query: string, + additionalFields: IAdditionalFields +) => { + const joinFields = await additionalFields.fromJoin(command); + return uniqBy([...(joinFields ?? []), ...previousColumns], 'name'); +}; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/index.ts index db2fac4bc5556..44d6674dc7b07 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/index.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/index.ts @@ -11,10 +11,12 @@ import type { ICommandMethods } from '../../registry'; import { autocomplete } from './autocomplete'; import { validate } from './validate'; import type { ICommandContext } from '../../types'; +import { columnsAfter } from './columns_after'; const joinCommandMethods: ICommandMethods = { validate, autocomplete, + columnsAfter, }; export const joinCommand = { diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/utils.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/utils.ts index 643cba847dcad..cd1c7bb2f923c 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/utils.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/join/utils.ts @@ -8,23 +8,24 @@ */ import { i18n } from '@kbn/i18n'; import { uniqBy } from 'lodash'; -import { handleFragment, columnExists } from '../../../definitions/utils/autocomplete/helpers'; +import { buildFieldsDefinitionsWithMetadata } from '../../../definitions/utils'; +import { isColumn } from '../../../ast/is'; +import { columnExists, handleFragment } from '../../../definitions/utils/autocomplete/helpers'; import { unescapeColumnName } from '../../../definitions/utils/shared'; import * as mutate from '../../../mutate'; import { LeafPrinter } from '../../../pretty_print/leaf_printer'; -import { pipeCompleteItem, commaCompleteItem } from '../../complete_items'; -import { buildFieldsDefinitionsWithMetadata } from '../../../definitions/utils/functions'; -import type { ICommand } from '../../registry'; import type { ESQLAstJoinCommand, ESQLCommand, ESQLCommandOption } from '../../../types'; +import { commaCompleteItem, pipeCompleteItem } from '../../complete_items'; +import { TRIGGER_SUGGESTION_COMMAND } from '../../constants'; +import type { ICommand } from '../../registry'; import type { + ESQLColumnData, ESQLFieldWithMetadata, GetColumnsByTypeFn, - ISuggestionItem, ICommandContext, + ISuggestionItem, } from '../../types'; import type { JoinCommandPosition, JoinPosition, JoinStaticPosition } from './types'; -import { TRIGGER_SUGGESTION_COMMAND } from '../../constants'; -import { isColumn } from '../../../ast/is'; const REGEX = /^(?\w+((?\s+((?(JOIN|JOI|JO|J)((?\s+((?\S+((?\s+(?(AS|A))?(?\s+(((?\S+)?(?\s+)?)?))?((?(ON|O)((?\s+(?[^\s])?)?))?))?))?))?))?))?))?/i; @@ -153,7 +154,7 @@ export const suggestFields = async ( innerText: string, command: ESQLCommand, getColumnsByType: GetColumnsByTypeFn, - getColumnsForQuery: (query: string) => Promise, + getColumnsForQuery: (query: string) => Promise, context?: ICommandContext ) => { if (!context) { @@ -163,7 +164,9 @@ export const suggestFields = async ( const { suggestions: fieldSuggestions, lookupIndexFieldExists } = await getFieldSuggestions( command, getColumnsByType, - getColumnsForQuery, + // this type cast is ok because getFieldSuggestions only ever fetches columns + // from a bare FROM clause, so they will always be fields, not user-defined columns + getColumnsForQuery as (query: string) => Promise, context ); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/autocomplete.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/autocomplete.test.ts index 1e71d0718b78d..3427910c9c9b6 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/autocomplete.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/autocomplete.test.ts @@ -33,15 +33,11 @@ describe('KEEP Autocomplete', () => { }); it('suggests available fields after KEEP', async () => { - const fieldsMap = mockContext.fields; - const userDefinedColumns = mockContext.userDefinedColumns; - const allFields = [ - ...Array.from(fieldsMap.values()), - ...Array.from(userDefinedColumns.values()).flat(), - ]; + const columns = mockContext.columns; + keepExpectSuggestions( 'FROM a | KEEP ', - allFields.map((field) => field.name) + Array.from(columns.values()).map((column) => column.name) ); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts index 9c602ba20e3b2..b94cd73b215cc 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.test.ts @@ -7,25 +7,18 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import { synth } from '../../../..'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; +import type { ESQLFieldWithMetadata } from '../../types'; import { columnsAfter } from './columns_after'; -describe('FORK', () => { - const context = { - userDefinedColumns: new Map([]), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }; +describe('KEEP', () => { it('should return the correct fields after the command', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - ] as ESQLFieldWithMetadata[]; + const previousCommandFields: ESQLFieldWithMetadata[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; - const result = columnsAfter(synth.cmd`KEEP field1`, previousCommandFields, context); + const result = columnsAfter(synth.cmd`KEEP field1`, previousCommandFields, ''); - expect(result).toEqual([{ name: 'field1', type: 'keyword' }]); + expect(result).toEqual([{ name: 'field1', type: 'keyword', userDefined: false }]); }); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts index 6d0ab16aba080..72235cd4a8f3c 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/columns_after.ts @@ -6,15 +6,14 @@ * 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 { walk } from '../../../walker'; import { type ESQLCommand } from '../../../types'; -import type { ESQLFieldWithMetadata } from '../../types'; -import type { ICommandContext } from '../../types'; +import { walk } from '../../../walker'; +import type { ESQLColumnData } from '../../types'; export const columnsAfter = ( command: ESQLCommand, - previousColumns: ESQLFieldWithMetadata[], - context?: ICommandContext + previousColumns: ESQLColumnData[], + query: string ) => { const columnsToKeep: string[] = []; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/validate.test.ts index 54f89a7341a80..b9ad1da948ff3 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/validate.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/keep/validate.test.ts @@ -20,24 +20,22 @@ describe('KEEP Validation', () => { }); test('validates the most basic query', () => { - const newUserDefinedColumns = new Map(mockContext.userDefinedColumns); - newUserDefinedColumns.set('MIN(doubleField * 10)', [ - { - name: 'MIN(doubleField * 10)', - type: 'double', - location: { min: 0, max: 10 }, - }, - ]); - newUserDefinedColumns.set('COUNT(*)', [ - { - name: 'COUNT(*)', - type: 'integer', - location: { min: 0, max: 10 }, - }, - ]); + const newColumns = new Map(mockContext.columns); + newColumns.set('MIN(doubleField * 10)', { + name: 'MIN(doubleField * 10)', + type: 'double', + location: { min: 0, max: 10 }, + userDefined: true, + }); + newColumns.set('COUNT(*)', { + name: 'COUNT(*)', + type: 'integer', + location: { min: 0, max: 10 }, + userDefined: true, + }); const context = { ...mockContext, - userDefinedColumns: newUserDefinedColumns, + columns: newColumns, }; keepExpectErrors('from index | keep keywordField, doubleField, integerField, dateField', []); keepExpectErrors( diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts index ae475e70550d2..bf63511267fb6 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.test.ts @@ -7,60 +7,53 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import { synth } from '../../../..'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; +import type { ESQLColumnData } from '../../types'; import { columnsAfter } from './columns_after'; describe('RENAME', () => { - const context = { - userDefinedColumns: new Map([]), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }; it('renames the given columns with the new names using AS', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - ] as ESQLFieldWithMetadata[]; + const previousCommandFields: ESQLColumnData[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; - const result = columnsAfter(synth.cmd`RENAME field1 as meow`, previousCommandFields, context); + const result = columnsAfter(synth.cmd`RENAME field1 as meow`, previousCommandFields, ''); - expect(result).toEqual([ - { name: 'meow', type: 'keyword' }, - { name: 'field2', type: 'double' }, + expect(result).toEqual([ + { name: 'meow', type: 'keyword', userDefined: true, location: { max: 0, min: 0 } }, + { name: 'field2', type: 'double', userDefined: false }, ]); }); it('renames the given columns with the new names using ASSIGN', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - ] as ESQLFieldWithMetadata[]; + const previousCommandFields: ESQLColumnData[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; - const result = columnsAfter(synth.cmd`RENAME meow = field1`, previousCommandFields, context); + const result = columnsAfter(synth.cmd`RENAME meow = field1`, previousCommandFields, ''); - expect(result).toEqual([ - { name: 'meow', type: 'keyword' }, - { name: 'field2', type: 'double' }, + expect(result).toEqual([ + { name: 'meow', type: 'keyword', userDefined: true, location: { max: 0, min: 0 } }, + { name: 'field2', type: 'double', userDefined: false }, ]); }); it('renames the given columns with the new names using a mix of ASSIGN and =', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - ] as ESQLFieldWithMetadata[]; + const previousCommandFields: ESQLColumnData[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; const result = columnsAfter( synth.cmd`RENAME meow = field1, field2 as woof`, previousCommandFields, - context + '' ); - expect(result).toEqual([ - { name: 'meow', type: 'keyword' }, - { name: 'woof', type: 'double' }, + expect(result).toEqual([ + { name: 'meow', type: 'keyword', userDefined: true, location: { max: 0, min: 0 } }, + { name: 'woof', type: 'double', userDefined: true, location: { max: 0, min: 0 } }, ]); }); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts index 5c15019ff13df..31b3915a54744 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/columns_after.ts @@ -7,15 +7,14 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import uniqBy from 'lodash/uniqBy'; -import type { ESQLCommand, ESQLFunction, ESQLAstBaseItem } from '../../../types'; -import { isFunctionExpression } from '../../../ast/is'; -import type { ESQLFieldWithMetadata } from '../../types'; -import type { ICommandContext } from '../../types'; +import { isColumn, isFunctionExpression } from '../../../ast/is'; +import type { ESQLAstBaseItem, ESQLCommand, ESQLFunction } from '../../../types'; +import type { ESQLColumnData } from '../../types'; export const columnsAfter = ( command: ESQLCommand, - previousColumns: ESQLFieldWithMetadata[], - context?: ICommandContext + previousColumns: ESQLColumnData[], + query: string ) => { const asRenamePairs: ESQLFunction[] = []; const assignRenamePairs: ESQLFunction[] = []; @@ -31,21 +30,31 @@ export const columnsAfter = ( } // rename the columns with the user defined name - const newFields = previousColumns.map((oldColumn) => { + const newFields = previousColumns.map((oldColumn) => { const asRenamePair = asRenamePairs.find( (pair) => (pair?.args?.[0] as ESQLAstBaseItem)?.name === oldColumn.name ); - if (asRenamePair?.args?.[1]) { - return { name: (asRenamePair.args[1] as ESQLAstBaseItem).name, type: oldColumn.type }; + if (isColumn(asRenamePair?.args?.[1])) { + return { + name: asRenamePair.args[1].name, + type: oldColumn.type, + location: asRenamePair.args[1].location, + userDefined: true, + }; } const assignRenamePair = assignRenamePairs.find( (pair) => (pair?.args?.[1] as ESQLAstBaseItem)?.name === oldColumn.name ); - if (assignRenamePair?.args?.[0]) { - return { name: (assignRenamePair.args[0] as ESQLAstBaseItem).name, type: oldColumn.type }; + if (isColumn(assignRenamePair?.args?.[0])) { + return { + name: assignRenamePair.args[0].name, + type: oldColumn.type, + location: assignRenamePair.args[0].location, + userDefined: true, + }; } return oldColumn; // No rename found, keep the old name diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/validate.test.ts index 4f59f186c9d11..f1ec956163339 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/validate.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rename/validate.test.ts @@ -7,8 +7,8 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import { mockContext } from '../../../__tests__/context_fixtures'; -import { validate } from './validate'; import { expectErrors } from '../../../__tests__/validation'; +import { validate } from './validate'; const renameExpectErrors = (query: string, expectedErrors: string[], context = mockContext) => { return expectErrors(query, expectedErrors, context, 'rename', validate); @@ -20,24 +20,22 @@ describe('RENAME Validation', () => { }); test('validates the most basic query', () => { - const newUserDefinedColumns = new Map(mockContext.userDefinedColumns); - newUserDefinedColumns.set('doubleField + 1', [ - { - name: 'doubleField + 1', - type: 'double', - location: { min: 0, max: 10 }, - }, - ]); - newUserDefinedColumns.set('avg(doubleField)', [ - { - name: 'avg(doubleField)', - type: 'double', - location: { min: 0, max: 10 }, - }, - ]); + const newColumns = new Map(mockContext.columns); + newColumns.set('doubleField + 1', { + name: 'doubleField + 1', + type: 'double', + location: { min: 0, max: 10 }, + userDefined: true, + }); + newColumns.set('avg(doubleField)', { + name: 'avg(doubleField)', + type: 'double', + location: { min: 0, max: 10 }, + userDefined: true, + }); const context = { ...mockContext, - userDefinedColumns: newUserDefinedColumns, + columns: newColumns, }; renameExpectErrors('from a_index | rename textField as', [ 'AS expected 2 arguments, but got 1.', @@ -65,6 +63,6 @@ describe('RENAME Validation', () => { renameExpectErrors('from a_index |eval doubleField + 1 | rename `doubleField + 1` as ', [ 'AS expected 2 arguments, but got 1.', ]); - renameExpectErrors('from a_index | rename key* as keywords', ['Unknown column "keywords"']); + renameExpectErrors('from a_index | rename key* as keywords', []); }); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/validate.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/validate.ts index 0cbe21feee4dd..956ffade901ce 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/validate.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/rerank/validate.ts @@ -23,12 +23,8 @@ export const validate = ( ): ESQLMessage[] => { const messages: ESQLMessage[] = []; - const { query, targetField, location, inferenceId } = command as ESQLAstRerankCommand; - const rerankExpressionType = getExpressionType( - query, - context?.fields, - context?.userDefinedColumns - ); + const { query, location, inferenceId } = command as ESQLAstRerankCommand; + const rerankExpressionType = getExpressionType(query, context?.columns); // check for supported query types if (!supportedQueryTypes.includes(rerankExpressionType)) { @@ -44,17 +40,6 @@ export const validate = ( messages.push(errors.byId('inferenceIdRequired', command.location, { command: 'RERANK' })); } - const targetName = targetField?.name || 'rerank'; - - // Sets the target field so the column is recognized after the command is applied - context?.userDefinedColumns.set(targetName, [ - { - name: targetName, - location: targetField?.location || location, - type: 'keyword', - }, - ]); - messages.push(...validateCommandArguments(command, ast, context, callbacks)); return messages; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.test.ts new file mode 100644 index 0000000000000..2f7d253913d9c --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.test.ts @@ -0,0 +1,50 @@ +/* + * 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 { Parser } from '../../../..'; +import type { ESQLUserDefinedColumn } from '../../types'; + +import { columnsAfter } from './columns_after'; + +describe('ROW > columnsAfter', () => { + it('adds new user-defined columns', () => { + const queryString = `ROW baz = 1, foo = FLOOR(2. + 2.), TO_LONG(23 * 4)`; + + // Can't use synth because it steps on the location information + // which is used to determine the name of the new column + const { + root: { + commands: [command], + }, + } = Parser.parseQuery(queryString); + + const result = columnsAfter(command, [], queryString); + + expect(result).toEqual([ + { + name: 'baz', + type: 'integer', + location: { min: 4, max: 6 }, + userDefined: true, + }, + { + name: 'foo', + type: 'double', + location: { min: 13, max: 15 }, + userDefined: true, + }, + { + name: 'TO_LONG(23 * 4)', + type: 'long', + location: { min: 35, max: 49 }, + userDefined: true, + }, + ]); + }); +}); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.ts new file mode 100644 index 0000000000000..2ff89393eba11 --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/columns_after.ts @@ -0,0 +1,46 @@ +/* + * 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 { isAssignment, isColumn } from '../../../ast/is'; +import { getExpressionType } from '../../../definitions/utils'; +import type { ESQLAstItem, ESQLCommand } from '../../../types'; +import type { ESQLColumnData, ESQLUserDefinedColumn } from '../../types'; + +export const columnsAfter = ( + command: ESQLCommand, + _previousColumns: ESQLColumnData[], // will always be empty for ROW + query: string +) => { + const typeOf = (thing: ESQLAstItem) => getExpressionType(thing, new Map()); + + const columns = []; + + for (const expression of command.args) { + if (isAssignment(expression) && isColumn(expression.args[0])) { + const name = expression.args[0].parts.join('.'); + const newColumn: ESQLUserDefinedColumn = { + name, + type: typeOf(expression.args[1]), + location: expression.args[0].location, + userDefined: true, + }; + columns.push(newColumn); + } else if (!Array.isArray(expression)) { + const newColumn: ESQLUserDefinedColumn = { + name: query.substring(expression.location.min, expression.location.max + 1), + type: typeOf(expression), + location: expression.location, + userDefined: true, + }; + columns.push(newColumn); + } + } + + return columns; +}; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/index.ts index 986c5932dbd57..9a454af9e2ac9 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/index.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/row/index.ts @@ -11,10 +11,12 @@ import type { ICommandMethods } from '../../registry'; import { autocomplete } from './autocomplete'; import type { ICommandContext } from '../../types'; import { validate } from './validate'; +import { columnsAfter } from './columns_after'; const rowCommandMethods: ICommandMethods = { autocomplete, validate, + columnsAfter, }; export const rowCommand = { diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.test.ts new file mode 100644 index 0000000000000..779cf69419ccc --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.test.ts @@ -0,0 +1,18 @@ +/* + * 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 { SHOW_INFO_FIELDS, columnsAfter } from './columns_after'; + +describe('SHOW columnsAfter', () => { + it('returns info', async () => { + const result = await columnsAfter(); + + expect(result).toEqual(SHOW_INFO_FIELDS); + }); +}); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.ts new file mode 100644 index 0000000000000..22e129fbdca58 --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/columns_after.ts @@ -0,0 +1,32 @@ +/* + * 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 type { ESQLFieldWithMetadata } from '../../types'; + +// from https://github.com/elastic/elasticsearch/blob/da50c723a43262a9f46228ca36099647706c90dd/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/show/ShowInfo.java#L56-L58 +export const SHOW_INFO_FIELDS: ESQLFieldWithMetadata[] = [ + { + name: 'version', + type: 'keyword', + userDefined: false, + }, + { + name: 'date', + type: 'keyword', + userDefined: false, + }, + { + name: 'hash', + type: 'keyword', + userDefined: false, + }, +]; + +export const columnsAfter = () => { + return [...SHOW_INFO_FIELDS]; +}; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/index.ts index dc8e0b16fe35a..0f930f7fc9064 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/index.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/show/index.ts @@ -10,9 +10,11 @@ import { i18n } from '@kbn/i18n'; import type { ICommandMethods } from '../../registry'; import { autocomplete } from './autocomplete'; import type { ICommandContext } from '../../types'; +import { columnsAfter } from './columns_after'; const showCommandMethods: ICommandMethods = { autocomplete, + columnsAfter, }; export const showCommand = { diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/autocomplete.ts index 05a79c9c5fc3e..5d3c6a260f782 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/autocomplete.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/autocomplete.ts @@ -78,12 +78,7 @@ export async function autocomplete( const columnExists = (name: string) => _columnExists(name, context); - if ( - isExpressionComplete( - getExpressionType(expressionRoot, context?.fields, context?.userDefinedColumns), - innerText - ) - ) { + if (isExpressionComplete(getExpressionType(expressionRoot, context?.columns), innerText)) { suggestions.push( ...getSuggestionsAfterCompleteExpression(innerText, expressionRoot, columnExists) ); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/validate.test.ts index 60b0263d96ea3..d41ffc2e48593 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/validate.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/sort/validate.test.ts @@ -10,6 +10,7 @@ import { mockContext } from '../../../__tests__/context_fixtures'; import { validate } from './validate'; import { expectErrors } from '../../../__tests__/validation'; import { getNoValidCallSignatureError } from '../../../definitions/utils/validation/utils'; +import type { ICommandContext } from '../../types'; const sortExpectErrors = (query: string, expectedErrors: string[], context = mockContext) => { return expectErrors(query, expectedErrors, context, 'sort', validate); @@ -21,18 +22,17 @@ describe('SORT Validation', () => { }); test('validates the most basic query', () => { - const newUserDefinedColumns = new Map(mockContext.userDefinedColumns); + const newColumns = new Map(mockContext.columns); - newUserDefinedColumns.set('COUNT(*)', [ - { - name: 'COUNT(*)', - type: 'integer', - location: { min: 0, max: 10 }, - }, - ]); - const context = { + newColumns.set('COUNT(*)', { + name: 'COUNT(*)', + type: 'integer', + location: { min: 0, max: 10 }, + userDefined: true, + }); + const context: ICommandContext = { ...mockContext, - userDefinedColumns: newUserDefinedColumns, + columns: newColumns, }; sortExpectErrors('from a_index | sort "field" ', []); sortExpectErrors('from a_index | sort wrongField ', ['Unknown column "wrongField"']); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/autocomplete.ts index 697541fe6ef9b..a83ddcaadfeb2 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/autocomplete.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/autocomplete.ts @@ -176,11 +176,7 @@ export async function autocomplete( // Is this a complete boolean expression? // If so, we can call it done and suggest a pipe - const expressionType = getExpressionType( - expressionRoot, - context?.fields, - context?.userDefinedColumns - ); + const expressionType = getExpressionType(expressionRoot, context?.columns); if (expressionType === 'boolean' && isExpressionComplete(expressionType, innerText)) { suggestions.push(pipeCompleteItem, { ...commaCompleteItem, text: ', ' }, byCompleteItem); } @@ -306,12 +302,7 @@ async function getExpressionSuggestions({ suggestions.push(...emptySuggestions); } - if ( - isExpressionComplete( - getExpressionType(expressionRoot, context?.fields, context?.userDefinedColumns), - innerText - ) - ) { + if (isExpressionComplete(getExpressionType(expressionRoot, context?.columns), innerText)) { suggestions.push(...afterCompleteSuggestions); } diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts index f6d74a8e3ee9b..847dcca5f8dd0 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.test.ts @@ -6,152 +6,94 @@ * 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 { synth } from '../../../..'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; +import { Parser, synth } from '../../../..'; +import type { ESQLColumnData } from '../../types'; +import { type ESQLFieldWithMetadata } from '../../types'; import { columnsAfter } from './columns_after'; describe('STATS', () => { it('adds the user defined column, when no grouping is given', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - ] as ESQLFieldWithMetadata[]; - const context = { - userDefinedColumns: new Map([ - [ - 'var0', - [ - { - name: 'var0', - type: 'double', - location: { min: 0, max: 10 }, - }, - ], - ], - ]), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }; - - const result = columnsAfter(synth.cmd`STATS var0=AVG(field2)`, previousCommandFields, context); - - expect(result).toEqual([{ name: 'var0', type: 'double' }]); + const previousCommandFields: ESQLFieldWithMetadata[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; + + const result = columnsAfter(synth.cmd`STATS var0=AVG(field2)`, previousCommandFields, ''); + + expect(result).toEqual([ + { name: 'var0', type: 'double', userDefined: true, location: { min: 0, max: 0 } }, + ]); }); it('adds the escaped column, when no grouping is given', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - ] as ESQLFieldWithMetadata[]; - - const context = { - userDefinedColumns: new Map([ - [ - 'AVG(field2)', - [ - { - name: 'AVG(field2)', - type: 'double', - location: { min: 0, max: 10 }, - }, - ], - ], - ]), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }; - - const result = columnsAfter(synth.cmd`STATS AVG(field2)`, previousCommandFields, context); - - expect(result).toEqual([{ name: 'AVG(field2)', type: 'double' }]); + const previousCommandFields: ESQLFieldWithMetadata[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; + + const queryString = `FROM index | STATS AVG(field2)`; + + // Can't use synth because it steps on the location information + // which is used to determine the name of the new column + const { + root: { + commands: [, command], + }, + } = Parser.parseQuery(queryString); + + const result = columnsAfter(command, previousCommandFields, queryString); + + expect(result).toEqual([ + { name: 'AVG(field2)', type: 'double', userDefined: true, location: { min: 19, max: 29 } }, + ]); }); it('adds the escaped and grouping columns', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - ] as ESQLFieldWithMetadata[]; - - const context = { - userDefinedColumns: new Map([ - [ - 'AVG(field2)', - [ - { - name: 'AVG(field2)', - type: 'double', - location: { min: 0, max: 10 }, - }, - ], - ], - ]), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }; - - const result = columnsAfter( - synth.cmd`STATS AVG(field2) BY field1`, - previousCommandFields, - context - ); + const previousCommandFields: ESQLFieldWithMetadata[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + ]; + + const queryString = `FROM a | STATS AVG(field2) BY field1`; + + // Can't use synth because it steps on the location information + // which is used to determine the name of the new column + const { + root: { + commands: [, command], + }, + } = Parser.parseQuery(queryString); + + const result = columnsAfter(command, previousCommandFields, queryString); expect(result).toEqual([ - { name: 'field1', type: 'keyword' }, - { name: 'AVG(field2)', type: 'double' }, + { name: 'AVG(field2)', type: 'double', userDefined: true, location: { min: 15, max: 25 } }, + { name: 'field1', type: 'keyword', userDefined: true, location: { min: 30, max: 35 } }, ]); }); it('adds the user defined and grouping columns', () => { - const previousCommandFields = [ - { name: 'field1', type: 'keyword' }, - { name: 'field2', type: 'double' }, - { name: '@timestamp', type: 'date' }, - ] as ESQLFieldWithMetadata[]; - - const context = { - userDefinedColumns: new Map([ - [ - 'AVG(field2)', - [ - { - name: 'AVG(field2)', - type: 'double', - location: { min: 0, max: 10 }, - }, - ], - ], - [ - 'buckets', - [ - { - name: 'buckets', - type: 'unknown', - location: { min: 0, max: 10 }, - }, - ], - ], - ]), - fields: new Map([ - ['field1', { name: 'field1', type: 'keyword' }], - ['count', { name: 'count', type: 'double' }], - ]), - }; - - const result = columnsAfter( - synth.cmd`STATS AVG(field2) BY buckets=BUCKET(@timestamp,50,?_tstart,?_tend)`, - previousCommandFields, - context - ); + const previousCommandFields: ESQLFieldWithMetadata[] = [ + { name: 'field1', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, + { name: '@timestamp', type: 'date', userDefined: false }, + ]; - expect(result).toEqual([ - { name: 'AVG(field2)', type: 'double' }, - { name: 'buckets', type: 'unknown' }, + const queryString = `FROM a | STATS AVG(field2) BY buckets=BUCKET(@timestamp,50,?_tstart,?_tend)`; + + // Can't use synth because it steps on the location information + // which is used to determine the name of the new column + const { + root: { + commands: [, command], + }, + } = Parser.parseQuery(queryString); + + const result = columnsAfter(command, previousCommandFields, queryString); + + expect(result).toEqual([ + { name: 'AVG(field2)', type: 'double', userDefined: true, location: { min: 15, max: 25 } }, + { name: 'buckets', type: 'date', userDefined: true, location: { min: 30, max: 36 } }, ]); }); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts index 1855367c5ec69..902b97406aa3d 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/columns_after.ts @@ -7,56 +7,61 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import uniqBy from 'lodash/uniqBy'; -import type { ESQLCommand } from '../../../types'; -import { walk } from '../../../walker'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../types'; -import type { ICommandContext } from '../../types'; -import type { FieldType } from '../../../definitions/types'; -import { isColumn } from '../../../ast/is'; - -function transformMapToESQLFields( - inputMap: Map -): ESQLFieldWithMetadata[] { - const esqlFields: ESQLFieldWithMetadata[] = []; - - for (const [, userDefinedColumns] of inputMap) { - for (const userDefinedColumn of userDefinedColumns) { - // Only include userDefinedColumns that have a known type - if (userDefinedColumn.type) { - esqlFields.push({ - name: userDefinedColumn.name, - type: userDefinedColumn.type as FieldType, - }); - } +import { isAssignment, isColumn, isOptionNode } from '../../../ast/is'; +import type { SupportedDataType } from '../../../definitions/types'; +import { getExpressionType } from '../../../definitions/utils'; +import type { ESQLAstItem, ESQLCommand, ESQLCommandOption } from '../../../types'; +import type { ESQLColumnData, ESQLUserDefinedColumn } from '../../types'; + +const getUserDefinedColumns = ( + command: ESQLCommand | ESQLCommandOption, + typeOf: (thing: ESQLAstItem) => SupportedDataType | 'unknown', + query: string +): ESQLUserDefinedColumn[] => { + const columns: ESQLUserDefinedColumn[] = []; + + for (const expression of command.args) { + if (isAssignment(expression) && isColumn(expression.args[0])) { + const name = expression.args[0].parts.join('.'); + const newColumn: ESQLUserDefinedColumn = { + name, + type: typeOf(expression.args[1]), + location: expression.args[0].location, + userDefined: true, + }; + columns.push(newColumn); + continue; + } + + if (isOptionNode(expression) && expression.name === 'by') { + columns.push(...getUserDefinedColumns(expression, typeOf, query)); + continue; + } + + if (!isOptionNode(expression) && !Array.isArray(expression)) { + const newColumn: ESQLUserDefinedColumn = { + name: query.substring(expression.location.min, expression.location.max + 1), + type: typeOf(expression), + location: expression.location, + userDefined: true, + }; + columns.push(newColumn); + continue; } } - return esqlFields; -} + return columns; +}; export const columnsAfter = ( command: ESQLCommand, - previousColumns: ESQLFieldWithMetadata[], - context?: ICommandContext + previousColumns: ESQLColumnData[], + query: string ) => { - const columns: string[] = []; - const userDefinedColumns = - context?.userDefinedColumns ?? new Map(); - - walk(command, { - visitCommandOption: (node) => { - const args = node.args.filter(isColumn); - const breakdownColumns = args.map((arg) => arg.name); - columns.push(...breakdownColumns); - }, - }); - - const columnsToKeep = previousColumns.filter((field) => { - return columns.some((column) => column === field.name); - }); - - const arrayOfUserDefinedColumns: ESQLFieldWithMetadata[] = - transformMapToESQLFields(userDefinedColumns); - - return uniqBy([...columnsToKeep, ...arrayOfUserDefinedColumns], 'name'); + const columnMap = new Map(); + previousColumns.forEach((col) => columnMap.set(col.name, col)); // TODO make this more efficient + + const typeOf = (thing: ESQLAstItem) => getExpressionType(thing, columnMap); + + return uniqBy([...getUserDefinedColumns(command, typeOf, query)], 'name'); }; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/validate.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/validate.test.ts index 4ff1a9ce9cf1b..75f66b4776f9d 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/validate.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/stats/validate.test.ts @@ -9,6 +9,7 @@ import { mockContext } from '../../../__tests__/context_fixtures'; import { validate } from './validate'; import { expectErrors } from '../../../__tests__/validation'; +import type { ICommandContext } from '../../types'; const statsExpectErrors = (query: string, expectedErrors: string[], context = mockContext) => { return expectErrors(query, expectedErrors, context, 'stats', validate); @@ -20,24 +21,22 @@ describe('STATS Validation', () => { }); describe('STATS [ BY ]', () => { - const newUserDefinedColumns = new Map(mockContext.userDefinedColumns); - newUserDefinedColumns.set('doubleField * 3.281', [ - { - name: 'doubleField * 3.281', - type: 'double', - location: { min: 0, max: 10 }, - }, - ]); - newUserDefinedColumns.set('avg_doubleField', [ - { - name: 'avg_doubleField', - type: 'double', - location: { min: 0, max: 10 }, - }, - ]); - const context = { + const newColumns = new Map(mockContext.columns); + newColumns.set('doubleField * 3.281', { + name: 'doubleField * 3.281', + type: 'double', + location: { min: 0, max: 10 }, + userDefined: true, + }); + newColumns.set('avg_doubleField', { + name: 'avg_doubleField', + type: 'double', + location: { min: 0, max: 10 }, + userDefined: true, + }); + const context: ICommandContext = { ...mockContext, - userDefinedColumns: newUserDefinedColumns, + columns: newColumns, }; test('no errors on correct usage', () => { statsExpectErrors('from a_index | stats by textField', []); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.test.ts new file mode 100644 index 0000000000000..4e166a8759076 --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.test.ts @@ -0,0 +1,28 @@ +/* + * 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 type { ESQLFieldWithMetadata } from '../../types'; +import { columnsAfter } from './columns_after'; + +describe('TS columnsAfter', () => { + it('returns fields from the source', async () => { + const sourceFields: ESQLFieldWithMetadata[] = [ + { name: 'field1', type: 'double', userDefined: false }, + { name: 'field2', type: 'keyword', userDefined: false }, + ]; + + const result = await columnsAfter({} as any, [], '', { + fromFrom: () => Promise.resolve(sourceFields), + fromJoin: () => Promise.resolve([]), + fromEnrich: () => Promise.resolve([]), + }); + + expect(result).toEqual(sourceFields); + }); +}); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.ts new file mode 100644 index 0000000000000..24a08d97529e1 --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/columns_after.ts @@ -0,0 +1,20 @@ +/* + * 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 { type ESQLCommand } from '../../../types'; +import type { ESQLColumnData } from '../../types'; +import type { IAdditionalFields } from '../../registry'; + +export const columnsAfter = ( + command: ESQLCommand, + _previousColumns: ESQLColumnData[], // will always be empty for TS + _query: string, + additionalFields: IAdditionalFields +) => { + return additionalFields.fromFrom(command); +}; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/index.ts index c8b106a8654f8..2270c6579dc4e 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/index.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/timeseries/index.ts @@ -10,11 +10,13 @@ import { i18n } from '@kbn/i18n'; import type { ICommandMethods } from '../../registry'; import { autocomplete } from './autocomplete'; import { validate } from './validate'; +import { columnsAfter } from './columns_after'; import type { ICommandContext } from '../../types'; const timeseriesCommandMethods: ICommandMethods = { autocomplete, validate, + columnsAfter, }; export const timeseriesCommand = { diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/where/autocomplete.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/where/autocomplete.ts index 0773410264679..f737337ea696e 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/where/autocomplete.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/commands/where/autocomplete.ts @@ -49,11 +49,7 @@ export async function autocomplete( // Is this a complete boolean expression? // If so, we can call it done and suggest a pipe - const expressionType = getExpressionType( - expressionRoot, - context?.fields, - context?.userDefinedColumns - ); + const expressionType = getExpressionType(expressionRoot, context?.columns); if (expressionType === 'boolean' && isExpressionComplete(expressionType, innerText)) { suggestions.push(pipeCompleteItem); } diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts index 6bb5037901abd..591bb3b484ff1 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/registry.ts @@ -8,7 +8,12 @@ */ import type { LicenseType } from '@kbn/licensing-types'; import type { ESQLMessage, ESQLCommand, ESQLAst } from '../types'; -import type { ISuggestionItem, ESQLFieldWithMetadata, ICommandCallbacks } from './types'; +import type { + ISuggestionItem, + ICommandCallbacks, + ESQLColumnData, + ESQLFieldWithMetadata, +} from './types'; /** * Interface defining the methods that each ES|QL command should register. @@ -59,9 +64,10 @@ export interface ICommandMethods { */ columnsAfter?: ( command: ESQLCommand, - previousColumns: ESQLFieldWithMetadata[], - context?: TContext - ) => ESQLFieldWithMetadata[]; + previousColumns: ESQLColumnData[], + query: string, + newFields: IAdditionalFields + ) => Promise | ESQLColumnData[]; } export interface ICommandMetadata { @@ -114,6 +120,12 @@ export interface ICommandRegistry { getCommandByName(commandName: string): ICommand | undefined; } +export interface IAdditionalFields { + fromJoin: (cmd: ESQLCommand) => Promise; + fromEnrich: (cmd: ESQLCommand) => Promise; + fromFrom: (cmd: ESQLCommand) => Promise; +} + /** * Implementation of the ESQL Command Registry. * This class manages the registration, storage, and retrieval of ESQL command methods. diff --git a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts index 0cfe9cf22207f..c10be856e91e1 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/commands_registry/types.ts @@ -97,9 +97,11 @@ export type GetColumnsByTypeFn = ( } ) => Promise; +// TODO consider not exporting this export interface ESQLFieldWithMetadata { name: string; type: FieldType; + userDefined: false; isEcs?: boolean; hasConflict?: boolean; metadata?: { @@ -107,15 +109,19 @@ export interface ESQLFieldWithMetadata { }; } +// TODO consider not exporting this export interface ESQLUserDefinedColumn { name: string; // invalid expressions produce columns of type "unknown" // also, there are some cases where we can't yet infer the type of // a valid expression as with `CASE` which can return union types type: SupportedDataType | 'unknown'; - location: ESQLLocation; + userDefined: true; + location: ESQLLocation; // TODO should this be optional? } +export type ESQLColumnData = ESQLUserDefinedColumn | ESQLFieldWithMetadata; + export interface ESQLPolicy { name: string; sourceIndices: string[]; @@ -126,13 +132,12 @@ export interface ESQLPolicy { export interface ICommandCallbacks { getByType?: GetColumnsByTypeFn; getSuggestedUserDefinedColumnName?: (extraFieldNames?: string[] | undefined) => string; - getColumnsForQuery?: (query: string) => Promise; + getColumnsForQuery?: (query: string) => Promise; hasMinimumLicenseRequired?: (minimumLicenseRequired: LicenseType) => boolean; } export interface ICommandContext { - userDefinedColumns: Map; - fields: Map; + columns: Map; sources?: ESQLSourceResult[]; joinSources?: IndexAutocompleteItem[]; timeSeriesSources?: IndexAutocompleteItem[]; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts index 49175c3f8d2e3..6f6c1972c0204 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/columns.ts @@ -6,202 +6,21 @@ * 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 { Visitor } from '../../../visitor'; -import type { - ESQLFieldWithMetadata, - ESQLUserDefinedColumn, - ICommandContext, -} from '../../../commands_registry/types'; -import type { - ESQLFunction, - ESQLAst, - ESQLAstItem, - ESQLCommand, - ESQLColumn, - ESQLIdentifier, -} from '../../../types'; -import { isColumn, isFunctionExpression } from '../../../ast/is'; -import { getExpressionType } from '../expressions'; -import { EDITOR_MARKER } from '../../constants'; +import type { ICommandContext } from '../../../commands_registry/types'; +import type { ESQLColumn, ESQLIdentifier } from '../../../types'; import { fuzzySearch } from '../shared'; -function addToUserDefinedColumnOccurrences( - userDefinedColumns: Map, - instance: ESQLUserDefinedColumn -) { - if (!userDefinedColumns.has(instance.name)) { - userDefinedColumns.set(instance.name, []); - } - const userDefinedColumnsOccurrencies = userDefinedColumns.get(instance.name)!; - userDefinedColumnsOccurrencies.push(instance); -} - -export function excludeUserDefinedColumnsFromCurrentCommand( - commands: ESQLCommand[], - currentCommand: ESQLCommand, - fieldsMap: Map, - queryString: string -) { - const anyUserDefinedColumns = collectUserDefinedColumns(commands, fieldsMap, queryString); - const currentCommandUserDefinedColumns = collectUserDefinedColumns( - [currentCommand], - fieldsMap, - queryString - ); - const resultUserDefinedColumns = new Map(); - anyUserDefinedColumns.forEach((value, key) => { - if (!currentCommandUserDefinedColumns.has(key)) { - resultUserDefinedColumns.set(key, value); - } - }); - return resultUserDefinedColumns; -} - -function addUserDefinedColumnFromAssignment( - assignOperation: ESQLFunction, - userDefinedColumns: Map, - fields: Map -) { - if (isColumn(assignOperation.args[0])) { - const rightHandSideArgType = getExpressionType( - assignOperation.args[1], - fields, - userDefinedColumns - ); - addToUserDefinedColumnOccurrences(userDefinedColumns, { - name: assignOperation.args[0].parts.join('.'), - type: rightHandSideArgType /* fallback to number */, - location: assignOperation.args[0].location, - }); - } -} - -function addUserDefinedColumnFromExpression( - expressionOperation: ESQLFunction, - queryString: string, - userDefinedColumns: Map, - fields: Map -) { - if (!expressionOperation.text.includes(EDITOR_MARKER)) { - const expressionText = queryString.substring( - expressionOperation.location.min, - expressionOperation.location.max + 1 - ); - const expressionType = getExpressionType(expressionOperation, fields, userDefinedColumns); - addToUserDefinedColumnOccurrences(userDefinedColumns, { - name: expressionText, - type: expressionType, - location: expressionOperation.location, - }); - } -} - -function addToUserDefinedColumns( - oldArg: ESQLAstItem, - newArg: ESQLAstItem, - fields: Map, - userDefinedColumns: Map -) { - if (isColumn(oldArg) && isColumn(newArg)) { - const newUserDefinedColumn: ESQLUserDefinedColumn = { - name: newArg.parts.join('.'), - type: 'double' /* fallback to number */, - location: newArg.location, - }; - // Now workout the exact type - // it can be a rename of another userDefinedColumn as well - const oldRef = - fields.get(oldArg.parts.join('.')) || userDefinedColumns.get(oldArg.parts.join('.')); - if (oldRef) { - addToUserDefinedColumnOccurrences(userDefinedColumns, newUserDefinedColumn); - newUserDefinedColumn.type = Array.isArray(oldRef) ? oldRef[0].type : oldRef.type; - } - } -} - -export function collectUserDefinedColumns( - ast: ESQLAst, - fields: Map, - queryString: string -): Map { - const userDefinedColumns = new Map(); - - const visitor = new Visitor() - .on('visitLiteralExpression', (ctx) => { - // TODO - add these as userDefinedColumns - }) - .on('visitExpression', (_ctx) => {}) // required for the types :shrug: - .on('visitFunctionCallExpression', (ctx) => { - const node = ctx.node; - - if (node.subtype === 'binary-expression' && node.name === 'where') { - ctx.visitArgument(0, undefined); - return; - } - - if (node.name === 'as') { - const [oldArg, newArg] = ctx.node.args; - addToUserDefinedColumns(oldArg, newArg, fields, userDefinedColumns); - } else if (node.name === '=') { - addUserDefinedColumnFromAssignment(node, userDefinedColumns, fields); - } else { - addUserDefinedColumnFromExpression(node, queryString, userDefinedColumns, fields); - } - }) - .on('visitCommandOption', (ctx) => { - if (ctx.node.name === 'by') { - return [...ctx.visitArguments()]; - } else if (ctx.node.name === 'with') { - for (const assignFn of ctx.node.args) { - if (isFunctionExpression(assignFn)) { - const [newArg, oldArg] = assignFn?.args || []; - // TODO why is oldArg an array? - if (Array.isArray(oldArg)) { - addToUserDefinedColumns(oldArg[0], newArg, fields, userDefinedColumns); - } - } - } - } - }) - .on('visitCommand', (ctx) => { - const ret = []; - if (['row', 'eval', 'stats', 'inlinestats', 'ts', 'rename'].includes(ctx.node.name)) { - ret.push(...ctx.visitArgs()); - } - if (['stats', 'inlinestats', 'enrich'].includes(ctx.node.name)) { - // BY and WITH can contain userDefinedColumns - ret.push(...ctx.visitOptions()); - } - if (ctx.node.name === 'fork') { - ret.push(...ctx.visitSubQueries()); - } - return ret; - }) - .on('visitQuery', (ctx) => [...ctx.visitCommands()]); - - visitor.visitQuery(ast); - - return userDefinedColumns; -} - /** * TODO - consider calling lookupColumn under the hood of this function. Seems like they should really do the same thing. */ -export function getColumnExists( - node: ESQLColumn | ESQLIdentifier, - { fields, userDefinedColumns }: ICommandContext -) { +export function getColumnExists(node: ESQLColumn | ESQLIdentifier, { columns }: ICommandContext) { const columnName = node.type === 'identifier' ? node.name : node.parts.join('.'); - if (fields.has(columnName) || userDefinedColumns.has(columnName)) { + if (columns.has(columnName)) { return true; } // TODO — I don't see this fuzzy searching in lookupColumn... should it be there? - if ( - Boolean( - fuzzySearch(columnName, fields.keys()) || fuzzySearch(columnName, userDefinedColumns.keys()) - ) - ) { + if (Boolean(fuzzySearch(columnName, columns.keys()))) { return true; } diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts index 3b5da9c6397bf..e82882eb0cb1d 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/functions.ts @@ -24,7 +24,7 @@ import { listCompleteItem, } from '../../../commands_registry/complete_items'; import type { - ESQLFieldWithMetadata, + ESQLColumnData, GetColumnsByTypeFn, ICommandCallbacks, ICommandContext, @@ -51,7 +51,6 @@ import { import { getCompatibleLiterals, getDateLiterals } from '../literals'; import { getSuggestionsToRightOfOperatorExpression } from '../operators'; import { buildValueDefinitions } from '../values'; -import { collectUserDefinedColumns, excludeUserDefinedColumnsFromCurrentCommand } from './columns'; import { extractTypeFromASTArg, getFieldsOrFunctionsSuggestions, @@ -167,19 +166,11 @@ export async function getFunctionArgsSuggestions( signatures: filterFunctionSignatures(fnDefinition.signatures, hasMinimumLicenseRequired), }; - const fieldsMap: Map = context?.fields || new Map(); - const anyUserDefinedColumns = collectUserDefinedColumns(commands, fieldsMap, innerText); + const columnMap: Map = context?.columns || new Map(); const references = { - fields: fieldsMap, - userDefinedColumns: anyUserDefinedColumns, + columns: columnMap, }; - const userDefinedColumnsExcludingCurrentCommandOnes = excludeUserDefinedColumnsFromCurrentCommand( - commands, - command, - fieldsMap, - innerText - ); const { typesToSuggestNext, hasMoreMandatoryArgs, enrichedArgs, argIndex } = getValidSignaturesAndTypesToSuggestNext( @@ -232,8 +223,7 @@ export async function getFunctionArgsSuggestions( arg && isColumn(arg) && !getColumnExists(arg, { - fields: fieldsMap, - userDefinedColumns: userDefinedColumnsExcludingCurrentCommandOnes, + columns: columnMap, }); if (noArgDefined || isUnknownColumn) { // ... | EVAL fn( ) @@ -437,7 +427,7 @@ async function getListArgsSuggestions( innerText: string, commands: ESQLCommand[], getFieldsByType: GetColumnsByTypeFn, - fieldsMap: Map, + columnMap: Map, offset: number, hasMinimumLicenseRequired?: (minimumLicenseRequired: LicenseType) => boolean, activeProduct?: PricingProduct @@ -460,18 +450,10 @@ async function getListArgsSuggestions( } } - const anyUserDefinedColumns = collectUserDefinedColumns(commands, fieldsMap, innerText); - // extract the current node from the userDefinedColumns inferred - anyUserDefinedColumns.forEach((values, key) => { - if (values.some((v) => v.location === node.location)) { - anyUserDefinedColumns.delete(key); - } - }); const [firstArg] = node.args; if (isColumn(firstArg)) { const argType = extractTypeFromASTArg(firstArg, { - fields: fieldsMap, - userDefinedColumns: anyUserDefinedColumns, + columns: columnMap, }); if (argType) { // do not propose existing columns again @@ -485,8 +467,7 @@ async function getListArgsSuggestions( getFieldsByType, { functions: true, - fields: true, - userDefinedColumns: anyUserDefinedColumns, + columns: true, }, { ignoreColumns: [firstArg.name, ...otherArgs.map(({ name }) => name)] }, hasMinimumLicenseRequired, @@ -536,8 +517,7 @@ export const getInsideFunctionsSuggestions = async ( queryText: innerText, location: getLocationFromCommandOrOptionName(command.name), rootOperator: node, - getExpressionType: (expression) => - getExpressionType(expression, context?.fields, context?.userDefinedColumns), + getExpressionType: (expression) => getExpressionType(expression, context?.columns), getColumnsByType: callbacks?.getByType ?? (() => Promise.resolve([])), hasMinimumLicenseRequired: callbacks?.hasMinimumLicenseRequired, activeProduct: context?.activeProduct, @@ -550,7 +530,7 @@ export const getInsideFunctionsSuggestions = async ( innerText, ast, callbacks?.getByType ?? (() => Promise.resolve([])), - context?.fields ?? new Map(), + context?.columns ?? new Map(), cursorPosition ?? 0, callbacks?.hasMinimumLicenseRequired, context?.activeProduct diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts index e4e71dc601b66..741af5ac94ac7 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/autocomplete/helpers.ts @@ -22,7 +22,6 @@ import type { import type { ISuggestionItem, GetColumnsByTypeFn, - ESQLUserDefinedColumn, ICommandContext, } from '../../../commands_registry/types'; import { Location } from '../../../commands_registry/types'; @@ -169,14 +168,12 @@ export async function getFieldsOrFunctionsSuggestions( getFieldsByType: GetColumnsByTypeFn, { functions, - fields, - userDefinedColumns, + columns: fields, values = false, literals = false, }: { functions: boolean; - fields: boolean; - userDefinedColumns?: Map; + columns: boolean; literals?: boolean; values?: boolean; }, @@ -201,34 +198,6 @@ export async function getFieldsOrFunctionsSuggestions( functions ); - const filteredColumnByType: string[] = []; - if (userDefinedColumns) { - for (const userDefinedColumn of userDefinedColumns.values()) { - if ( - (types.includes('any') || types.includes(userDefinedColumn[0].type)) && - !ignoreColumns.includes(userDefinedColumn[0].name) - ) { - filteredColumnByType.push(userDefinedColumn[0].name); - } - } - // due to a bug on the ES|QL table side, filter out fields list with underscored userDefinedColumns names (??) - // avg( numberField ) => avg_numberField_ - const ALPHANUMERIC_REGEXP = /[^a-zA-Z\d]/g; - if ( - filteredColumnByType.length && - filteredColumnByType.some((v) => ALPHANUMERIC_REGEXP.test(v)) - ) { - for (const userDefinedColumn of filteredColumnByType) { - const underscoredName = userDefinedColumn.replace(ALPHANUMERIC_REGEXP, '_'); - const index = filteredFieldsByType.findIndex( - ({ label }) => underscoredName === label || `_${underscoredName}_` === label - ); - if (index >= 0) { - filteredFieldsByType.splice(index); - } - } - } - } // could also be in stats (bucket) but our autocomplete is not great yet const displayDateSuggestions = types.includes('date') && [Location.WHERE, Location.EVAL].includes(location); @@ -246,9 +215,6 @@ export async function getFieldsOrFunctionsSuggestions( activeProduct ) : [], - userDefinedColumns - ? pushItUpInTheList(buildUserDefinedColumnsDefinitions(filteredColumnByType), functions) - : [], literals ? getCompatibleLiterals(types) : [] ); @@ -380,11 +346,7 @@ export async function suggestForExpression({ case 'after_literal': case 'after_column': case 'after_function': - const expressionType = getExpressionType( - expressionRoot, - context?.fields, - context?.userDefinedColumns - ); + const expressionType = getExpressionType(expressionRoot, context?.columns); if (!isParameterType(expressionType)) { break; @@ -472,8 +434,7 @@ export async function suggestForExpression({ location, rootOperator: rightmostOperator, preferredExpressionType, - getExpressionType: (expression) => - getExpressionType(expression, context?.fields, context?.userDefinedColumns), + getExpressionType: (expression) => getExpressionType(expression, context?.columns), getColumnsByType, hasMinimumLicenseRequired, activeProduct, diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/columns.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/columns.ts index c6661f38f173d..5e59f24313cd3 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/columns.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/columns.ts @@ -16,12 +16,18 @@ import { fuzzySearch } from './shared'; */ export function getColumnExists( node: ESQLColumn | ESQLIdentifier, - { fields, userDefinedColumns }: Pick, + { columns }: Pick, excludeFields = false ) { const columnName = node.type === 'identifier' ? node.name : node.parts.join('.'); - const set = new Set([...(!excludeFields ? fields.keys() : []), ...userDefinedColumns.keys()]); + const set = new Set( + !excludeFields + ? columns.keys() + : Array.from(columns.values()) + .filter((col) => col.userDefined) + .map((col) => col.name) + ); if (set.has(columnName)) { return true; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.test.ts index 31e68a6f444d2..777783f941db6 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.test.ts @@ -9,6 +9,7 @@ import { Parser } from '../../parser'; import type { SupportedDataType } from '../types'; import { FunctionDefinitionTypes } from '../types'; +import type { ESQLColumnData } from '../../commands_registry/types'; import { Location } from '../../commands_registry/types'; import { buildPartialMatcher, getExpressionType } from './expressions'; import { setTestFunctions } from './test_functions'; @@ -127,27 +128,25 @@ describe('getExpressionType', () => { { name: 'fieldName', type: 'geo_shape', + userDefined: false, }, ], - ]), - new Map() + ]) ) ).toBe('geo_shape'); expect( getExpressionType( getASTForExpression('col0'), - new Map(), new Map([ [ 'col0', - [ - { - name: 'col0', - type: 'long', - location: { min: 0, max: 0 }, - }, - ], + { + name: 'col0', + type: 'long', + location: { min: 0, max: 0 }, + userDefined: true, + }, ], ]) ) @@ -155,9 +154,7 @@ describe('getExpressionType', () => { }); it('handles fields and userDefinedColumns which do not exist', () => { - expect(getExpressionType(getASTForExpression('fieldName'), new Map(), new Map())).toBe( - 'unknown' - ); + expect(getExpressionType(getASTForExpression('fieldName'), new Map())).toBe('unknown'); }); // this is here to fix https://github.com/elastic/kibana/issues/215157 @@ -172,16 +169,16 @@ describe('getExpressionType', () => { name: 'fieldName', type: 'unsupported', hasConflict: true, + userDefined: false, }, ], - ]), - new Map() + ]) ) ).toBe('unknown'); }); it('handles fields defined by a named param', () => { - expect(getExpressionType(getASTForExpression('??field'), new Map(), new Map())).toBe('param'); + expect(getExpressionType(getASTForExpression('??field'), new Map())).toBe('param'); }); }); @@ -279,17 +276,23 @@ describe('getExpressionType', () => { expect( getExpressionType( getASTForExpression('CASE(true, "", true, "", keywordField)'), - new Map([[`keywordField`, { name: 'keywordField', type: 'keyword' }]]), - new Map() + new Map([[`keywordField`, { name: 'keywordField', type: 'keyword', userDefined: false }]]) ) ).toBe('keyword'); expect( getExpressionType( getASTForExpression('CASE(true, "", true, "", keywordVar)'), - new Map(), - new Map([ - [`keywordVar`, [{ name: 'keywordVar', type: 'keyword', location: { min: 0, max: 0 } }]], + new Map([ + [ + `keywordVar`, + { + name: 'keywordVar', + type: 'keyword', + location: { min: 0, max: 0 }, + userDefined: true, + }, + ], ]) ) ).toBe('keyword'); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.ts index fd39885e28e3f..aef58a357413e 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/expressions.ts @@ -14,7 +14,7 @@ import { isFunctionExpression, } from '../../ast/is'; import type { ESQLAstItem, ESQLFunction } from '../../types'; -import type { ESQLFieldWithMetadata, ESQLUserDefinedColumn } from '../../commands_registry/types'; +import type { ESQLColumnData } from '../../commands_registry/types'; import type { SupportedDataType, FunctionDefinition } from '../types'; import { lastItem } from '../../visitor/utils'; import { getFunctionDefinition } from './functions'; @@ -127,8 +127,7 @@ export function isExpressionComplete( */ export function getExpressionType( root: ESQLAstItem | undefined, - fields?: Map, - userDefinedColumns?: Map + columns?: Map ): SupportedDataType | 'unknown' { if (!root) { return 'unknown'; @@ -138,7 +137,7 @@ export function getExpressionType( if (root.length === 0) { return 'unknown'; } - return getExpressionType(root[0], fields, userDefinedColumns); + return getExpressionType(root[0], columns); } if (isLiteral(root)) { @@ -163,8 +162,8 @@ export function getExpressionType( } } - if (isColumn(root) && fields && userDefinedColumns) { - const column = getColumnForASTNode(root, { fields, userDefinedColumns }); + if (isColumn(root) && columns) { + const column = getColumnForASTNode(root, { columns }); const lastArg = lastItem(root.args); // If the last argument is a param, we return its type (param literal type) // This is useful for cases like `where ??field` @@ -181,7 +180,7 @@ export function getExpressionType( } if (root.type === 'list') { - return getExpressionType(root.values[0], fields, userDefinedColumns); + return getExpressionType(root.values[0], columns); } if (isFunctionExpression(root)) { @@ -215,7 +214,7 @@ export function getExpressionType( * will be null, which we aren't detecting. But this is ok because we consider * userDefinedColumns and fields to be nullable anyways and account for that during validation. */ - return getExpressionType(root.args[root.args.length - 1], fields, userDefinedColumns); + return getExpressionType(root.args[root.args.length - 1], columns); } const signaturesWithCorrectArity = getSignaturesWithMatchingArity(fnDefinition, root); @@ -223,7 +222,7 @@ export function getExpressionType( if (!signaturesWithCorrectArity.length) { return 'unknown'; } - const argTypes = root.args.map((arg) => getExpressionType(arg, fields, userDefinedColumns)); + const argTypes = root.args.map((arg) => getExpressionType(arg, columns)); // When functions are passed null for any argument, they generally return null // This is a special case that is not reflected in our function definitions diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/functions.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/functions.ts index e2054c7ff8351..986edff787959 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/functions.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/functions.ts @@ -24,7 +24,7 @@ import { aggFunctionDefinitions } from '../generated/aggregation_functions'; import { timeSeriesAggFunctionDefinitions } from '../generated/time_series_agg_functions'; import { groupingFunctionDefinitions } from '../generated/grouping_functions'; import { scalarFunctionDefinitions } from '../generated/scalar_functions'; -import type { ESQLFieldWithMetadata, ISuggestionItem } from '../../commands_registry/types'; +import type { ESQLColumnData, ISuggestionItem } from '../../commands_registry/types'; import { TRIGGER_SUGGESTION_COMMAND } from '../../commands_registry/constants'; import { buildFunctionDocumentation } from './documentation'; import { getSafeInsertText, getControlSuggestion } from './autocomplete/helpers'; @@ -374,8 +374,8 @@ const getVariablePrefix = (variableType: ESQLVariableType) => ? '??' : '?'; -export const buildFieldsDefinitionsWithMetadata = ( - fields: ESQLFieldWithMetadata[], +export const buildColumnSuggestions = ( + columns: ESQLColumnData[], recommendedFieldsFromExtensions: RecommendedField[] = [], options?: { advanceCursor?: boolean; @@ -386,20 +386,23 @@ export const buildFieldsDefinitionsWithMetadata = ( }, variables?: ESQLControlVariable[] ): ISuggestionItem[] => { - const fieldsSuggestions = fields.map((field) => { - const fieldType = field.type.charAt(0).toUpperCase() + field.type.slice(1); - const titleCaseType = `${field.name} (${fieldType})`; + const fieldsSuggestions = columns.map((column) => { + const fieldType = column.type.charAt(0).toUpperCase() + column.type.slice(1); + const titleCaseType = `${column.name} (${fieldType})`; // Check if the field is in the recommended fields from extensions list // and if so, mark it as recommended. This also ensures that recommended fields // that are registered wrongly, won't be shown as suggestions. const fieldIsRecommended = recommendedFieldsFromExtensions.some( - (recommendedField) => recommendedField.name === field.name + (recommendedField) => recommendedField.name === column.name + ); + const sortText = getFieldsSortText( + !column.userDefined && Boolean(column.isEcs), + Boolean(fieldIsRecommended) ); - const sortText = getFieldsSortText(Boolean(field.isEcs), Boolean(fieldIsRecommended)); return { - label: field.name, + label: column.name, text: - getSafeInsertText(field.name) + + getSafeInsertText(column.name) + (options?.addComma ? ',' : '') + (options?.advanceCursor ? ' ' : ''), kind: 'Variable', @@ -415,7 +418,7 @@ export const buildFieldsDefinitionsWithMetadata = ( const userDefinedColumns = variables?.filter((variable) => variable.type === variableType) ?? []; - const controlSuggestions = fields.length + const controlSuggestions = columns.length ? getControlSuggestion( variableType, userDefinedColumns?.map((v) => `${getVariablePrefix(variableType)}${v.key}`) diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/index.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/index.ts index d0e53570e115b..3eaa965567fcb 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/index.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/index.ts @@ -16,7 +16,7 @@ export { } from './autocomplete/helpers'; export { getSuggestionsToRightOfOperatorExpression } from './operators'; export { - buildFieldsDefinitionsWithMetadata, + buildColumnSuggestions as buildFieldsDefinitionsWithMetadata, getFunctionSuggestions, getFunctionSignatures, getFunctionDefinition, diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/operators.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/operators.ts index c857cf1bc6972..fededccf342f1 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/operators.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/operators.ts @@ -192,7 +192,7 @@ export async function getSuggestionsToRightOfOperatorExpression({ getColumnsByType, { functions: true, - fields: true, + columns: true, values: Boolean(operator.subtype === 'binary-expression'), }, {}, diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/shared.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/shared.ts index fff968191b220..e046d52e7bd5c 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/shared.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/shared.ts @@ -126,10 +126,10 @@ export function unescapeColumnName(columnName: string) { */ export function getColumnByName( columnName: string, - { fields, userDefinedColumns }: ICommandContext + { columns }: ICommandContext ): ESQLFieldWithMetadata | ESQLUserDefinedColumn | undefined { const unescaped = unescapeColumnName(columnName); - return fields.get(unescaped) || userDefinedColumns.get(unescaped)?.[0]; + return columns.get(unescaped); } /** @@ -137,10 +137,10 @@ export function getColumnByName( */ export function getColumnForASTNode( node: ESQLColumn | ESQLIdentifier, - { fields, userDefinedColumns }: ICommandContext + { columns }: ICommandContext ): ESQLFieldWithMetadata | ESQLUserDefinedColumn | undefined { const formatted = node.type === 'identifier' ? node.name : node.parts.join('.'); - return getColumnByName(formatted, { fields, userDefinedColumns }); + return getColumnByName(formatted, { columns }); } function hasWildcard(name: string) { diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/commands.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/commands.ts index 514c98158d5c2..4c104cba12bed 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/commands.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/commands.ts @@ -18,8 +18,7 @@ export const validateCommandArguments = ( command: ESQLCommand, ast: ESQLAst, context: ICommandContext = { - userDefinedColumns: new Map(), // Ensure context is always defined - fields: new Map(), + columns: new Map(), // Ensure context is always defined }, callbacks: ICommandCallbacks = {} ) => { diff --git a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts index 0904f03358853..787edbe01228f 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/definitions/utils/validation/function.ts @@ -7,7 +7,6 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import type { LicenseType } from '@kbn/licensing-types'; -import { uniqBy } from 'lodash'; import { errors, getFunctionDefinition } from '..'; import { FunctionDefinitionTypes, within } from '../../../..'; import { @@ -35,8 +34,8 @@ import { getSignaturesWithMatchingArity, matchesArity, } from '../expressions'; -import { ColumnValidator } from './column'; import { isArrayType } from '../operators'; +import { ColumnValidator } from './column'; export function validateFunction({ fn, @@ -74,9 +73,7 @@ class FunctionValidator { ) { this.definition = getFunctionDefinition(fn.name); for (const arg of this.fn.args) { - this.argTypes.push( - getExpressionType(arg, this.context.fields, this.context.userDefinedColumns) - ); + this.argTypes.push(getExpressionType(arg, this.context.columns)); this.argLiteralsMask.push(isLiteral(arg)); } } @@ -155,20 +152,24 @@ class FunctionValidator { } // Validate column arguments - const columnMessages = this.fn.args.flat().flatMap((arg) => { - if (isColumn(arg) || isIdentifier(arg)) { - return new ColumnValidator(arg, this.context, this.parentCommand.name).validate(); + const columnsToValidate = []; + const flatArgs = this.fn.args.flat(); + for (let i = 0; i < flatArgs.length; i++) { + const arg = flatArgs[i]; + if ( + (isColumn(arg) || isIdentifier(arg)) && + !(this.definition.name === '=' && i === 0) && // don't validate left-hand side of assignment + !(this.definition.name === 'as' && i === 1) // don't validate right-hand side of AS + ) { + columnsToValidate.push(arg); } - return []; + } + + const columnMessages = columnsToValidate.flatMap((arg) => { + return new ColumnValidator(arg, this.context, this.parentCommand.name).validate(); }); - // uniqBy is used to cover a special case in ENRICH where an implicit assignment is possible - // so the AST actually stores an explicit "columnX = columnX" which duplicates the message - // - // @TODO - we will no longer need to store an assignment in the AST approach when we - // align field availability detection with the system used by autocomplete - // (start using columnsAfter instead of collectUserDefinedColumns) - this.report(...uniqBy(columnMessages, ({ location }) => `${location.min}-${location.max}`)); + this.report(...columnMessages); } /** diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/index.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/index.ts index b3507dca1e0ff..973717d840d86 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/index.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/index.ts @@ -21,10 +21,8 @@ export { suggest } from './src/autocomplete/autocomplete'; * Some utility functions that can be useful to build more feature * for the ES|QL language */ -export { collectUserDefinedColumns } from './src/shared/user_defined_columns'; - export { - getFieldsByTypeHelper, + getColumnsByTypeHelper as getFieldsByTypeHelper, getPolicyHelper, getSourcesHelper, } from './src/shared/resources_helpers'; diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/__tests__/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/__tests__/helpers.ts index a18dfbcb50824..0202eb6489f55 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/__tests__/helpers.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/__tests__/helpers.ts @@ -17,23 +17,28 @@ import type { ESQLCallbacks } from '../shared/types'; export const metadataFields: ESQLFieldWithMetadata[] = METADATA_FIELDS.map((field) => ({ name: field, type: 'keyword', + userDefined: false as false, })); export const fields: ESQLFieldWithMetadata[] = [ - ...fieldTypes.map((type) => ({ name: `${camelCase(type)}Field`, type })), - { name: 'any#Char$Field', type: 'double' }, - { name: 'kubernetes.something.something', type: 'double' }, - { name: '@timestamp', type: 'date' }, + ...fieldTypes.map((type) => ({ + name: `${camelCase(type)}Field`, + type, + userDefined: false as false, + })), + { name: 'any#Char$Field', type: 'double', userDefined: false }, + { name: 'kubernetes.something.something', type: 'double', userDefined: false }, + { name: '@timestamp', type: 'date', userDefined: false }, ]; export const enrichFields: ESQLFieldWithMetadata[] = [ - { name: 'otherField', type: 'text' }, - { name: 'yetAnotherField', type: 'double' }, + { name: 'otherField', type: 'text', userDefined: false }, + { name: 'yetAnotherField', type: 'double', userDefined: false }, ]; // eslint-disable-next-line @typescript-eslint/naming-convention export const unsupported_field: ESQLFieldWithMetadata[] = [ - { name: 'unsupported_field', type: 'unsupported' }, + { name: 'unsupported_field', type: 'unsupported', userDefined: false }, ]; export const indexes = [ @@ -140,6 +145,7 @@ export function getCallbackMocks(): ESQLCallbacks { name: 'keywordField', type: 'unsupported', hasConflict: true, + userDefined: false, }; return [field]; } diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.variables.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.variables.test.ts index 4fb4366355bb3..dc510dc939336 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.variables.test.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/autocomplete.command.variables.test.ts @@ -24,7 +24,8 @@ describe('autocomplete.suggest', () => { type: ESQLVariableType.VALUES, }, ], - getColumnsFor: () => Promise.resolve([{ name: 'agent.name', type: 'keyword' }]), + getColumnsFor: () => + Promise.resolve([{ name: 'agent.name', type: 'keyword', userDefined: false }]), }, }); @@ -54,7 +55,8 @@ describe('autocomplete.suggest', () => { callbacks: { canSuggestVariables: () => true, getVariables: () => [], - getColumnsFor: () => Promise.resolve([{ name: 'clientip', type: 'ip' }]), + getColumnsFor: () => + Promise.resolve([{ name: 'clientip', type: 'ip', userDefined: false }]), }, }); @@ -81,7 +83,8 @@ describe('autocomplete.suggest', () => { type: ESQLVariableType.FUNCTIONS, }, ], - getColumnsFor: () => Promise.resolve([{ name: 'clientip', type: 'ip' }]), + getColumnsFor: () => + Promise.resolve([{ name: 'clientip', type: 'ip', userDefined: false }]), }, }); @@ -102,7 +105,8 @@ describe('autocomplete.suggest', () => { callbacks: { canSuggestVariables: () => true, getVariables: () => [], - getColumnsFor: () => Promise.resolve([{ name: 'clientip', type: 'ip' }]), + getColumnsFor: () => + Promise.resolve([{ name: 'clientip', type: 'ip', userDefined: false }]), }, }); @@ -129,7 +133,8 @@ describe('autocomplete.suggest', () => { type: ESQLVariableType.FIELDS, }, ], - getColumnsFor: () => Promise.resolve([{ name: 'clientip', type: 'ip' }]), + getColumnsFor: () => + Promise.resolve([{ name: 'clientip', type: 'ip', userDefined: false }]), }, }); @@ -156,7 +161,8 @@ describe('autocomplete.suggest', () => { type: ESQLVariableType.TIME_LITERAL, }, ], - getColumnsFor: () => Promise.resolve([{ name: '@timestamp', type: 'date' }]), + getColumnsFor: () => + Promise.resolve([{ name: '@timestamp', type: 'date', userDefined: false }]), }, }); @@ -177,7 +183,8 @@ describe('autocomplete.suggest', () => { callbacks: { canSuggestVariables: () => true, getVariables: () => [], - getColumnsFor: () => Promise.resolve([{ name: 'bytes', type: 'double' }]), + getColumnsFor: () => + Promise.resolve([{ name: 'bytes', type: 'double', userDefined: false }]), }, }); @@ -198,7 +205,8 @@ describe('autocomplete.suggest', () => { callbacks: { canSuggestVariables: () => true, getVariables: () => [], - getColumnsFor: () => Promise.resolve([{ name: 'agent.name', type: 'keyword' }]), + getColumnsFor: () => + Promise.resolve([{ name: 'agent.name', type: 'keyword', userDefined: false }]), }, }); @@ -225,7 +233,8 @@ describe('autocomplete.suggest', () => { type: ESQLVariableType.VALUES, }, ], - getColumnsFor: () => Promise.resolve([{ name: 'agent.name', type: 'keyword' }]), + getColumnsFor: () => + Promise.resolve([{ name: 'agent.name', type: 'keyword', userDefined: false }]), }, }); @@ -246,7 +255,8 @@ describe('autocomplete.suggest', () => { callbacks: { canSuggestVariables: () => true, getVariables: () => [], - getColumnsFor: () => Promise.resolve([{ name: 'agent.name', type: 'keyword' }]), + getColumnsFor: () => + Promise.resolve([{ name: 'agent.name', type: 'keyword', userDefined: false }]), }, }); diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts index c18d4b92e9629..74ad4ea8e47c2 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/__tests__/helpers.ts @@ -46,15 +46,27 @@ export const fields: TestField[] = [ ...fieldTypes.map((type) => ({ name: `${camelCase(type)}Field`, type, + userDefined: false as const, + // suggestedAs is optional and omitted here })), - { name: 'any#Char$Field', type: 'double', suggestedAs: '`any#Char$Field`' }, - { name: 'kubernetes.something.something', type: 'double' }, + { + name: 'any#Char$Field', + type: 'double', + suggestedAs: '`any#Char$Field`', + userDefined: false as const, + }, + { + name: 'kubernetes.something.something', + type: 'double', + suggestedAs: undefined, + userDefined: false as const, + }, ]; export const lookupIndexFields: TestField[] = [ - { name: 'booleanField', type: 'boolean' }, - { name: 'dateField', type: 'date' }, - { name: 'joinIndexOnlyField', type: 'text' }, + { name: 'booleanField', type: 'boolean', userDefined: false }, + { name: 'dateField', type: 'date', userDefined: false }, + { name: 'joinIndexOnlyField', type: 'text', userDefined: false }, ]; export const indexes = ( diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts index d29d130cce494..ed7737f5c7cc8 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.test.ts @@ -124,82 +124,57 @@ describe('autocomplete', () => { testSuggestions('from a metadata _id | eval col0 = a | /', commands); }); - for (const command of ['keep', 'drop']) { - describe(command, () => { - testSuggestions(`from a | ${command} /`, getFieldNamesByType('any')); - testSuggestions( - `from a | ${command} keywordField, /`, - getFieldNamesByType('any').filter((name) => name !== 'keywordField') - ); + describe.each(['keep', 'drop'])('%s', (command) => { + testSuggestions(`from a | ${command} /`, getFieldNamesByType('any')); + testSuggestions( + `from a | ${command} keywordField, /`, + getFieldNamesByType('any').filter((name) => name !== 'keywordField') + ); - testSuggestions( - `from a | ${command} keywordField,/`, - getFieldNamesByType('any').filter((name) => name !== 'keywordField'), - ',' - ); + testSuggestions( + `from a | ${command} keywordField,/`, + getFieldNamesByType('any').filter((name) => name !== 'keywordField'), + ',' + ); - testSuggestions( - `from a_index | eval round(doubleField) + 1 | eval \`round(doubleField) + 1\` + 1 | eval \`\`\`round(doubleField) + 1\`\` + 1\` + 1 | eval \`\`\`\`\`\`\`round(doubleField) + 1\`\`\`\` + 1\`\` + 1\` + 1 | eval \`\`\`\`\`\`\`\`\`\`\`\`\`\`\`round(doubleField) + 1\`\`\`\`\`\`\`\` + 1\`\`\`\` + 1\`\` + 1\` + 1 | ${command} /`, - [ - ...getFieldNamesByType('any'), - '`round(doubleField) + 1`', - '```round(doubleField) + 1`` + 1`', - '```````round(doubleField) + 1```` + 1`` + 1`', - '```````````````round(doubleField) + 1```````` + 1```` + 1`` + 1`', - '```````````````````````````````round(doubleField) + 1```````````````` + 1```````` + 1```` + 1`` + 1`', - ], - undefined, - [ - [ - ...fields, - // the following non-field columns will come over the wire as part of the response - { - name: 'round(doubleField) + 1', - type: 'double', - }, - { - name: '`round(doubleField) + 1` + 1', - type: 'double', - }, - { - name: '```round(doubleField) + 1`` + 1` + 1', - type: 'double', - }, - { - name: '```````round(doubleField) + 1```` + 1`` + 1` + 1', - type: 'double', - }, - { - name: '```````````````round(doubleField) + 1```````` + 1```` + 1`` + 1` + 1', - type: 'double', - }, - ], - ] - ); + testSuggestions( + `from a_index | eval round(doubleField) + 1 | eval \`round(doubleField) + 1\` + 1 | eval \`\`\`round(doubleField) + 1\`\` + 1\` + 1 | eval \`\`\`\`\`\`\`round(doubleField) + 1\`\`\`\` + 1\`\` + 1\` + 1 | eval \`\`\`\`\`\`\`\`\`\`\`\`\`\`\`round(doubleField) + 1\`\`\`\`\`\`\`\` + 1\`\`\`\` + 1\`\` + 1\` + 1 | ${command} /`, + [ + ...getFieldNamesByType('any'), + '`round(doubleField) + 1`', + '```round(doubleField) + 1`` + 1`', + '```````round(doubleField) + 1```` + 1`` + 1`', + '```````````````round(doubleField) + 1```````` + 1```` + 1`` + 1`', + '```````````````````````````````round(doubleField) + 1```````````````` + 1```````` + 1```` + 1`` + 1`', + ] + ); - it('should not suggest already-used fields and user defined columns', async () => { - const { suggest: suggestTest } = await setup(); - const getSuggestions = async (query: string, opts?: SuggestOptions) => - (await suggestTest(query, opts)).map((value) => value.text); - - expect( - await getSuggestions('from a_index | EVAL foo = 1 | KEEP /', { - callbacks: { getColumnsFor: () => [...fields, { name: 'foo', type: 'integer' }] }, - }) - ).toContain('foo'); - expect( - await getSuggestions('from a_index | EVAL foo = 1 | KEEP foo, /', { - callbacks: { getColumnsFor: () => [...fields, { name: 'foo', type: 'integer' }] }, - }) - ).not.toContain('foo'); - - expect(await getSuggestions('from a_index | KEEP /')).toContain('doubleField'); - expect(await getSuggestions('from a_index | KEEP doubleField, /')).not.toContain( - 'doubleField' - ); - }); + it('should not suggest already-used fields and user defined columns', async () => { + const { suggest: suggestTest } = await setup(); + const getSuggestions = async (query: string, opts?: SuggestOptions) => + (await suggestTest(query, opts)).map((value) => value.text); + + expect( + await getSuggestions('from a_index | EVAL foo = 1 | KEEP /', { + callbacks: { + getColumnsFor: () => [...fields, { name: 'foo', type: 'integer', userDefined: false }], + }, + }) + ).toContain('foo'); + expect( + await getSuggestions('from a_index | EVAL foo = 1 | KEEP foo, /', { + callbacks: { + getColumnsFor: () => [...fields, { name: 'foo', type: 'integer', userDefined: false }], + }, + }) + ).not.toContain('foo'); + + expect(await getSuggestions('from a_index | KEEP /')).toContain('doubleField'); + expect(await getSuggestions('from a_index | KEEP doubleField, /')).not.toContain( + 'doubleField' + ); }); - } + }); // @TODO: get updated eval block from main describe('values suggestions', () => { @@ -226,7 +201,7 @@ describe('autocomplete', () => { const triggerOffset = statement.lastIndexOf(' '); await suggest(statement, triggerOffset + 1, callbackMocks); expect(callbackMocks.getColumnsFor).toHaveBeenCalledWith({ - query: 'from index_b', + query: 'FROM index_b', }); }); it('should send the fields query aware of the location', async () => { @@ -234,7 +209,7 @@ describe('autocomplete', () => { const statement = 'from index_d | drop | eval col0 = abs(doubleField) '; const triggerOffset = statement.lastIndexOf('p') + 1; // drop await suggest(statement, triggerOffset + 1, callbackMocks); - expect(callbackMocks.getColumnsFor).toHaveBeenCalledWith({ query: 'from index_d' }); + expect(callbackMocks.getColumnsFor).toHaveBeenCalledWith({ query: 'FROM index_d' }); }); }); @@ -1091,21 +1066,21 @@ describe('autocomplete', () => { 'FROM index_a | KEEP field.nam/', [{ text: 'field.name', rangeToReplace: { start: 20, end: 29 } }], undefined, - [[{ name: 'field.name', type: 'double' }]] + [[{ name: 'field.name', type: 'double', userDefined: false }]] ); // multi-line testSuggestions( 'FROM index_a\n| KEEP field.nam/', [{ text: 'field.name', rangeToReplace: { start: 20, end: 29 } }], undefined, - [[{ name: 'field.name', type: 'double' }]] + [[{ name: 'field.name', type: 'double', userDefined: false }]] ); // triple separator testSuggestions( 'FROM index_c\n| KEEP field.name.f/', [{ text: 'field.name.foo', rangeToReplace: { start: 20, end: 32 } }], undefined, - [[{ name: 'field.name.foo', type: 'double' }]] + [[{ name: 'field.name.foo', type: 'double', userDefined: false }]] ); // whitespace — we can't support this case yet because // we are relying on string checking instead of the AST :( @@ -1113,7 +1088,7 @@ describe('autocomplete', () => { 'FROM index_a | KEEP field . n/', [{ text: 'field . name', rangeToReplace: { start: 14, end: 22 } }], undefined, - [[{ name: 'field.name', type: 'double' }]] + [[{ name: 'field.name', type: 'double', userDefined: false }]] ); }); }); diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts index 4c06fdc67e9d5..0b6b41b5a187f 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/autocomplete.ts @@ -7,41 +7,39 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ +import type { ESQLAstQueryExpression } from '@kbn/esql-ast'; import { + ESQL_VARIABLES_PREFIX, + EsqlQuery, + esqlCommandRegistry, + getCommandAutocompleteDefinitions, parse, type ESQLAstItem, type ESQLCommand, type ESQLCommandOption, type ESQLFunction, - esqlCommandRegistry, - getCommandAutocompleteDefinitions, - ESQL_VARIABLES_PREFIX, } from '@kbn/esql-ast'; -import { EDITOR_MARKER } from '@kbn/esql-ast/src/definitions/constants'; -import { correctQuerySyntax } from '@kbn/esql-ast/src/definitions/utils/ast'; -import { - getControlSuggestionIfSupported, - buildFieldsDefinitionsWithMetadata, -} from '@kbn/esql-ast/src/definitions/utils'; import { getRecommendedQueriesSuggestionsFromStaticTemplates } from '@kbn/esql-ast/src/commands_registry/options/recommended_queries'; import type { - ESQLUserDefinedColumn, - ESQLFieldWithMetadata, + ESQLColumnData, GetColumnsByTypeFn, ISuggestionItem, } from '@kbn/esql-ast/src/commands_registry/types'; +import { + buildFieldsDefinitionsWithMetadata, + getControlSuggestionIfSupported, +} from '@kbn/esql-ast/src/definitions/utils'; +import { correctQuerySyntax } from '@kbn/esql-ast/src/definitions/utils/ast'; import { ESQLVariableType } from '@kbn/esql-types'; import type { LicenseType } from '@kbn/licensing-types'; -import { isSourceCommand } from '../shared/helpers'; -import { collectUserDefinedColumns } from '../shared/user_defined_columns'; import { getAstContext } from '../shared/context'; -import { getFieldsByTypeHelper, getSourcesHelper } from '../shared/resources_helpers'; +import { isSourceCommand } from '../shared/helpers'; +import { getColumnsByTypeHelper, getSourcesHelper } from '../shared/resources_helpers'; import type { ESQLCallbacks } from '../shared/types'; -import { getQueryForFields } from './helper'; -import { mapRecommendedQueriesFromExtensions } from './utils/recommended_queries_helpers'; import { getCommandContext } from './get_command_context'; +import { mapRecommendedQueriesFromExtensions } from './utils/recommended_queries_helpers'; -type GetFieldsMapFn = () => Promise>; +type GetColumnMapFn = () => Promise>; export async function suggest( fullText: string, @@ -58,14 +56,12 @@ export async function suggest( return []; } - // build the correct query to fetch the list of fields - const queryForFields = getQueryForFields(correctedQuery, root); - - const { getFieldsByType, getFieldsMap } = getFieldsByTypeRetriever( - queryForFields.replace(EDITOR_MARKER, ''), - resourceRetriever, - innerText + const { getColumnsByType, getColumnMap } = getColumnsByTypeRetriever( + root, + innerText, + resourceRetriever ); + const supportsControls = resourceRetriever?.canSuggestVariables?.() ?? false; const getVariables = resourceRetriever?.getVariables; const getSources = getSourcesHelper(resourceRetriever); @@ -113,10 +109,10 @@ export async function suggest( fromCommand = `FROM ${visibleSources[0].name}`; } - const { getFieldsByType: getFieldsByTypeEmptyState } = getFieldsByTypeRetriever( - fromCommand, - resourceRetriever, - innerText + const { getColumnsByType: getColumnsByTypeEmptyState } = getColumnsByTypeRetriever( + EsqlQuery.fromSrc(fromCommand).ast, + innerText, + resourceRetriever ); const editorExtensions = (await resourceRetriever?.getEditorExtensions?.(fromCommand)) ?? { recommendedQueries: [], @@ -127,7 +123,7 @@ export async function suggest( const recommendedQueriesSuggestionsFromStaticTemplates = await getRecommendedQueriesSuggestionsFromStaticTemplates( - getFieldsByTypeEmptyState, + getColumnsByTypeEmptyState, fromCommand ); recommendedQueriesSuggestions.push( @@ -161,8 +157,8 @@ export async function suggest( fullText, ast, astContext, - getFieldsByType, - getFieldsMap, + getColumnsByType, + getColumnMap, resourceRetriever, offset, hasMinimumLicenseRequired @@ -172,20 +168,20 @@ export async function suggest( return []; } -export function getFieldsByTypeRetriever( - queryForFields: string, - resourceRetriever?: ESQLCallbacks, - fullQuery?: string -): { getFieldsByType: GetColumnsByTypeFn; getFieldsMap: GetFieldsMapFn } { - const helpers = getFieldsByTypeHelper(queryForFields, resourceRetriever); +export function getColumnsByTypeRetriever( + query: ESQLAstQueryExpression, + queryText: string, + resourceRetriever?: ESQLCallbacks +): { getColumnsByType: GetColumnsByTypeFn; getColumnMap: GetColumnMapFn } { + const helpers = getColumnsByTypeHelper(query, queryText, resourceRetriever); const getVariables = resourceRetriever?.getVariables; const canSuggestVariables = resourceRetriever?.canSuggestVariables?.() ?? false; - const queryString = fullQuery ?? queryForFields; + const queryString = queryText; const lastCharacterTyped = queryString[queryString.length - 1]; const lastCharIsQuestionMark = lastCharacterTyped === ESQL_VARIABLES_PREFIX; return { - getFieldsByType: async ( + getColumnsByType: async ( expectedType: Readonly | Readonly = 'any', ignored: string[] = [], options @@ -194,24 +190,24 @@ export function getFieldsByTypeRetriever( ...options, supportsControls: canSuggestVariables && !lastCharIsQuestionMark, }; - const editorExtensions = (await resourceRetriever?.getEditorExtensions?.(queryForFields)) ?? { + const editorExtensions = (await resourceRetriever?.getEditorExtensions?.(queryText)) ?? { recommendedQueries: [], recommendedFields: [], }; const recommendedFieldsFromExtensions = editorExtensions.recommendedFields; - const fields = await helpers.getFieldsByType(expectedType, ignored); + const columns = await helpers.getColumnsByType(expectedType, ignored); return buildFieldsDefinitionsWithMetadata( - fields, + columns, recommendedFieldsFromExtensions, updatedOptions, await getVariables?.() ); }, - getFieldsMap: helpers.getFieldsMap, + getColumnMap: helpers.getColumnMap, }; } -function findNewUserDefinedColumn(userDefinedColumns: Map) { +function findNewUserDefinedColumn(userDefinedColumns: Map) { let autoGeneratedColumnCounter = 0; let name = `col${autoGeneratedColumnCounter++}`; while (userDefinedColumns.has(name)) { @@ -230,7 +226,7 @@ async function getSuggestionsWithinCommandExpression( containingFunction?: ESQLFunction; }, getColumnsByType: GetColumnsByTypeFn, - getFieldsMap: GetFieldsMapFn, + getColumnMap: GetColumnMapFn, callbacks?: ESQLCallbacks, offset?: number, hasMinimumLicenseRequired?: (minimumLicenseRequired: LicenseType) => boolean @@ -243,23 +239,19 @@ async function getSuggestionsWithinCommandExpression( } // collect all fields + userDefinedColumns to suggest - const fieldsMap: Map = await getFieldsMap(); - const anyUserDefinedColumns = collectUserDefinedColumns(commands, fieldsMap, innerText); - - const references = { fields: fieldsMap, userDefinedColumns: anyUserDefinedColumns }; + const columnMap: Map = await getColumnMap(); + const references = { columns: columnMap }; - const getSuggestedUserDefinedColumnName = (extraFieldNames?: string[]) => { - if (!extraFieldNames?.length) { - return findNewUserDefinedColumn(anyUserDefinedColumns); + const getSuggestedUserDefinedColumnName = (extraColumnNames?: string[]) => { + if (!extraColumnNames?.length) { + return findNewUserDefinedColumn(columnMap); } - const augmentedFieldsMap = new Map(fieldsMap); - extraFieldNames.forEach((name) => { - augmentedFieldsMap.set(name, { name, type: 'double' }); + const augmentedColumnsMap = new Map(columnMap); + extraColumnNames.forEach((name) => { + augmentedColumnsMap.set(name, { name, type: 'double', userDefined: false }); }); - return findNewUserDefinedColumn( - collectUserDefinedColumns(commands, augmentedFieldsMap, innerText) - ); + return findNewUserDefinedColumn(augmentedColumnsMap); }; const additionalCommandContext = await getCommandContext( diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts deleted file mode 100644 index 36954fa514f92..0000000000000 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.ts +++ /dev/null @@ -1,79 +0,0 @@ -/* - * 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 type { ESQLAstQueryExpression } from '@kbn/esql-ast'; -import { type ESQLCommand, BasicPrettyPrinter, Builder, EDITOR_MARKER } from '@kbn/esql-ast'; - -/** - * This function is used to build the query that will be used to compute the - * available fields for the current cursor location. - * - * Generally, this is the user's query up to the end of the previous command. - * - * @param queryString The original query string - * @param commands - * @returns - */ -export function getQueryForFields(queryString: string, root: ESQLAstQueryExpression): string { - const commands = root.commands; - const lastCommand = commands[commands.length - 1]; - if (lastCommand && lastCommand.name === 'fork' && lastCommand.args.length > 0) { - /** - * This translates the current fork command branch into a simpler but equivalent - * query that is compatible with the existing field computation/caching strategy. - * - * The intuition here is that if the cursor is within a fork branch, the - * previous context is equivalent to a query without the FORK command.: - * - * Original query: FROM lolz | EVAL foo = 1 | FORK (EVAL bar = 2) (EVAL baz = 3 | WHERE /) - * Simplified: FROM lolz | EVAL foo = 1 | EVAL baz = 3 - */ - const currentBranch = lastCommand.args[lastCommand.args.length - 1] as ESQLAstQueryExpression; - const newCommands = commands.slice(0, -1).concat(currentBranch.commands.slice(0, -1)); - return BasicPrettyPrinter.print({ ...root, commands: newCommands }); - } - - if (lastCommand && lastCommand.name === 'eval') { - const endsWithComma = queryString.replace(EDITOR_MARKER, '').trim().endsWith(','); - if (lastCommand.args.length > 1 || endsWithComma) { - /** - * If we get here, we know that we have a multi-expression EVAL statement. - * - * e.g. EVAL foo = 1, bar = foo + 1, baz = bar + 1 - * - * In order for this to work with the caching system which expects field availability to be - * delineated by pipes, we need to split the current EVAL command into an equivalent - * set of single-expression EVAL commands. - * - * Original query: FROM lolz | EVAL foo = 1, bar = foo + 1, baz = bar + 1, / - * Simplified: FROM lolz | EVAL foo = 1 | EVAL bar = foo + 1 | EVAL baz = bar + 1 - */ - const individualEvalCommands: ESQLCommand[] = []; - for (const expression of lastCommand.args) { - individualEvalCommands.push(Builder.command({ name: 'eval', args: [expression] })); - } - const newCommands = commands - .slice(0, -1) - .concat(endsWithComma ? individualEvalCommands : individualEvalCommands.slice(0, -1)); - return BasicPrettyPrinter.print({ ...root, commands: newCommands }); - } - } - - // If there is only one source command and it does not require fields, do not - // fetch fields, hence return an empty string. - return commands.length === 1 && ['row', 'show'].includes(commands[0].name) - ? '' - : buildQueryUntilPreviousCommand(queryString, commands); -} - -// TODO consider replacing this with a pretty printer-based solution -function buildQueryUntilPreviousCommand(queryString: string, commands: ESQLCommand[]) { - const prevCommand = commands[Math.max(commands.length - 2, 0)]; - return prevCommand ? queryString.substring(0, prevCommand.location.max + 1) : queryString; -} diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts index bfffbfef723d7..47493ec786800 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/utils/ecs_metadata_helper.test.ts @@ -7,15 +7,16 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ -import type { ESQLFieldWithMetadata } from '@kbn/esql-ast/src/commands_registry/types'; +import type { ESQLColumnData } from '@kbn/esql-ast/src/commands_registry/types'; +import { type ESQLFieldWithMetadata } from '@kbn/esql-ast/src/commands_registry/types'; import { type ECSMetadata, enrichFieldsWithECSInfo } from './ecs_metadata_helper'; describe('enrichFieldsWithECSInfo', () => { it('should return original columns if fieldsMetadata is not provided', async () => { const columns: ESQLFieldWithMetadata[] = [ - { name: 'ecs.version', type: 'keyword' }, - { name: 'field1', type: 'text' }, - { name: 'field2', type: 'double' }, + { name: 'ecs.version', type: 'keyword', userDefined: false }, + { name: 'field1', type: 'text', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, ]; const result = await enrichFieldsWithECSInfo(columns); @@ -24,80 +25,56 @@ describe('enrichFieldsWithECSInfo', () => { it('should return columns with metadata if both name and type match with ECS fields', async () => { const columns: ESQLFieldWithMetadata[] = [ - { name: 'ecs.field', type: 'text' }, - { name: 'ecs.fakeBooleanField', type: 'boolean' }, - { name: 'field2', type: 'double' }, + { name: 'ecs.field', type: 'text', userDefined: false }, + { name: 'ecs.fakeBooleanField', type: 'boolean', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, ]; - const fieldsMetadata = { - getClient: jest.fn().mockResolvedValue({ - find: jest.fn().mockResolvedValue({ - fields: { - 'ecs.version': { description: 'ECS version field', type: 'keyword' }, - 'ecs.field': { description: 'ECS field description', type: 'text' }, - 'ecs.fakeBooleanField': { - description: 'ECS fake boolean field description', - type: 'keyword', - }, - }, - }), - }), + + const fieldsMetadataCache: ECSMetadata = { + 'ecs.version': { description: 'ECS version field', type: 'keyword' }, + 'ecs.field': { description: 'ECS field description', type: 'text' }, + 'ecs.fakeBooleanField': { + description: 'ECS fake boolean field description', + type: 'keyword', + }, }; - const fieldsMetadataCache = await ( - await fieldsMetadata.getClient() - ).find({ - attributes: ['type'], - }); - const result = await enrichFieldsWithECSInfo( - columns, - fieldsMetadataCache.fields as ECSMetadata - ); + const result = enrichFieldsWithECSInfo(columns, fieldsMetadataCache); expect(result).toEqual([ { name: 'ecs.field', type: 'text', isEcs: true, + userDefined: false, }, - { name: 'ecs.fakeBooleanField', type: 'boolean' }, - { name: 'field2', type: 'double' }, + { name: 'ecs.fakeBooleanField', type: 'boolean', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, ]); }); it('should handle keyword suffix correctly', async () => { const columns: ESQLFieldWithMetadata[] = [ - { name: 'ecs.version', type: 'keyword' }, - { name: 'ecs.version.keyword', type: 'keyword' }, - { name: 'field2', type: 'double' }, + { name: 'ecs.version', type: 'keyword', userDefined: false }, + { name: 'ecs.version.keyword', type: 'keyword', userDefined: false }, + { name: 'field2', type: 'double', userDefined: false }, ]; - const fieldsMetadata = { - getClient: jest.fn().mockResolvedValue({ - find: jest.fn().mockResolvedValue({ - fields: { - 'ecs.version': { description: 'ECS version field', type: 'keyword' }, - } as unknown as ECSMetadata, - }), - }), + + const fieldsMetadataCache: ECSMetadata = { + 'ecs.version': { description: 'ECS version field', type: 'keyword' }, }; - const fieldsMetadataCache = await ( - await fieldsMetadata.getClient() - ).find({ - attributes: ['type'], - }); - const result = await enrichFieldsWithECSInfo( - columns, - fieldsMetadataCache.fields as ECSMetadata - ); + const result = enrichFieldsWithECSInfo(columns, fieldsMetadataCache); - expect(result).toEqual([ - { name: 'ecs.version', type: 'keyword', isEcs: true }, + expect(result).toEqual([ + { name: 'ecs.version', type: 'keyword', isEcs: true, userDefined: false }, { name: 'ecs.version.keyword', type: 'keyword', isEcs: true, + userDefined: false, }, - { name: 'field2', type: 'double' }, + { name: 'field2', type: 'double', userDefined: false }, ]); }); }); diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/expand_evals.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/expand_evals.ts new file mode 100644 index 0000000000000..07338095d16ef --- /dev/null +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/expand_evals.ts @@ -0,0 +1,41 @@ +/* + * 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 type { ESQLCommand } from '@kbn/esql-ast'; +import { Builder } from '@kbn/esql-ast'; + +/** + * Expands EVAL commands into separate single-expression EVAL commands. + * + * E.g. EVAL foo = 1, bar = 2 => [EVAL foo = 1, EVAL bar = 2] + * + * This is logically equivalent and makes validation and field existence detection much easier. + * + * @param commands The list of commands to expand. + * @returns The expanded list of commands. + */ +export function expandEvals(commands: ESQLCommand[]): ESQLCommand[] { + const expanded: ESQLCommand[] = []; + for (const command of commands) { + if (command.name.toLowerCase() === 'eval') { + for (const arg of command.args) { + expanded.push( + Builder.command({ + name: 'eval', + args: [arg], + location: command.location, + }) + ); + } + } else { + expanded.push(command); + } + } + return expanded; +} diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts index 1eb2d013ad01f..8d77fb6432707 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/helpers.ts @@ -7,21 +7,24 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ import { - type ESQLAstCommand, + BasicPrettyPrinter, esqlCommandRegistry, - type FieldType, + isSource, + mutate, + synth, + type ESQLAstCommand, type FunctionDefinition, } from '@kbn/esql-ast'; import type { + ESQLColumnData, ESQLFieldWithMetadata, - ESQLUserDefinedColumn, + ESQLPolicy, } from '@kbn/esql-ast/src/commands_registry/types'; -import type { ESQLParamLiteral } from '@kbn/esql-ast/src/types'; -import { uniqBy } from 'lodash'; +import type { ESQLAstQueryExpression, ESQLParamLiteral } from '@kbn/esql-ast/src/types'; +import type { IAdditionalFields } from '@kbn/esql-ast/src/commands_registry/registry'; import { enrichFieldsWithECSInfo } from '../autocomplete/utils/ecs_metadata_helper'; import type { ESQLCallbacks } from './types'; -import { collectUserDefinedColumns } from './user_defined_columns'; export function nonNullable(v: T): v is NonNullable { return v != null; @@ -77,28 +80,6 @@ export function getParamAtPosition( return params.length > position ? params[position] : minParams ? params[params.length - 1] : null; } -// --- Fields helpers --- - -export function transformMapToESQLFields( - inputMap: Map -): ESQLFieldWithMetadata[] { - const esqlFields: ESQLFieldWithMetadata[] = []; - - for (const [, userDefinedColumns] of inputMap) { - for (const userDefinedColumn of userDefinedColumns) { - // Only include userDefinedColumns that have a known type - if (userDefinedColumn.type) { - esqlFields.push({ - name: userDefinedColumn.name, - type: userDefinedColumn.type as FieldType, - }); - } - } - } - - return esqlFields; -} - async function getEcsMetadata(resourceRetriever?: ESQLCallbacks) { if (!resourceRetriever?.getFieldsMetadata) { return undefined; @@ -125,30 +106,66 @@ export async function getFieldsFromES(query: string, resourceRetriever?: ESQLCal * @param previousPipeFields, the fields from the previous pipe * @returns a list of fields that are available for the current pipe */ -export async function getCurrentQueryAvailableFields( - query: string, +export async function getCurrentQueryAvailableColumns( commands: ESQLAstCommand[], - previousPipeFields: ESQLFieldWithMetadata[] + previousPipeFields: ESQLColumnData[], + fetchFields: (query: string) => Promise, + getPolicies: () => Promise>, + originalQueryText: string ) { - const cacheCopy = new Map(); - previousPipeFields.forEach((field) => cacheCopy.set(field.name, field)); const lastCommand = commands[commands.length - 1]; - const commandDefinition = esqlCommandRegistry.getCommandByName(lastCommand.name); - - // If the command has a columnsAfter function, use it to get the fields - if (commandDefinition?.methods.columnsAfter) { - const userDefinedColumns = collectUserDefinedColumns([lastCommand], cacheCopy, query); - - return commandDefinition.methods.columnsAfter(lastCommand, previousPipeFields, { - userDefinedColumns, - }); - } else { - // If the command doesn't have a columnsAfter function, use the default behavior - const userDefinedColumns = collectUserDefinedColumns(commands, cacheCopy, query); - const arrayOfUserDefinedColumns: ESQLFieldWithMetadata[] = transformMapToESQLFields( - userDefinedColumns ?? new Map() + const commandDef = esqlCommandRegistry.getCommandByName(lastCommand.name); + + const getJoinFields = (command: ESQLAstCommand): Promise => { + const joinSummary = mutate.commands.join.summarize({ + type: 'query', + commands: [command], + } as ESQLAstQueryExpression); + const joinIndices = joinSummary.map(({ target: { index } }) => index); + if (joinIndices.length > 0) { + const joinFieldQuery = synth.cmd`FROM ${joinIndices}`.toString(); + return fetchFields(joinFieldQuery); + } + return Promise.resolve([]); + }; + + const getEnrichFields = async (command: ESQLAstCommand): Promise => { + if (!isSource(command.args[0])) { + return []; + } + + const policyName = command.args[0].name; + + const policies = await getPolicies(); + const policy = policies.get(policyName); + + if (policy) { + const fieldsQuery = `FROM ${policy.sourceIndices.join( + ', ' + )} | KEEP ${policy.enrichFields.join(', ')}`; + return fetchFields(fieldsQuery); + } + + return []; + }; + + const getFromFields = (command: ESQLAstCommand): Promise => { + return fetchFields(BasicPrettyPrinter.command(command)); + }; + + const additionalFields: IAdditionalFields = { + fromJoin: getJoinFields, + fromEnrich: getEnrichFields, + fromFrom: getFromFields, + }; + + if (commandDef?.methods.columnsAfter) { + return commandDef.methods.columnsAfter( + lastCommand, + previousPipeFields, + originalQueryText, + additionalFields ); - const allFields = uniqBy([...(previousPipeFields ?? []), ...arrayOfUserDefinedColumns], 'name'); - return allFields; } + return previousPipeFields; } diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.test.ts deleted file mode 100644 index c69767d4a3cb1..0000000000000 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.test.ts +++ /dev/null @@ -1,168 +0,0 @@ -/* - * 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 { removeLastPipe, processPipes, toSingleLine, removeComments } from './query_string_utils'; - -describe('query_string_utils', () => { - describe('removeComments', () => { - it('should remove single-line comments', () => { - const text = ` - FROM users // This is a comment - | WHERE age > 18 - `; - const expected = `FROM users | WHERE age > 18`; - expect(removeComments(text)).toBe(expected); - }); - - it('should remove multi-line comments', () => { - const text = ` - FROM /* This is a - multi-line - comment */ - products - `; - const expected = `FROM - products`; - expect(removeComments(text)).toBe(expected); - }); - - it('should remove both single-line and multi-line comments', () => { - const text = ` - FROM items // Get the name - /* The price of the - product */ - | KEEP name, price - `; - const expected = `FROM items - | KEEP name, price`; - expect(removeComments(text)).toBe(expected); - }); - - it('should handle text with no comments', () => { - const text = ` - FROM orders - | KEEP order_id, status - | WHERE status = 'pending'; - `; - expect(removeComments(text)).toBe(text.trim()); - }); - - it('should handle comments at the beginning and end of the text', () => { - const text = ` - // Initial comment - FROM logs - /* Final - comment */ - `; - const expected = `FROM logs`; - expect(removeComments(text)).toBe(expected); - }); - - it('should handle consecutive single-line comments', () => { - const text = ` - // Comment line 1 - // Comment line 2 - FROM events | STATS COUNT(*) - `; - const expected = ` - FROM events | STATS COUNT(*) - `; - expect(removeComments(text)).toBe(expected.trim()); - }); - }); - - describe('removeLastPipe', () => { - it('should remove the last pipe and any trailing whitespace', () => { - expect(removeLastPipe('value1|value2|')).toBe('value1|value2'); - expect(removeLastPipe('value1|value2| ')).toBe('value1|value2'); - }); - - it('should return the original string if there is no pipe', () => { - expect(removeLastPipe('FROM index')).toBe('FROM index'); - expect(removeLastPipe('FROM index ')).toBe('FROM index'); - }); - - it('should handle strings with multiple pipes correctly', () => { - expect(removeLastPipe('FROM index | STATS count() | DROP field1 ')).toBe( - 'FROM index | STATS count()' - ); - }); - - it('should handle an empty string', () => { - expect(removeLastPipe('')).toBe(''); - }); - - it('should handle a string with only whitespace', () => { - expect(removeLastPipe(' ')).toBe(''); - }); - }); - - describe('processPipes', () => { - it('should return an array of strings, each progressively including parts separated by " | "', () => { - const input = 'FROM index|EVAL col = ABS(numeric) | KEEP col'; - const expected = [ - 'FROM index', - 'FROM index | EVAL col = ABS(numeric)', - 'FROM index | EVAL col = ABS(numeric) | KEEP col', - ]; - expect(processPipes(input)).toEqual(expected); - }); - - it('should handle leading and trailing whitespace in parts', () => { - const input = ' FROM index | EVAL col = ABS(numeric) | KEEP col '; - const expected = [ - 'FROM index', - 'FROM index | EVAL col = ABS(numeric)', - 'FROM index | EVAL col = ABS(numeric) | KEEP col', - ]; - expect(processPipes(input)).toEqual(expected); - }); - - it('should return an array with the trimmed input if there are no pipes', () => { - const input = 'FROM index'; - const expected = ['FROM index']; - expect(processPipes(input)).toEqual(expected); - - const inputWithWhitespace = ' FROM index '; - const expectedWithWhitespace = ['FROM index']; - expect(processPipes(inputWithWhitespace)).toEqual(expectedWithWhitespace); - }); - - it('should ignore comments', () => { - const input = '// This is an ES|QL query \n FROM index'; - const expected = ['FROM index']; - expect(processPipes(input)).toEqual(expected); - }); - }); - - describe('toSingleLine', () => { - it('should convert a multi-line pipe-separated string to a single line with " | " as separator', () => { - const input = 'FROM index \n|EVAL col = ABS(numeric)\n|KEEP col'; - const expected = 'FROM index | EVAL col = ABS(numeric) | KEEP col'; - expect(toSingleLine(input)).toBe(expected); - }); - - it('should trim whitespace from each part', () => { - const input = ' FROM index | EVAL col = ABS(numeric) | KEEP col '; - const expected = 'FROM index | EVAL col = ABS(numeric) | KEEP col'; - expect(toSingleLine(input)).toBe(expected); - }); - - it('should trim whitespace from each part for multi-line strings', () => { - const input = ' FROM index \n| EVAL col = ABS(numeric) \n| KEEP col '; - const expected = 'FROM index | EVAL col = ABS(numeric) | KEEP col'; - expect(toSingleLine(input)).toBe(expected); - }); - - it('should ignore comments', () => { - const input = '// This is an ES|QL query \n FROM index'; - const expected = 'FROM index'; - expect(toSingleLine(input)).toEqual(expected); - }); - }); -}); diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.ts deleted file mode 100644 index 43229074c02df..0000000000000 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/query_string_utils.ts +++ /dev/null @@ -1,51 +0,0 @@ -/* - * 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". - */ -export function removeComments(text: string): string { - // Remove single-line comments - const withoutSingleLineComments = text.replace(/\/\/.*?(?:\r\n|\r|\n|$)/gm, ''); - // Remove multi-line comments - const withoutMultiLineComments = withoutSingleLineComments.replace(/\/\*[\s\S]*?\*\//g, ''); - return withoutMultiLineComments.trim(); -} - -export function removeLastPipe(inputString: string): string { - const queryNoComments = removeComments(inputString); - const lastPipeIndex = queryNoComments.lastIndexOf('|'); - if (lastPipeIndex !== -1) { - return queryNoComments.substring(0, lastPipeIndex).trimEnd(); - } - return queryNoComments.trimEnd(); -} - -export function processPipes(inputString: string) { - const queryNoComments = removeComments(inputString); - const parts = queryNoComments.split('|'); - const results = []; - let currentString = ''; - - for (let i = 0; i < parts.length; i++) { - const part = parts[i].trim(); - if (i === 0) { - currentString = part; - } else { - currentString += ' | ' + part; - } - results.push(currentString.trim()); - } - return results; -} - -export function toSingleLine(inputString: string): string { - const queryNoComments = removeComments(inputString); - return queryNoComments - .split('|') - .map((line) => line.trim()) - .filter((line) => line !== '') - .join(' | '); -} diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.test.ts similarity index 90% rename from src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.test.ts rename to src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.test.ts index bdaa985966c16..b623e75087d0e 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/autocomplete/helper.test.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.test.ts @@ -7,8 +7,8 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ -import { EDITOR_MARKER, parse } from '@kbn/esql-ast'; -import { getQueryForFields } from './helper'; +import { parse, EDITOR_MARKER } from '@kbn/esql-ast'; +import { getQueryForFields } from './resources_helpers'; describe('getQueryForFields', () => { const assert = (query: string, expected: string) => { @@ -59,9 +59,4 @@ describe('getQueryForFields', () => { | EVAL foo = 1, ${EDITOR_MARKER}`; assert(query3, 'FROM index | EVAL foo = 1'); }); - - it('should return empty string if non-FROM source command', () => { - assert('ROW field1 = 1', ''); - assert('SHOW INFO', ''); - }); }); diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts index b4e15d99eaf1c..01151fec9ac14 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts @@ -7,15 +7,20 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ -import { parse } from '@kbn/esql-ast'; -import type { ESQLFieldWithMetadata } from '@kbn/esql-ast/src/commands_registry/types'; +import type { ESQLAstQueryExpression } from '@kbn/esql-ast'; +import { BasicPrettyPrinter, Builder, EDITOR_MARKER, EsqlQuery } from '@kbn/esql-ast'; +import type { + ESQLColumnData, + ESQLFieldWithMetadata, + ESQLPolicy, +} from '@kbn/esql-ast/src/commands_registry/types'; import type { ESQLCallbacks } from './types'; -import { getFieldsFromES, getCurrentQueryAvailableFields } from './helpers'; -import { removeLastPipe, processPipes, toSingleLine } from './query_string_utils'; +import { getFieldsFromES, getCurrentQueryAvailableColumns } from './helpers'; +import { expandEvals } from './expand_evals'; export const NOT_SUGGESTED_TYPES = ['unsupported']; -const cache = new Map(); +const cache = new Map(); // Function to check if a key exists in the cache, ignoring case function checkCacheInsensitive(keyToCheck: string) { @@ -42,57 +47,102 @@ function getValueInsensitive(keyToCheck: string) { * for the next time the same query is used. * @param queryText */ -async function cacheFieldsForQuery(queryText: string) { - const existsInCache = checkCacheInsensitive(queryText); +async function cacheColumnsForQuery( + query: ESQLAstQueryExpression, + fetchFields: (query: string) => Promise, + getPolicies: () => Promise>, + originalQueryText: string +) { + let cacheKey: string; + try { + cacheKey = BasicPrettyPrinter.print(query); + } catch { + // for some syntactically incorrect queries + // the printer will throw. They're incorrect + // anyways, so just move on — ANTLR errors will + // be reported. + return; + } + + const existsInCache = checkCacheInsensitive(cacheKey); if (existsInCache) { // this is already in the cache return; } - const queryTextWithoutLastPipe = removeLastPipe(queryText); - // retrieve the user defined fields from the query without an extra call - const fieldsAvailableAfterPreviousCommand = getValueInsensitive(queryTextWithoutLastPipe); - if (fieldsAvailableAfterPreviousCommand && fieldsAvailableAfterPreviousCommand?.length) { - const { root } = parse(queryText); - const availableFields = await getCurrentQueryAvailableFields( - queryText, - root.commands, - fieldsAvailableAfterPreviousCommand - ); - cache.set(queryText, availableFields); - } + + const queryBeforeCurrentCommand = BasicPrettyPrinter.print({ + ...query, + commands: query.commands.slice(0, -1), + }); + const fieldsAvailableAfterPreviousCommand = getValueInsensitive(queryBeforeCurrentCommand) ?? []; + + const availableFields = await getCurrentQueryAvailableColumns( + query.commands, + fieldsAvailableAfterPreviousCommand, + fetchFields, + getPolicies, + originalQueryText + ); + + cache.set(cacheKey, availableFields); } -export function getFieldsByTypeHelper(queryText: string, resourceRetriever?: ESQLCallbacks) { - const getFields = async () => { - // in some cases (as in the case of ROW or SHOW) the query is not set - if (!queryText) { +/** + * Efficiently computes the list of columns available after the given query + * and returns column querying utilities. + * + * @param root + * @param originalQueryText the text of the original query, used to infer column names from expressions + * @param resourceRetriever + * @returns + */ +export function getColumnsByTypeHelper( + query: ESQLAstQueryExpression, + originalQueryText: string, + resourceRetriever?: ESQLCallbacks +) { + const queryForFields = getQueryForFields(originalQueryText, query); + const root = EsqlQuery.fromSrc(queryForFields).ast; + + const cacheColumns = async () => { + if (!queryForFields) { return; } - const [sourceCommand, ...partialQueries] = processPipes(queryText); + const getFields = async (queryToES: string) => { + const cached = getValueInsensitive(queryToES); + if (cached) { + return cached as ESQLFieldWithMetadata[]; + } + const fields = await getFieldsFromES(queryToES, resourceRetriever); + cache.set(queryToES, fields); + return fields; + }; - // retrieve the index fields from ES ONLY if the source command is not in the cache - const existsInCache = getValueInsensitive(sourceCommand); - if (!existsInCache) { - const fieldsWithMetadata = await getFieldsFromES(sourceCommand, resourceRetriever); - cache.set(sourceCommand, fieldsWithMetadata); + const subqueries = []; + for (let i = 0; i < root.commands.length; i++) { + subqueries.push(Builder.expression.query(root.commands.slice(0, i + 1))); } + const getPolicies = async () => { + const policies = (await resourceRetriever?.getPolicies?.()) ?? []; + return new Map(policies.map((p) => [p.name, p])); + }; + // build fields cache for every partial query - for (const query of partialQueries) { - await cacheFieldsForQuery(query); + for (const subquery of subqueries) { + await cacheColumnsForQuery(subquery, getFields, getPolicies, originalQueryText); } }; return { - getFieldsByType: async ( + getColumnsByType: async ( expectedType: Readonly | Readonly = 'any', ignored: string[] = [] - ): Promise => { + ): Promise => { const types = Array.isArray(expectedType) ? expectedType : [expectedType]; - await getFields(); - const queryTextForCacheSearch = toSingleLine(queryText); - const cachedFields = getValueInsensitive(queryTextForCacheSearch); + await cacheColumns(); + const cachedFields = getValueInsensitive(queryForFields); return ( cachedFields?.filter(({ name, type }) => { const ts = Array.isArray(type) ? type : [type]; @@ -105,11 +155,10 @@ export function getFieldsByTypeHelper(queryText: string, resourceRetriever?: ESQ }) || [] ); }, - getFieldsMap: async () => { - await getFields(); - const queryTextForCacheSearch = toSingleLine(queryText); - const cachedFields = getValueInsensitive(queryTextForCacheSearch); - const cacheCopy = new Map(); + getColumnMap: async (): Promise> => { + await cacheColumns(); + const cachedFields = getValueInsensitive(queryForFields); + const cacheCopy = new Map(); cachedFields?.forEach((field) => cacheCopy.set(field.name, field)); return cacheCopy; }, @@ -137,3 +186,66 @@ export function getSourcesHelper(resourceRetriever?: ESQLCallbacks) { return (await resourceRetriever?.getSources?.()) || []; }; } + +/** + * This function is used to build the query that will be used to compute the + * fields available at the final position. It is robust to final partial commands + * e.g. "FROM logs* | EVAL foo = 1 | EVAL " + * + * Generally, this is the user's query up to the end of the previous command, but there + * are special cases for multi-expression EVAL and FORK branches. + * + * @param queryString The original query string + * @param commands + * @returns + */ +export function getQueryForFields(queryString: string, root: ESQLAstQueryExpression): string { + const commands = root.commands; + const lastCommand = commands[commands.length - 1]; + if (lastCommand && lastCommand.name === 'fork' && lastCommand.args.length > 0) { + /** + * This flattens the current fork branch into a simpler but equivalent + * query that is compatible with the existing field computation/caching strategy. + * + * The intuition here is that if the cursor is within a fork branch, the + * previous context is equivalent to a query without the FORK command: + * + * Original query: FROM lolz | EVAL foo = 1 | FORK (EVAL bar = 2) (EVAL baz = 3 | WHERE /) + * Simplified: FROM lolz | EVAL foo = 1 | EVAL baz = 3 + */ + const currentBranch = lastCommand.args[lastCommand.args.length - 1] as ESQLAstQueryExpression; + const newCommands = commands.slice(0, -1).concat(currentBranch.commands.slice(0, -1)); + return BasicPrettyPrinter.print({ ...root, commands: newCommands }); + } + + if (lastCommand && lastCommand.name === 'eval') { + const endsWithComma = queryString.replace(EDITOR_MARKER, '').trim().endsWith(','); + if (lastCommand.args.length > 1 || endsWithComma) { + /** + * If we get here, we know that we have a multi-expression EVAL statement. + * + * e.g. EVAL foo = 1, bar = foo + 1, baz = bar + 1 + * + * In order for this to work with the caching system which expects field availability to be + * delineated by pipes, we need to split the current EVAL command into an equivalent + * set of single-expression EVAL commands. + * + * Original query: FROM lolz | EVAL foo = 1, bar = foo + 1, baz = bar + 1, / + * Simplified: FROM lolz | EVAL foo = 1 | EVAL bar = foo + 1 | EVAL baz = bar + 1 + */ + const expanded = expandEvals(commands); + const newCommands = expanded.slice(0, endsWithComma ? undefined : -1); + return BasicPrettyPrinter.print({ ...root, commands: newCommands }); + } + } + + return buildQueryUntilPreviousCommand(root); +} + +function buildQueryUntilPreviousCommand(root: ESQLAstQueryExpression) { + if (root.commands.length === 1) { + return BasicPrettyPrinter.print({ ...root.commands[0] }); + } else { + return BasicPrettyPrinter.print({ ...root, commands: root.commands.slice(0, -1) }); + } +} diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts deleted file mode 100644 index 0ae6ba0f2a47e..0000000000000 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/user_defined_columns.ts +++ /dev/null @@ -1,177 +0,0 @@ -/* - * 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 type { ESQLAst, ESQLAstItem, ESQLCommand, ESQLFunction } from '@kbn/esql-ast'; -import { isColumn, isFunctionExpression } from '@kbn/esql-ast'; -import { getExpressionType } from '@kbn/esql-ast/src/definitions/utils'; -import { EDITOR_MARKER } from '@kbn/esql-ast/src/definitions/constants'; -import { Visitor } from '@kbn/esql-ast/src/visitor'; -import type { - ESQLUserDefinedColumn, - ESQLFieldWithMetadata, -} from '@kbn/esql-ast/src/commands_registry/types'; - -function addToUserDefinedColumnOccurrences( - userDefinedColumns: Map, - instance: ESQLUserDefinedColumn -) { - if (!userDefinedColumns.has(instance.name)) { - userDefinedColumns.set(instance.name, []); - } - const userDefinedColumnsOccurrencies = userDefinedColumns.get(instance.name)!; - userDefinedColumnsOccurrencies.push(instance); -} - -function addToUserDefinedColumns( - oldArg: ESQLAstItem, - newArg: ESQLAstItem, - fields: Map, - userDefinedColumns: Map -) { - if (isColumn(oldArg) && isColumn(newArg)) { - const newUserDefinedColumn: ESQLUserDefinedColumn = { - name: newArg.parts.join('.'), - type: 'double' /* fallback to number */, - location: newArg.location, - }; - // Now workout the exact type - // it can be a rename of another userDefinedColumn as well - const oldRef = - fields.get(oldArg.parts.join('.')) || userDefinedColumns.get(oldArg.parts.join('.')); - if (oldRef) { - addToUserDefinedColumnOccurrences(userDefinedColumns, newUserDefinedColumn); - newUserDefinedColumn.type = Array.isArray(oldRef) ? oldRef[0].type : oldRef.type; - } - } -} - -export function excludeUserDefinedColumnsFromCurrentCommand( - commands: ESQLCommand[], - currentCommand: ESQLCommand, - fieldsMap: Map, - queryString: string -) { - const anyUserDefinedColumns = collectUserDefinedColumns(commands, fieldsMap, queryString); - const currentCommandUserDefinedColumns = collectUserDefinedColumns( - [currentCommand], - fieldsMap, - queryString - ); - const resultUserDefinedColumns = new Map(); - anyUserDefinedColumns.forEach((value, key) => { - if (!currentCommandUserDefinedColumns.has(key)) { - resultUserDefinedColumns.set(key, value); - } - }); - return resultUserDefinedColumns; -} - -function addUserDefinedColumnFromAssignment( - assignOperation: ESQLFunction, - userDefinedColumns: Map, - fields: Map -) { - if (isColumn(assignOperation.args[0])) { - const rightHandSideArgType = getExpressionType( - assignOperation.args[1], - fields, - userDefinedColumns - ); - addToUserDefinedColumnOccurrences(userDefinedColumns, { - name: assignOperation.args[0].parts.join('.'), - type: rightHandSideArgType /* fallback to number */, - location: assignOperation.args[0].location, - }); - } -} - -function addUserDefinedColumnFromExpression( - expressionOperation: ESQLFunction, - queryString: string, - userDefinedColumns: Map, - fields: Map -) { - if (!expressionOperation.text.includes(EDITOR_MARKER)) { - const expressionText = queryString.substring( - expressionOperation.location.min, - expressionOperation.location.max + 1 - ); - const expressionType = getExpressionType(expressionOperation, fields, userDefinedColumns); - addToUserDefinedColumnOccurrences(userDefinedColumns, { - name: expressionText, - type: expressionType, - location: expressionOperation.location, - }); - } -} - -export function collectUserDefinedColumns( - ast: ESQLAst, - fields: Map, - queryString: string -): Map { - const userDefinedColumns = new Map(); - - const visitor = new Visitor() - .on('visitLiteralExpression', (ctx) => { - // TODO - add these as userDefinedColumns - }) - .on('visitExpression', (_ctx) => {}) // required for the types :shrug: - .on('visitFunctionCallExpression', (ctx) => { - const node = ctx.node; - - if (node.subtype === 'binary-expression' && node.name === 'where') { - ctx.visitArgument(0, undefined); - return; - } - - if (node.name === 'as') { - const [oldArg, newArg] = ctx.node.args; - addToUserDefinedColumns(oldArg, newArg, fields, userDefinedColumns); - } else if (node.name === '=') { - addUserDefinedColumnFromAssignment(node, userDefinedColumns, fields); - } else { - addUserDefinedColumnFromExpression(node, queryString, userDefinedColumns, fields); - } - }) - .on('visitCommandOption', (ctx) => { - if (ctx.node.name === 'by') { - return [...ctx.visitArguments()]; - } else if (ctx.node.name === 'with') { - for (const assignFn of ctx.node.args) { - if (isFunctionExpression(assignFn)) { - const [newArg, oldArg] = assignFn?.args || []; - // TODO why is oldArg an array? - if (Array.isArray(oldArg)) { - addToUserDefinedColumns(oldArg[0], newArg, fields, userDefinedColumns); - } - } - } - } - }) - .on('visitCommand', (ctx) => { - const ret = []; - if (['row', 'eval', 'stats', 'inlinestats', 'ts', 'rename'].includes(ctx.node.name)) { - ret.push(...ctx.visitArgs()); - } - if (['stats', 'inlinestats', 'enrich'].includes(ctx.node.name)) { - // BY and WITH can contain userDefinedColumns - ret.push(...ctx.visitOptions()); - } - if (ctx.node.name === 'fork') { - ret.push(...ctx.visitSubQueries()); - } - return ret; - }) - .on('visitQuery', (ctx) => [...ctx.visitCommands()]); - - visitor.visitQuery(ast); - - return userDefinedColumns; -} diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/callbacks.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/callbacks.test.ts deleted file mode 100644 index 16898cc737e58..0000000000000 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/callbacks.test.ts +++ /dev/null @@ -1,64 +0,0 @@ -/* - * 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 { setup } from './helpers'; - -describe('FROM', () => { - test('does not load fields when validating only a single FROM, SHOW, ROW command', async () => { - const { validate, callbacks } = await setup(); - - await validate('FROM kib'); - await validate('FROM kibana_ecommerce METADATA _i'); - await validate('FROM kibana_ecommerce METADATA _id | '); - await validate('SHOW'); - await validate('ROW \t'); - - expect((callbacks.getColumnsFor as any).mock.calls.length).toBe(0); - }); - - test('loads fields with FROM source when commands after pipe present', async () => { - const { validate, callbacks } = await setup(); - - await validate('FROM kibana_ecommerce METADATA _id | eval'); - - expect((callbacks.getColumnsFor as any).mock.calls.length).toBe(1); - }); - - test('loads fields from JOIN index', async () => { - const { validate, callbacks } = await setup(); - - await validate('FROM index1 | LOOKUP JOIN index2 ON field1 | LIMIT 123'); - - expect((callbacks.getColumnsFor as any).mock.calls.length).toBe(1); - - const query = (callbacks.getColumnsFor as any).mock.calls[0][0].query as string; - - expect(query.includes('index1')).toBe(true); - expect(query.includes('index2')).toBe(true); - }); - - test('includes all "from" and "join" index for loading fields', async () => { - const { validate, callbacks } = await setup(); - - await validate( - 'FROM index1, index2, index3 | LOOKUP JOIN index4 ON field1 | KEEP abc | LOOKUP JOIN index5 ON field2 | LIMIT 123' - ); - - expect((callbacks.getColumnsFor as any).mock.calls.length).toBe(1); - - const query = (callbacks.getColumnsFor as any).mock.calls[0][0].query as string; - - expect(query.includes('index1')).toBe(true); - expect(query.includes('index2')).toBe(true); - expect(query.includes('index3')).toBe(true); - expect(query.includes('index4')).toBe(true); - expect(query.includes('index5')).toBe(true); - expect(query.includes('index6')).toBe(false); - }); -}); diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/fields_and_variables.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/fields_and_variables.test.ts index 9d62dcb95da83..4d99dc0d172b7 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/fields_and_variables.test.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/fields_and_variables.test.ts @@ -16,7 +16,7 @@ import { Location } from '@kbn/esql-ast/src/commands_registry/types'; import { setTestFunctions } from '@kbn/esql-ast/src/definitions/utils/test_functions'; import { setup } from './helpers'; -describe('field and userDefinedColumn escaping', () => { +describe('column escaping', () => { it('recognizes escaped fields', async () => { const { expectErrors } = await setup(); // command level @@ -39,7 +39,7 @@ describe('field and userDefinedColumn escaping', () => { ); }); - it('recognizes escaped userDefinedColumns', async () => { + it('recognizes escaped user-defined columns', async () => { const { expectErrors } = await setup(); // command level await expectErrors('ROW `var$iable` = 1 | EVAL `var$iable`', []); @@ -56,12 +56,12 @@ describe('field and userDefinedColumn escaping', () => { [] ); - // expression userDefinedColumn + // expression user-defined column await expectErrors('FROM index | EVAL doubleField + 20 | EVAL `doubleField + 20`', []); await expectErrors('ROW 21 + 20 | STATS AVG(`21 + 20`)', []); }); - it('recognizes userDefinedColumns with spaces and comments', async () => { + it('recognizes user-defined columns with spaces and comments', async () => { const { expectErrors } = await setup(); // command level await expectErrors( @@ -114,13 +114,13 @@ describe('field and userDefinedColumn escaping', () => { }); }); -describe('userDefinedColumn support', () => { - describe('userDefinedColumn data type detection', () => { +describe('user-defined column support', () => { + describe('user-defined column data type detection', () => { // most of these tests are aspirational (and skipped) because we don't have // a good way to compute the type of an expression yet. beforeAll(() => { setTestFunctions([ - // this test function is just used to test the type of the userDefinedColumn + // this test function is just used to test the type of the user-defined column { type: FunctionDefinitionTypes.SCALAR, description: 'Test function', @@ -131,7 +131,7 @@ describe('userDefinedColumn support', () => { ], }, // this test function is used to check that the correct return type is used - // when determining userDefinedColumn types + // when determining user-defined column types { type: FunctionDefinitionTypes.SCALAR, description: 'Test function', @@ -157,10 +157,9 @@ describe('userDefinedColumn support', () => { }); const expectType = (type: FunctionParameterType) => - `Argument of [test] must be [cartesian_point], found value [var] type [${type}]`; + getNoValidCallSignatureError('test', [type]); - // @todo unskip after https://github.com/elastic/kibana/issues/195682 - test.skip('literals', async () => { + test('literals', async () => { const { expectErrors } = await setup(); // literal assignment await expectErrors('FROM index | EVAL var = 1, TEST(var)', [expectType('integer')]); @@ -176,16 +175,14 @@ describe('userDefinedColumn support', () => { ]); }); - // @todo unskip after https://github.com/elastic/kibana/issues/195682 - test.skip('userDefinedColumns', async () => { + test('user-defined columns', async () => { const { expectErrors } = await setup(); await expectErrors('FROM index | EVAL var = textField, col2 = var, TEST(col2)', [ - `Argument of [test] must be [cartesian_point], found value [col2] type [text]`, + getNoValidCallSignatureError('test', ['text']), ]); }); - // @todo unskip after https://github.com/elastic/kibana/issues/195682 - test.skip('inline casting', async () => { + test('inline casting', async () => { const { expectErrors } = await setup(); // inline cast assignment await expectErrors('FROM index | EVAL var = doubleField::long, TEST(var)', [ @@ -197,8 +194,7 @@ describe('userDefinedColumn support', () => { ]); }); - // @todo unskip after https://github.com/elastic/kibana/issues/195682 - test.skip('function results', async () => { + test('function results', async () => { const { expectErrors } = await setup(); // function assignment await expectErrors('FROM index | EVAL var = RETURN_VALUE(doubleField), TEST(var)', [ diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/functions.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/functions.test.ts index 87b2f7e8971fe..4da35fc502e4b 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/functions.test.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/functions.test.ts @@ -391,6 +391,24 @@ describe('function validation', () => { ]); }); }); + + it('skips column validation for left assignment arg', async () => { + const { expectErrors } = await setup(); + + await expectErrors('FROM a_index | EVAL lolz = 2', []); + await expectErrors('FROM a_index | EVAL lolz = nonexistent', [ + 'Unknown column "nonexistent"', + ]); + }); + + it('skips column validation for right arg to AS', async () => { + const { expectErrors } = await setup(); + + await expectErrors('FROM a_index | RENAME keywordField AS lolz', []); + await expectErrors('FROM a_index | RENAME nonexistent AS lolz', [ + 'Unknown column "nonexistent"', + ]); + }); }); }); diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/esql_validation_meta_tests.json b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/esql_validation_meta_tests.json index 74209aa860b54..a21932393cefd 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/esql_validation_meta_tests.json +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/esql_validation_meta_tests.json @@ -10,113 +10,140 @@ "fields": [ { "name": "booleanField", - "type": "boolean" + "type": "boolean", + "userDefined": false }, { "name": "dateField", - "type": "date" + "type": "date", + "userDefined": false }, { "name": "doubleField", - "type": "double" + "type": "double", + "userDefined": false }, { "name": "ipField", - "type": "ip" + "type": "ip", + "userDefined": false }, { "name": "keywordField", - "type": "keyword" + "type": "keyword", + "userDefined": false }, { "name": "integerField", - "type": "integer" + "type": "integer", + "userDefined": false }, { "name": "longField", - "type": "long" + "type": "long", + "userDefined": false }, { "name": "textField", - "type": "text" + "type": "text", + "userDefined": false }, { "name": "unsignedLongField", - "type": "unsigned_long" + "type": "unsigned_long", + "userDefined": false }, { "name": "versionField", - "type": "version" + "type": "version", + "userDefined": false }, { "name": "cartesianPointField", - "type": "cartesian_point" + "type": "cartesian_point", + "userDefined": false }, { "name": "cartesianShapeField", - "type": "cartesian_shape" + "type": "cartesian_shape", + "userDefined": false }, { "name": "geoPointField", - "type": "geo_point" + "type": "geo_point", + "userDefined": false }, { "name": "geoShapeField", - "type": "geo_shape" + "type": "geo_shape", + "userDefined": false }, { "name": "counterIntegerField", - "type": "counter_integer" + "type": "counter_integer", + "userDefined": false }, { "name": "counterLongField", - "type": "counter_long" + "type": "counter_long", + "userDefined": false }, { "name": "counterDoubleField", - "type": "counter_double" + "type": "counter_double", + "userDefined": false }, { "name": "unsupportedField", - "type": "unsupported" + "type": "unsupported", + "userDefined": false }, { "name": "dateNanosField", - "type": "date_nanos" + "type": "date_nanos", + "userDefined": false }, { "name": "functionNamedParametersField", - "type": "function_named_parameters" + "type": "function_named_parameters", + "userDefined": false }, { "name": "any#Char$Field", - "type": "double" + "type": "double", + "userDefined": false }, { "name": "kubernetes.something.something", - "type": "double" + "type": "double", + "userDefined": false }, { "name": "@timestamp", - "type": "date" + "type": "date", + "userDefined": false }, { "name": "otherStringField", - "type": "keyword" + "type": "keyword", + "userDefined": false } ], "enrichFields": [ { "name": "otherField", - "type": "text" + "type": "text", + "userDefined": false }, { "name": "yetAnotherField", - "type": "double" + "type": "double", + "userDefined": false }, { "name": "otherStringField", - "type": "keyword" + "type": "keyword", + "userDefined": false } ], "policies": [ @@ -146,7 +173,8 @@ "unsupported_field": [ { "name": "unsupported_field", - "type": "unsupported" + "type": "unsupported", + "userDefined": false } ], "testCases": [ diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts index c04e8bfbf76cd..7296cfbeb6aa9 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts @@ -7,67 +7,9 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ -import { - type ESQLAst, - type ESQLAstQueryExpression, - type ESQLColumn, - type ESQLSource, - type ESQLCommand, - type FunctionDefinition, - Walker, -} from '@kbn/esql-ast'; -import { mutate, synth } from '@kbn/esql-ast'; -import type { ESQLPolicy } from '@kbn/esql-ast/src/commands_registry/types'; - -export function buildQueryForFieldsFromSource(queryString: string, ast: ESQLAst) { - const firstCommand = ast[0]; - if (!firstCommand) return ''; - - const sources: ESQLSource[] = []; - const metadataFields: ESQLColumn[] = []; - - if (firstCommand.name === 'ts') { - const timeseries = firstCommand as ESQLCommand<'ts'>; - const tsSources = timeseries.args as ESQLSource[]; - - sources.push(...tsSources); - } else if (firstCommand.name === 'from') { - const fromSources = mutate.commands.from.sources.list(firstCommand as any); - const fromMetadataColumns = [...mutate.commands.from.metadata.list(firstCommand as any)].map( - ([column]) => column - ); - - sources.push(...fromSources); - if (fromMetadataColumns.length) metadataFields.push(...fromMetadataColumns); - } - - const joinSummary = mutate.commands.join.summarize({ - type: 'query', - commands: ast, - } as ESQLAstQueryExpression); - const joinIndices = joinSummary.map(({ target: { index } }) => index); - - if (joinIndices.length > 0) { - sources.push(...joinIndices); - } - - if (sources.length === 0) { - return queryString.substring(0, firstCommand.location.max + 1); - } - - const from = - metadataFields.length > 0 - ? synth.cmd`FROM ${sources} METADATA ${metadataFields}` - : synth.cmd`FROM ${sources}`; - - return from.toString(); -} - -export function buildQueryForFieldsInPolicies(policies: ESQLPolicy[]) { - return `from ${policies - .flatMap(({ sourceIndices }) => sourceIndices) - .join(', ')} | keep ${policies.flatMap(({ enrichFields }) => enrichFields).join(', ')}`; -} +import type { ESQLAstQueryExpression } from '@kbn/esql-ast'; +import { type ESQLCommand, type FunctionDefinition, Walker, Builder } from '@kbn/esql-ast'; +import { expandEvals } from '../shared/expand_evals'; /** * Returns the maximum and minimum number of parameters allowed by a function @@ -96,3 +38,49 @@ export function getMaxMinNumberOfParams(definition: FunctionDefinition) { */ export const getEnrichCommands = (commands: ESQLCommand[]): ESQLCommand[] => Walker.matchAll(commands, { type: 'command', name: 'enrich' }) as ESQLCommand[]; + +/** + * Returns a list of subqueries to validate + * @param rootCommands + */ +export function getSubqueriesToValidate(rootCommands: ESQLCommand[]) { + const subsequences = []; + const expandedCommands = expandEvals(rootCommands); + for (let i = 0; i < expandedCommands.length; i++) { + const command = expandedCommands[i]; + + // every command within FORK's branches is its own subquery to be validated + if (command.name.toLowerCase() === 'fork') { + const branchSubqueries = getForkBranchSubqueries(command as ESQLCommand<'fork'>); + for (const subquery of branchSubqueries) { + subsequences.push([...expandedCommands.slice(0, i), ...subquery]); + } + } + + subsequences.push(expandedCommands.slice(0, i + 1)); + } + + return subsequences.map((subsequence) => Builder.expression.query(subsequence)); +} + +/** + * Expands a FORK command into flat subqueries for each command in each branch. + * + * E.g. FORK (EVAL 1 | LIMIT 10) (RENAME foo AS bar | DROP lolz) + * + * becomes [`EVAL 1`, `EVAL 1 | LIMIT 10`, `RENAME foo AS bar`, `RENAME foo AS bar | DROP lolz`] + * + * @param command a FORK command + * @returns an array of expanded subqueries + */ +function getForkBranchSubqueries(command: ESQLCommand<'fork'>): ESQLCommand[][] { + const expanded: ESQLCommand[][] = []; + const branches = command.args as ESQLAstQueryExpression[]; + for (let j = 0; j < branches.length; j++) { + for (let k = 0; k < branches[j].commands.length; k++) { + const partialQuery = branches[j].commands.slice(0, k + 1); + expanded.push(partialQuery); + } + } + return expanded; +} diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts index 006293ed895af..777984ca23a3f 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/resources.ts @@ -7,46 +7,12 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ -import { isSource, type ESQLCommand } from '@kbn/esql-ast'; -import type { ESQLFieldWithMetadata, ESQLPolicy } from '@kbn/esql-ast/src/commands_registry/types'; -import { createMapFromList, nonNullable } from '../shared/helpers'; -import { - getFieldsByTypeHelper, - getPolicyHelper, - getSourcesHelper, -} from '../shared/resources_helpers'; +import { type ESQLCommand } from '@kbn/esql-ast'; +import type { ESQLPolicy } from '@kbn/esql-ast/src/commands_registry/types'; +import { createMapFromList } from '../shared/helpers'; +import { getPolicyHelper, getSourcesHelper } from '../shared/resources_helpers'; import type { ESQLCallbacks } from '../shared/types'; -import { - buildQueryForFieldsFromSource, - buildQueryForFieldsInPolicies, - getEnrichCommands, -} from './helpers'; - -export async function retrieveFields( - queryString: string, - commands: ESQLCommand[], - callbacks?: ESQLCallbacks -): Promise> { - if (!callbacks || commands.length < 1) { - return new Map(); - } - // Do not fetch fields, if query has only one source command and that command - // does not require fields. - if (commands.length === 1) { - switch (commands[0].name) { - case 'from': - case 'show': - case 'row': { - return new Map(); - } - } - } - if (commands[0].name === 'row') { - return new Map(); - } - const customQuery = buildQueryForFieldsFromSource(queryString, commands); - return await getFieldsByTypeHelper(customQuery, callbacks).getFieldsMap(); -} +import { getEnrichCommands } from './helpers'; export async function retrievePolicies( commands: ESQLCommand[], @@ -74,28 +40,3 @@ export async function retrieveSources( const sources = await getSourcesHelper(callbacks)(); return new Set(sources.map(({ name }) => name)); } - -export async function retrievePoliciesFields( - commands: ESQLCommand[], - policies: Map, - callbacks?: ESQLCallbacks -): Promise> { - if (!callbacks) { - return new Map(); - } - const enrichCommands = getEnrichCommands(commands); - if (!enrichCommands.length) { - return new Map(); - } - const policyNames = enrichCommands - .map(({ args }) => (isSource(args[0]) ? args[0].name : undefined)) - .filter(nonNullable); - if (!policyNames.every((name) => policies.has(name))) { - return new Map(); - } - - const customQuery = buildQueryForFieldsInPolicies( - policyNames.map((name) => policies.get(name)) as ESQLPolicy[] - ); - return await getFieldsByTypeHelper(customQuery, callbacks).getFieldsMap(); -} diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/types.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/types.ts index 5b690fc68135c..cfa50bfddcc1f 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/types.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/types.ts @@ -8,15 +8,14 @@ */ import type { ESQLMessage } from '@kbn/esql-ast'; -import type { ESQLFieldWithMetadata } from '@kbn/esql-ast/src/commands_registry/types'; -import type { ESQLPolicy, ESQLUserDefinedColumn } from '@kbn/esql-ast/src/commands_registry/types'; +import type { ESQLColumnData } from '@kbn/esql-ast/src/commands_registry/types'; +import type { ESQLPolicy } from '@kbn/esql-ast/src/commands_registry/types'; import type { IndexAutocompleteItem } from '@kbn/esql-types'; import type { EditorError } from '../types'; export interface ReferenceMaps { sources: Set; - userDefinedColumns: Map; - fields: Map; + columns: Map; policies: Map; query: string; joinIndices: IndexAutocompleteItem[]; diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.test.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.test.ts index 3f581967a9e4f..a3ceb6df4fda2 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.test.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.test.ts @@ -178,9 +178,11 @@ describe('validation logic', () => { JSON.stringify( { indexes, - fields: fields.concat([{ name: policies[0].matchField, type: 'keyword' }]), + fields: fields.concat([ + { name: policies[0].matchField, type: 'keyword', userDefined: false }, + ]), enrichFields: enrichFields.concat([ - { name: policies[0].matchField, type: 'keyword' }, + { name: policies[0].matchField, type: 'keyword', userDefined: false }, ]), policies, unsupported_field, @@ -599,17 +601,6 @@ describe('validation logic', () => { expect(callbackMocks.getColumnsFor).toHaveBeenCalledTimes(0); }); - it('should call fields callbacks also for show command', async () => { - const callbackMocks = getCallbackMocks(); - await validateQuery(`show info | keep name`, undefined, callbackMocks); - expect(callbackMocks.getSources).not.toHaveBeenCalled(); - expect(callbackMocks.getPolicies).not.toHaveBeenCalled(); - expect(callbackMocks.getColumnsFor).toHaveBeenCalledTimes(1); - expect(callbackMocks.getColumnsFor).toHaveBeenLastCalledWith({ - query: 'show info', - }); - }); - it(`should fetch additional fields if an enrich command is found`, async () => { const callbackMocks = getCallbackMocks(); await validateQuery( diff --git a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts index 8f4ec700ac4b4..38d78fc552cd6 100644 --- a/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts +++ b/src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/validation.ts @@ -8,23 +8,19 @@ */ import type { ESQLAst, ESQLCommand, ESQLMessage, ErrorTypes } from '@kbn/esql-ast'; -import { EsqlQuery, walk, esqlCommandRegistry } from '@kbn/esql-ast'; -import { getMessageFromId } from '@kbn/esql-ast/src/definitions/utils'; +import { EsqlQuery, esqlCommandRegistry, walk } from '@kbn/esql-ast'; import type { ESQLFieldWithMetadata, ICommandCallbacks, } from '@kbn/esql-ast/src/commands_registry/types'; +import { getMessageFromId } from '@kbn/esql-ast/src/definitions/utils'; import type { LicenseType } from '@kbn/licensing-types'; +import { getColumnsByTypeHelper } from '../shared/resources_helpers'; import type { ESQLCallbacks } from '../shared/types'; -import { collectUserDefinedColumns } from '../shared/user_defined_columns'; -import { - retrieveFields, - retrievePolicies, - retrievePoliciesFields, - retrieveSources, -} from './resources'; +import { retrievePolicies, retrieveSources } from './resources'; import type { ReferenceMaps, ValidationOptions, ValidationResult } from './types'; +import { getSubqueriesToValidate } from './helpers'; /** * ES|QL validation public API @@ -109,45 +105,60 @@ async function validateAst( const parsingResult = EsqlQuery.fromSrc(queryString); const rootCommands = parsingResult.ast.commands; - const [sources, availableFields, availablePolicies, joinIndices] = await Promise.all([ + const [sources, availablePolicies, joinIndices] = await Promise.all([ // retrieve the list of available sources retrieveSources(rootCommands, callbacks), - // retrieve available fields (if a source command has been defined) - retrieveFields(queryString, rootCommands, callbacks), // retrieve available policies (if an enrich command has been defined) retrievePolicies(rootCommands, callbacks), // retrieve indices for join command callbacks?.getJoinIndices?.(), ]); - if (availablePolicies.size) { - const fieldsFromPoliciesMap = await retrievePoliciesFields( - rootCommands, - availablePolicies, - callbacks - ); - fieldsFromPoliciesMap.forEach((value, key) => availableFields.set(key, value)); - } - - const userDefinedColumns = collectUserDefinedColumns(rootCommands, availableFields, queryString); - messages.push(...validateUnsupportedTypeFields(availableFields, rootCommands)); + const sourceQuery = queryString.split('|')[0]; + const sourceFields = await getColumnsByTypeHelper( + EsqlQuery.fromSrc(sourceQuery).ast, + sourceQuery, + callbacks + ).getColumnMap(); - const references: ReferenceMaps = { - sources, - fields: availableFields, - policies: availablePolicies, - userDefinedColumns, - query: queryString, - joinIndices: joinIndices?.indices || [], - }; + messages.push( + ...validateUnsupportedTypeFields( + sourceFields as Map, + rootCommands + ) + ); const license = await callbacks?.getLicense?.(); const hasMinimumLicenseRequired = license?.hasAtLeast; - for (const [_, command] of rootCommands.entries()) { - const commandMessages = validateCommand(command, references, rootCommands, { - ...callbacks, - hasMinimumLicenseRequired, - }); + + /** + * Even though we are validating single commands, we work with subqueries. + * + * The reason is that building the list of columns available in each command requires + * the full command subsequence that precedes that command. + */ + const subqueries = getSubqueriesToValidate(rootCommands); + for (const subquery of subqueries) { + const { getColumnMap } = getColumnsByTypeHelper(subquery, queryString, callbacks); + const availableColumns = await getColumnMap(); + + const references: ReferenceMaps = { + sources, + columns: availableColumns, + policies: availablePolicies, + query: queryString, + joinIndices: joinIndices?.indices || [], + }; + + const commandMessages = validateCommand( + subquery.commands[subquery.commands.length - 1], + references, + rootCommands, + { + ...callbacks, + hasMinimumLicenseRequired, + } + ); messages.push(...commandMessages); } @@ -206,9 +217,8 @@ function validateCommand( } const context = { - fields: references.fields, + columns: references.columns, policies: references.policies, - userDefinedColumns: references.userDefinedColumns, sources: [...references.sources].map((source) => ({ name: source, })), diff --git a/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/__tests__/fixtures.ts b/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/__tests__/fixtures.ts index 52ce9bfe8d22a..b9e8cd967c68a 100644 --- a/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/__tests__/fixtures.ts +++ b/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/__tests__/fixtures.ts @@ -18,9 +18,10 @@ const fields: Array = [ ...types.map((type) => ({ name: `${type}Field`, type, + userDefined: false as false, })), - { name: 'any#Char$Field', type: 'double', suggestedAs: '`any#Char$Field`' }, - { name: 'kubernetes.something.something', type: 'double' }, + { name: 'any#Char$Field', type: 'double', suggestedAs: '`any#Char$Field`', userDefined: false }, + { name: 'kubernetes.something.something', type: 'double', userDefined: false }, ]; const indexes = ( diff --git a/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts b/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts index c849bdae0dfad..fa81edfe2df7b 100644 --- a/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts +++ b/src/platform/packages/shared/kbn-monaco/src/languages/esql/lib/hover/hover.ts @@ -12,7 +12,6 @@ import { ENRICH_MODES, modeDescription, } from '@kbn/esql-ast/src/commands_registry/commands/enrich/util'; -import type { ESQLFieldWithMetadata } from '@kbn/esql-ast/src/commands_registry/types'; import { getFunctionDefinition, getFunctionSignatures, @@ -25,9 +24,8 @@ import { type ESQLSingleAstItem, type ESQLSource, } from '@kbn/esql-ast/src/types'; -import { collectUserDefinedColumns, type ESQLCallbacks } from '@kbn/esql-validation-autocomplete'; -import { getFieldsByTypeRetriever } from '@kbn/esql-validation-autocomplete/src/autocomplete/autocomplete'; -import { getQueryForFields } from '@kbn/esql-validation-autocomplete/src/autocomplete/helper'; +import { type ESQLCallbacks } from '@kbn/esql-validation-autocomplete'; +import { getColumnsByTypeRetriever } from '@kbn/esql-validation-autocomplete/src/autocomplete/autocomplete'; import { getPolicyHelper } from '@kbn/esql-validation-autocomplete/src/shared/resources_helpers'; import { i18n } from '@kbn/i18n'; import type { monaco } from '../../../../monaco_imports'; @@ -187,20 +185,17 @@ async function getHintForFunctionArg( offset: number, resourceRetriever?: ESQLCallbacks ) { - const queryForFields = getQueryForFields(query, root); - const { getFieldsMap } = getFieldsByTypeRetriever(queryForFields, resourceRetriever); + const { getColumnMap } = getColumnsByTypeRetriever(root, query, resourceRetriever); const fnDefinition = getFunctionDefinition(fnNode.name); // early exit on no hit if (!fnDefinition) { return []; } - const fieldsMap: Map = await getFieldsMap(); - const anyUserDefinedColumns = collectUserDefinedColumns(root.commands, fieldsMap, query); + const columnsMap = await getColumnMap(); const references = { - fields: fieldsMap, - userDefinedColumns: anyUserDefinedColumns, + columns: columnsMap, }; const { typesToSuggestNext, enrichedArgs } = getValidSignaturesAndTypesToSuggestNext( diff --git a/src/platform/plugins/shared/console/public/application/containers/editor/monaco_editor.tsx b/src/platform/plugins/shared/console/public/application/containers/editor/monaco_editor.tsx index 2cebacd4dc669..7470122bf9757 100644 --- a/src/platform/plugins/shared/console/public/application/containers/editor/monaco_editor.tsx +++ b/src/platform/plugins/shared/console/public/application/containers/editor/monaco_editor.tsx @@ -158,6 +158,7 @@ export const MonacoEditor = ({ localStorageValue, value, setValue }: EditorProps name: c.name, type: c.meta.esType as FieldType, hasConflict: c.meta.type === KBN_FIELD_TYPES.CONFLICT, + userDefined: false, }; }) || [] );