[ES|QL] Update other AST Tools (elastic#233130)

## Summary
part of elastic#232885

- Remove inferenceId from ESQLAstRerankCommand type
- Update Pretty printer, visitor, walker, mutate ....

note: most of these updates are tests

Author: bartoval

src/platform/packages/shared/kbn-esql-ast/src/mutate/commands/rerank/index.test.ts

@@ -10,16 +10,13 @@
 import { BasicPrettyPrinter } from '../../../pretty_print';
 import * as commands from '..';
 import { EsqlQuery } from '../../../query';
+import type { ESQLAstItem, ESQLCommandOption, ESQLMap } from '../../../types';
 
-/**
- * @todo Tests skipped, while RERANK command grammar is being stabilized. We will
- * get back to it after 9.1 release.
- */
-describe.skip('commands.rerank', () => {
+describe('commands.rerank', () => {
   describe('.list()', () => {
     it('lists the only "RERANK" commands', () => {
       const src =
-        'FROM index | LIMIT 1 | RERANK "star wars" ON title, overview=SUBSTRING(overview, 0, 100), actors WITH rerankerInferenceId | LIMIT 2';
+        'FROM index | LIMIT 1 | RERANK "star wars" ON title, overview=SUBSTRING(overview, 0, 100), actors WITH  { "inference_id": "model_id" } | LIMIT 2';
       const query = EsqlQuery.fromSrc(src);
 
       const nodes = [...commands.rerank.list(query.ast)];
@@ -35,22 +32,74 @@ describe.skip('commands.rerank', () => {
 
   describe('.setQuery()', () => {
     it('can change query', () => {
-      const src = 'FROM index | LIMIT 1 | RERANK "star wars" ON field WITH id | LIMIT 2';
+      const src =
+        'FROM index | RERANK "star wars" ON field WITH { "inference_id": "model_id" } | LIMIT 2';
       const query = EsqlQuery.fromSrc(src);
 
       const cmd = [...commands.rerank.list(query.ast)][0];
       commands.rerank.setQuery(cmd, 'new query');
 
       expect(BasicPrettyPrinter.expression(cmd.query)).toBe('"new query"');
-      expect(query.print()).toBe(
-        'FROM index | LIMIT 1 | RERANK "new query" ON field WITH id | LIMIT 2'
+      expect(query.print({ wrap: Infinity })).toBe(
+        'FROM index | RERANK "new query" ON field WITH {"inference_id": "model_id"} | LIMIT 2'
       );
     });
   });
 
+  it('can change query on a command with a target assignment', () => {
+    const src =
+      'FROM index | RERANK my_target = "star wars" ON field WITH { "inference_id": "model_id" }';
+    const query = EsqlQuery.fromSrc(src);
+
+    const cmd = [...commands.rerank.list(query.ast)][0];
+    commands.rerank.setQuery(cmd, 'new query');
+
+    expect(BasicPrettyPrinter.expression(cmd.query)).toBe('"new query"');
+    expect(query.print({ wrap: Infinity })).toBe(
+      'FROM index | RERANK my_target = "new query" ON field WITH {"inference_id": "model_id"}'
+    );
+  });
+});
+
+describe('.setTargetField()', () => {
+  it('can add a target field to a simple command', () => {
+    const src = 'FROM index | RERANK "star wars" ON field';
+    const query = EsqlQuery.fromSrc(src);
+
+    const cmd = [...commands.rerank.list(query.ast)][0];
+    commands.rerank.setTargetField(cmd, 'my_target');
+
+    expect(query.print({ wrap: Infinity })).toBe(
+      'FROM index | RERANK my_target = "star wars" ON field'
+    );
+  });
+
+  it('can change an existing target field', () => {
+    const src = 'FROM index | RERANK old_target = "star wars" ON field';
+    const query = EsqlQuery.fromSrc(src);
+
+    const cmd = [...commands.rerank.list(query.ast)][0];
+    commands.rerank.setTargetField(cmd, 'new_target');
+
+    expect(query.print({ wrap: Infinity })).toBe(
+      'FROM index | RERANK new_target = "star wars" ON field'
+    );
+  });
+
+  it('can remove an existing target field', () => {
+    const src = 'FROM index | RERANK my_target = "star wars" ON field';
+    const query = EsqlQuery.fromSrc(src);
+
+    const cmd = [...commands.rerank.list(query.ast)][0];
+    commands.rerank.setTargetField(cmd, null);
+
+    expect(query.print({ wrap: Infinity })).toBe('FROM index | RERANK "star wars" ON field');
+  });
+
   describe('.setFields()', () => {
     it('can change query', () => {
-      const src = 'FROM index | LIMIT 1 | RERANK "star wars" ON field WITH id | LIMIT 2';
+      const src =
+        'FROM index | RERANK "star wars" ON field WITH { "inference_id": "model_id" } | LIMIT 2';
       const query = EsqlQuery.fromSrc(src);
 
       const cmd = [...commands.rerank.list(query.ast)][0];
@@ -61,23 +110,44 @@ describe.skip('commands.rerank', () => {
         'b',
         '@timestamp',
       ]);
-      expect(query.print()).toBe(
-        'FROM index | LIMIT 1 | RERANK "star wars" ON a, b, @timestamp WITH id | LIMIT 2'
+      expect(query.print({ wrap: Infinity })).toBe(
+        'FROM index | RERANK "star wars" ON a, b, @timestamp WITH {"inference_id": "model_id"} | LIMIT 2'
+      );
+    });
+
+    it('should throw error when ON option is missing', () => {
+      const src = 'FROM index | RERANK "star wars"';
+      const query = EsqlQuery.fromSrc(src);
+
+      const cmd = [...commands.rerank.list(query.ast)][0];
+
+      expect(() => commands.rerank.setFields(cmd, ['newField'])).toThrow(
+        'RERANK command must have a ON option'
       );
     });
   });
 
-  describe('.setInferenceId()', () => {
-    it('can change query', () => {
-      const src = 'FROM index | LIMIT 1 | RERANK "star wars" ON field WITH id | LIMIT 2';
+  describe('.setWithParameter()', () => {
+    it('can add and update a new parameter to WITH map', () => {
+      const src =
+        'FROM index | RERANK "star wars" ON field WITH { "inference_id": "model_id" } | LIMIT 2';
       const query = EsqlQuery.fromSrc(src);
 
       const cmd = [...commands.rerank.list(query.ast)][0];
-      commands.rerank.setInferenceId(cmd, 'new_id');
+      commands.rerank.setWithParameter(cmd, 'scoreColumn', 'first_rank_score'); // create
+      commands.rerank.setWithParameter(cmd, 'scoreColumn', 'rank_score'); // update
+
+      const isWithOption = (arg: ESQLAstItem): arg is ESQLCommandOption =>
+        !!arg && !Array.isArray(arg) && arg.type === 'option' && arg.name === 'with';
+
+      const map = cmd.args.find(isWithOption)!.args[0] as ESQLMap;
+      const scoreColumnEntry = map?.entries?.find(
+        (entry) => entry.key.valueUnquoted === 'scoreColumn'
+      );
 
-      expect(BasicPrettyPrinter.expression(cmd.inferenceId)).toBe('new_id');
-      expect(query.print()).toBe(
-        'FROM index | LIMIT 1 | RERANK "star wars" ON field WITH new_id | LIMIT 2'
+      expect(BasicPrettyPrinter.expression(scoreColumnEntry!.value)).toBe('"rank_score"');
+      expect(query.print({ wrap: Infinity })).toBe(
+        'FROM index | RERANK "star wars" ON field WITH {"inference_id": "model_id", "scoreColumn": "rank_score"} | LIMIT 2'
       );
     });
   });

src/platform/packages/shared/kbn-esql-ast/src/mutate/commands/rerank/index.ts

@@ -12,8 +12,12 @@ import type {
   ESQLAstQueryExpression,
   ESQLAstRerankCommand,
   ESQLCommandOption,
-  ESQLIdentifierOrParam,
   ESQLStringLiteral,
+  ESQLParamLiteral,
+  ESQLMap,
+  ESQLAstItem,
+  ESQLMapEntry,
+  ESQLLiteral,
 } from '../../../types';
 import * as generic from '../../generic';
 
@@ -30,15 +34,79 @@ export const list = (ast: ESQLAstQueryExpression): IterableIterator<ESQLAstReran
   ) as IterableIterator<ESQLAstRerankCommand>;
 };
 
+/**
+ * Sets or updates the query text for the RERANK command.
+ *
+ * @param cmd The RERANK command AST node to modify.
+ * @param query The query text to set.
+ */
 export const setQuery = (cmd: ESQLAstRerankCommand, query: string | ESQLStringLiteral) => {
-  if (typeof query === 'string') {
-    query = Builder.expression.literal.string(query);
+  const queryLiteral = typeof query === 'string' ? Builder.expression.literal.string(query) : query;
+  const firstArg = cmd.args[0];
+
+  cmd.query = queryLiteral;
+
+  if (
+    firstArg &&
+    !Array.isArray(firstArg) &&
+    firstArg.type === 'function' &&
+    firstArg.name === '='
+  ) {
+    // It's an assignment, update the right side of the expression
+    firstArg.args[1] = queryLiteral;
+  } else {
+    // It's a simple query, replace the first argument
+    cmd.args[0] = queryLiteral;
   }
+};
+
+/**
+ * Sets, updates, or removes the target field for the RERANK command.
+ * This refers to the `targetField =` portion of the command, which specifies                                  │
+ * the new column where the rerank score will be stored.                                                       │
+ *
+ * @param cmd The RERANK command AST node to modify.
+ * @param target The name of the target field to set. If `null` the assignment is removed.
+ */
+export const setTargetField = (cmd: ESQLAstRerankCommand, target: string | null) => {
+  const firstArg = cmd.args[0];
+  const isAssignment =
+    firstArg && !Array.isArray(firstArg) && firstArg.type === 'function' && firstArg.name === '=';
 
-  cmd.query = query;
-  cmd.args[0] = query;
+  // Case 1: Set a new target field
+  if (target !== null) {
+    const newTargetColumn = Builder.expression.column(target);
+
+    if (isAssignment) {
+      // An assignment already exists => update the target field
+      firstArg.args[0] = newTargetColumn;
+    } else {
+      // No assignment exists => create one
+      const queryLiteral = cmd.query;
+      const assignment = Builder.expression.func.binary('=', [newTargetColumn, queryLiteral]);
+      cmd.args[0] = assignment;
+    }
+
+    cmd.targetField = newTargetColumn;
+  }
+  // Case 2: Remove the target field
+  else {
+    if (isAssignment) {
+      // An assignment exists, keep only the query
+      const queryLiteral = cmd.query;
+      cmd.args[0] = queryLiteral;
+    }
+    // If no assignment exists, do nothing
+    cmd.targetField = undefined;
+  }
 };
 
+/**
+ * Sets or updates the fields to be used for reranking in the ON clause.
+ *
+ * @param cmd The RERANK command AST node to modify.
+ * @param fields An array of field names or field nodes.
+ */
 export const setFields = (
   cmd: ESQLAstRerankCommand,
   fields: string[] | ESQLAstRerankCommand['fields']
@@ -51,21 +119,93 @@ export const setFields = (
     });
   }
 
+  const isOnOption = (arg: ESQLAstItem): arg is ESQLCommandOption =>
+    !!arg && !Array.isArray(arg) && arg.type === 'option' && arg.name === 'on';
+
+  const onOption = cmd.args.find(isOnOption);
+
+  if (!onOption) {
+    throw new Error('RERANK command must have a ON option');
+  }
+
   cmd.fields.length = 0;
   cmd.fields.push(...(fields as ESQLAstRerankCommand['fields']));
+
+  onOption.args.length = 0;
+  onOption.args.push(...(fields as ESQLAstRerankCommand['fields']));
 };
 
-export const setInferenceId = (cmd: ESQLAstRerankCommand, id: string | ESQLIdentifierOrParam) => {
-  if (typeof id === 'string') {
-    id = id[0] === '?' ? Builder.param.build(id) : Builder.identifier(id);
-  }
+/**
+ * Sets a parameter in the WITH clause of the RERANK command (e.g., 'inference_id').
+ * If the parameter already exists, its value is updated. Otherwise, it is added.
+ *
+ * @param cmd The RERANK command AST node to modify.
+ * @param key The name of the parameter to set.
+ * @param value The value of the parameter.
+ */
+export const setWithParameter = (
+  cmd: ESQLAstRerankCommand,
+  key: string,
+  value: string | ESQLStringLiteral | ESQLParamLiteral
+) => {
+  // Converts a value to an appropriate ESQLExpression based on its type
+  const toExpression = (val: string | ESQLLiteral | ESQLParamLiteral) => {
+    if (typeof val === 'string') {
+      return val.startsWith('?')
+        ? Builder.param.build(val)
+        : Builder.expression.literal.string(val);
+    }
+    return val;
+  };
 
-  if (id.type !== 'identifier' && id.type !== 'literal') {
-    throw new Error(`Invalid RERANK inferenceId: ${id}`);
+  const isWithOption = (arg: ESQLAstItem): arg is ESQLCommandOption =>
+    !!arg && !Array.isArray(arg) && arg.type === 'option' && arg.name === 'with';
+
+  // Validates and retrieves the map from a WITH option
+  const getWithOptionMap = (withOption: ESQLCommandOption): ESQLMap => {
+    const mapArg = withOption.args[0];
+
+    if (!mapArg || typeof mapArg === 'string' || Array.isArray(mapArg) || mapArg.type !== 'map') {
+      throw new Error('WITH option must contain a map');
+    }
+
+    return mapArg as ESQLMap;
+  };
+
+  // Normalizes a key for comparison by removing quotes
+  const normalizeKey = (keyValue: string | undefined): string => {
+    return keyValue?.replace(/"/g, '') ?? '';
+  };
+
+  // Checks if an entry in the map has the specified key
+  // normalizeKey avoid cases like: "inference_id" === "'inferenceId"' -> this is false
+  const getExistingEntry = (entry: ESQLMapEntry) =>
+    normalizeKey(entry.key.valueUnquoted ?? entry.key.value) === key;
+
+  // Creates a new map entry with the given key and value
+  const createMapEntry = (entryKey: string, entryValue: ESQLLiteral | ESQLParamLiteral) => ({
+    type: 'map-entry' as const,
+    name: 'map-entry' as const,
+    key: Builder.expression.literal.string(entryKey),
+    value: entryValue,
+    location: { min: 0, max: 0 },
+    text: '',
+    incomplete: false,
+  });
+
+  const withOption = cmd.args.find(isWithOption);
+
+  if (!withOption) {
+    throw new Error('RERANK command must have a WITH option');
   }
 
-  cmd.inferenceId = id;
+  const valueExpression = toExpression(value);
+  const map = getWithOptionMap(withOption);
+  const existingEntry = map.entries.find(getExistingEntry);
 
-  const withOption = cmd.args[2] as ESQLCommandOption;
-  withOption.args[0] = id;
+  if (existingEntry) {
+    existingEntry.value = valueExpression;
+  } else {
+    map.entries.push(createMapEntry(key, valueExpression));
+  }
 };