[ES|QL] fix column name detection

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/expand_evals.ts

@@ -26,11 +26,15 @@ export function expandEvals(commands: ESQLCommand[]): ESQLCommand[] {
     if (command.name.toLowerCase() === 'eval') {
       for (const arg of command.args) {
         expanded.push(
-          Builder.command({
-            name: 'eval',
-            args: [arg],
-            location: command.location,
-          })
+          Builder.command(
+            {
+              name: 'eval',
+              args: [arg],
+            },
+            {
+              location: command.location,
+            }
+          )
         );
       }
     } else {

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.test.ts

@@ -7,7 +7,7 @@
  * License v3.0 only", or the "Server Side Public License, v 1".
  */
 
-import { parse, EDITOR_MARKER } from '@kbn/esql-ast';
+import { parse, EDITOR_MARKER, BasicPrettyPrinter } from '@kbn/esql-ast';
 import { getQueryForFields } from './resources_helpers';
 
 describe('getQueryForFields', () => {
@@ -26,7 +26,7 @@ describe('getQueryForFields', () => {
 
     const result = getQueryForFields(query, root);
 
-    expect(result).toEqual(expected);
+    expect(BasicPrettyPrinter.print(result)).toEqual(expected);
   };
 
   it('should return everything up till the last command', () => {

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/shared/resources_helpers.ts

@@ -8,40 +8,20 @@
  */
 
 import type { ESQLAstQueryExpression } from '@kbn/esql-ast';
-import { BasicPrettyPrinter, Builder, EDITOR_MARKER, EsqlQuery } from '@kbn/esql-ast';
+import { BasicPrettyPrinter, Builder, EDITOR_MARKER } from '@kbn/esql-ast';
 import type {
   ESQLColumnData,
   ESQLFieldWithMetadata,
   ESQLPolicy,
 } from '@kbn/esql-ast/src/commands_registry/types';
-import type { ESQLCallbacks } from './types';
-import { getFieldsFromES, getCurrentQueryAvailableColumns } from './helpers';
 import { expandEvals } from './expand_evals';
+import { getCurrentQueryAvailableColumns, getFieldsFromES } from './helpers';
+import type { ESQLCallbacks } from './types';
 
 export const NOT_SUGGESTED_TYPES = ['unsupported'];
 
 const cache = new Map<string, ESQLColumnData[]>();
 
-// Function to check if a key exists in the cache, ignoring case
-function checkCacheInsensitive(keyToCheck: string) {
-  for (const key of cache.keys()) {
-    if (key.toLowerCase() === keyToCheck.toLowerCase()) {
-      return true; // Or return the value associated with the key if needed: return cache.get(key);
-    }
-  }
-  return false;
-}
-
-// Function to get a value from the cache, ignoring case
-function getValueInsensitive(keyToCheck: string) {
-  for (const key of cache.keys()) {
-    if (key.toLowerCase() === keyToCheck.toLowerCase()) {
-      return cache.get(key);
-    }
-  }
-  return undefined;
-}
-
 /**
  * Given a query, this function will compute the available fields and cache them
  * for the next time the same query is used.
@@ -53,18 +33,9 @@ async function cacheColumnsForQuery(
   getPolicies: () => Promise<Map<string, ESQLPolicy>>,
   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 cacheKey = originalQueryText.slice(query.location.min, query.location.max + 1);
 
-  const existsInCache = checkCacheInsensitive(cacheKey);
+  const existsInCache = cache.has(cacheKey);
   if (existsInCache) {
     // this is already in the cache
     return;
@@ -74,7 +45,7 @@ async function cacheColumnsForQuery(
     ...query,
     commands: query.commands.slice(0, -1),
   });
-  const fieldsAvailableAfterPreviousCommand = getValueInsensitive(queryBeforeCurrentCommand) ?? [];
+  const fieldsAvailableAfterPreviousCommand = cache.get(queryBeforeCurrentCommand) ?? [];
 
   const availableFields = await getCurrentQueryAvailableColumns(
     query.commands,
@@ -101,16 +72,18 @@ export function getColumnsByTypeHelper(
   originalQueryText: string,
   resourceRetriever?: ESQLCallbacks
 ) {
-  const queryForFields = getQueryForFields(originalQueryText, query);
-  const root = EsqlQuery.fromSrc(queryForFields).ast;
+  const root = getQueryForFields(originalQueryText, query);
+
+  // IMPORTANT: cache key can't be case-insensitive because column names are case-sensitive
+  const cacheKey = originalQueryText.slice(root.location.min, root.location.max + 1);
 
   const cacheColumns = async () => {
-    if (!queryForFields) {
+    if (!cacheKey) {
       return;
     }
 
     const getFields = async (queryToES: string) => {
-      const cached = getValueInsensitive(queryToES);
+      const cached = cache.get(queryToES);
       if (cached) {
         return cached as ESQLFieldWithMetadata[];
       }
@@ -121,7 +94,11 @@ export function getColumnsByTypeHelper(
 
     const subqueries = [];
     for (let i = 0; i < root.commands.length; i++) {
-      subqueries.push(Builder.expression.query(root.commands.slice(0, i + 1)));
+      subqueries.push(
+        Builder.expression.query(root.commands.slice(0, i + 1), {
+          location: { min: 0, max: root.commands[i].location.max },
+        })
+      );
     }
 
     const getPolicies = async () => {
@@ -142,7 +119,7 @@ export function getColumnsByTypeHelper(
     ): Promise<ESQLColumnData[]> => {
       const types = Array.isArray(expectedType) ? expectedType : [expectedType];
       await cacheColumns();
-      const cachedFields = getValueInsensitive(queryForFields);
+      const cachedFields = cache.get(cacheKey);
       return (
         cachedFields?.filter(({ name, type }) => {
           const ts = Array.isArray(type) ? type : [type];
@@ -157,7 +134,7 @@ export function getColumnsByTypeHelper(
     },
     getColumnMap: async (): Promise<Map<string, ESQLColumnData>> => {
       await cacheColumns();
-      const cachedFields = getValueInsensitive(queryForFields);
+      const cachedFields = cache.get(cacheKey);
       const cacheCopy = new Map<string, ESQLColumnData>();
       cachedFields?.forEach((field) => cacheCopy.set(field.name, field));
       return cacheCopy;
@@ -195,11 +172,16 @@ export function getSourcesHelper(resourceRetriever?: ESQLCallbacks) {
  * 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.
  *
+ * IMPORTANT: the AST nodes in the new query still reference locations in the original query text
+ *
  * @param queryString The original query string
  * @param commands
  * @returns
  */
-export function getQueryForFields(queryString: string, root: ESQLAstQueryExpression): string {
+export function getQueryForFields(
+  queryString: string,
+  root: ESQLAstQueryExpression
+): ESQLAstQueryExpression {
   const commands = root.commands;
   const lastCommand = commands[commands.length - 1];
   if (lastCommand && lastCommand.name === 'fork' && lastCommand.args.length > 0) {
@@ -215,7 +197,8 @@ export function getQueryForFields(queryString: string, root: ESQLAstQueryExpress
      */
     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 });
+    const newLocation = { min: 0, max: newCommands[newCommands.length - 1]?.location.max ?? 0 };
+    return { ...root, commands: newCommands, location: newLocation };
   }
 
   if (lastCommand && lastCommand.name === 'eval') {
@@ -235,7 +218,8 @@ export function getQueryForFields(queryString: string, root: ESQLAstQueryExpress
        */
       const expanded = expandEvals(commands);
       const newCommands = expanded.slice(0, endsWithComma ? undefined : -1);
-      return BasicPrettyPrinter.print({ ...root, commands: newCommands });
+      const newLocation = { min: 0, max: newCommands[newCommands.length - 1]?.location.max ?? 0 };
+      return { ...root, commands: newCommands, location: newLocation };
     }
   }
 
@@ -244,8 +228,15 @@ export function getQueryForFields(queryString: string, root: ESQLAstQueryExpress
 
 function buildQueryUntilPreviousCommand(root: ESQLAstQueryExpression) {
   if (root.commands.length === 1) {
-    return BasicPrettyPrinter.print({ ...root.commands[0] });
+    return { ...root, commands: [root.commands[0]] };
   } else {
-    return BasicPrettyPrinter.print({ ...root, commands: root.commands.slice(0, -1) });
+    const newCommands = root.commands.slice(0, -1);
+    const newLocation = { min: 0, max: newCommands[newCommands.length - 1]?.location.max ?? 0 };
+
+    return {
+      ...root,
+      commands: newCommands,
+      location: newLocation,
+    };
   }
 }

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/__tests__/column_caching.test.ts

@@ -0,0 +1,29 @@
+/*
+ * 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('column caching and casing', () => {
+  it('expression columns are cached case-sensitively', async () => {
+    const { expectErrors } = await setup();
+    await expectErrors('FROM index | EVAL TrIm("") | EVAL `TrIm("")`', []);
+    await expectErrors('FROM index | EVAL TrIm("") | EVAL `TRIM("")`', [
+      'Unknown column "TRIM("")"',
+    ]);
+  });
+
+  it('case changes bust the cache', async () => {
+    const { expectErrors } = await setup();
+    // first populate the cache with capitalized TRIM query
+    await expectErrors('FROM index | EVAL ABS(1) | EVAL `ABS(1)`', []);
+
+    // change the case of the TRIM to trim, which should bust the cache and create an error
+    await expectErrors('FROM index | EVAL abs(1) | EVAL `ABS(1)`', ['Unknown column "ABS(1)"']);
+  });
+});

src/platform/packages/shared/kbn-esql-validation-autocomplete/src/validation/helpers.ts

@@ -60,7 +60,11 @@ export function getSubqueriesToValidate(rootCommands: ESQLCommand[]) {
     subsequences.push(expandedCommands.slice(0, i + 1));
   }
 
-  return subsequences.map((subsequence) => Builder.expression.query(subsequence));
+  return subsequences.map((subsequence) =>
+    Builder.expression.query(subsequence, {
+      location: { min: 0, max: subsequence[subsequence.length - 1].location.max },
+    })
+  );
 }
 
 /**