[8.x] [ES|QL] Add`SORT` command mutation APIs (#197185)

examples/esql_ast_inspector/public/components/esql_inspector/helpers.tsx

@@ -82,6 +82,8 @@ export const highlight = (query: EsqlQuery): Annotation[] => {
   });
 
   Walker.visitComments(query.ast, (comment) => {
+    if (!comment.location) return;
+
     annotations.push([
       comment.location.min,
       comment.location.max,

packages/kbn-esql-ast/src/builder/builder.ts

@@ -11,6 +11,8 @@
 
 import {
   ESQLAstComment,
+  ESQLAstCommentMultiLine,
+  ESQLAstCommentSingleLine,
   ESQLAstQueryExpression,
   ESQLColumn,
   ESQLCommand,
@@ -20,6 +22,7 @@ import {
   ESQLIntegerLiteral,
   ESQLList,
   ESQLLocation,
+  ESQLOrderExpression,
   ESQLSource,
 } from '../types';
 import { AstNodeParserFields, AstNodeTemplate, PartialFields } from './types';
@@ -63,17 +66,17 @@ export namespace Builder {
     };
   };
 
-  export const comment = (
-    subtype: ESQLAstComment['subtype'],
+  export const comment = <S extends ESQLAstComment['subtype']>(
+    subtype: S,
     text: string,
-    location: ESQLLocation
-  ): ESQLAstComment => {
+    location?: ESQLLocation
+  ): S extends 'multi-line' ? ESQLAstCommentMultiLine : ESQLAstCommentSingleLine => {
     return {
       type: 'comment',
       subtype,
       text,
       location,
-    };
+    } as S extends 'multi-line' ? ESQLAstCommentMultiLine : ESQLAstCommentSingleLine;
   };
 
   export namespace expression {
@@ -130,6 +133,20 @@ export namespace Builder {
       };
     };
 
+    export const order = (
+      operand: ESQLColumn,
+      template: Omit<AstNodeTemplate<ESQLOrderExpression>, 'name' | 'args'>,
+      fromParser?: Partial<AstNodeParserFields>
+    ): ESQLOrderExpression => {
+      return {
+        ...template,
+        ...Builder.parserFields(fromParser),
+        name: '',
+        args: [operand],
+        type: 'order',
+      };
+    };
+
     export const inlineCast = (
       template: Omit<AstNodeTemplate<ESQLInlineCast>, 'name'>,
       fromParser?: Partial<AstNodeParserFields>

packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts

@@ -106,7 +106,7 @@ export const removeByPredicate = (
   option.args.splice(index, 1);
 
   if (option.args.length === 0) {
-    generic.removeCommandOption(ast, option);
+    generic.commands.options.remove(ast, option);
   }
 
   return tuple;
@@ -148,16 +148,16 @@ export const insert = (
   fieldName: string | string[],
   index: number = -1
 ): [column: ESQLColumn, option: ESQLCommandOption] | undefined => {
-  let option = generic.findCommandOptionByName(ast, 'from', 'metadata');
+  let option = generic.commands.options.findByName(ast, 'from', 'metadata');
 
   if (!option) {
-    const command = generic.findCommandByName(ast, 'from');
+    const command = generic.commands.findByName(ast, 'from');
 
     if (!command) {
       return;
     }
 
-    option = generic.appendCommandOption(command, 'metadata');
+    option = generic.commands.options.append(command, 'metadata');
   }
 
   const parts: string[] = typeof fieldName === 'string' ? [fieldName] : fieldName;
@@ -189,7 +189,7 @@ export const upsert = (
   fieldName: string | string[],
   index: number = -1
 ): [column: ESQLColumn, option: ESQLCommandOption] | undefined => {
-  const option = generic.findCommandOptionByName(ast, 'from', 'metadata');
+  const option = generic.commands.options.findByName(ast, 'from', 'metadata');
 
   if (option) {
     const parts = Array.isArray(fieldName) ? fieldName : [fieldName];

packages/kbn-esql-ast/src/mutate/commands/from/sources.ts

@@ -67,7 +67,7 @@ export const remove = (
     return undefined;
   }
 
-  const success = generic.removeCommandArgument(ast, node);
+  const success = generic.commands.args.remove(ast, node);
 
   return success ? node : undefined;
 };
@@ -78,7 +78,7 @@ export const insert = (
   clusterName?: string,
   index: number = -1
 ): ESQLSource | undefined => {
-  const command = generic.findCommandByName(ast, 'from');
+  const command = generic.commands.findByName(ast, 'from');
 
   if (!command) {
     return;
@@ -87,7 +87,7 @@ export const insert = (
   const source = Builder.expression.indexSource(indexName, clusterName);
 
   if (index === -1) {
-    generic.appendCommandArgument(command, source);
+    generic.commands.args.append(command, source);
   } else {
     command.args.splice(index, 0, source);
   }

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

@@ -9,5 +9,6 @@
 
 import * as from from './from';
 import * as limit from './limit';
+import * as sort from './sort';
 
-export { from, limit };
+export { from, limit, sort };

packages/kbn-esql-ast/src/mutate/commands/limit/index.ts

@@ -19,7 +19,7 @@ import { Predicate } from '../../types';
  * @returns A collection of "LIMIT" commands.
  */
 export const list = (ast: ESQLAstQueryExpression): IterableIterator<ESQLCommand> => {
-  return generic.listCommands(ast, (cmd) => cmd.name === 'limit');
+  return generic.commands.list(ast, (cmd) => cmd.name === 'limit');
 };
 
 /**
@@ -55,13 +55,13 @@ export const find = (
  * @returns The removed "LIMIT" command, if any.
  */
 export const remove = (ast: ESQLAstQueryExpression, index: number = 0): ESQLCommand | undefined => {
-  const command = generic.findCommandByName(ast, 'limit', index);
+  const command = generic.commands.findByName(ast, 'limit', index);
 
   if (!command) {
     return;
   }
 
-  const success = generic.removeCommand(ast, command);
+  const success = !!generic.commands.remove(ast, command);
 
   if (!success) {
     return;
@@ -128,7 +128,7 @@ export const upsert = (
     args: [literal],
   });
 
-  generic.appendCommand(ast, command);
+  generic.commands.append(ast, command);
 
   return command;
 };