[ES|QL] Fix field loading when, both, METADATA and LOOKUP JOIN present (#211375)

## Summary

Closes https://github.com/elastic/kibana/issues/210080

- Constructs field lookup ES|QL query more robustly. Instead of
extracting the text of the first command, now it uses the AST to get the
source and metadata nodes from the `FROM` command and indices from the
`JOIN` command. Then it constructs a new `FROM` query using the `synth`
API.

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

@@ -8,7 +8,7 @@
  */
 
 import { Walker } from '../../../walker';
-import { ESQLAstQueryExpression, ESQLColumn, ESQLCommandOption } from '../../../types';
+import { ESQLAstQueryExpression, ESQLColumn, ESQLCommand, ESQLCommandOption } from '../../../types';
 import { Visitor } from '../../../visitor';
 import { cmpArr, findByPredicate } from '../../util';
 import * as generic from '../../generic';
@@ -23,12 +23,12 @@ import type { Predicate } from '../../types';
  * @returns A collection of [column, option] pairs for each metadata field found.
  */
 export const list = (
-  ast: ESQLAstQueryExpression
+  ast: ESQLAstQueryExpression | ESQLCommand<'from'>
 ): IterableIterator<[ESQLColumn, ESQLCommandOption]> => {
   type ReturnExpression = IterableIterator<ESQLColumn>;
   type ReturnCommand = IterableIterator<[ESQLColumn, ESQLCommandOption]>;
 
-  return new Visitor()
+  const visitor = new Visitor()
     .on('visitExpression', function* (): ReturnExpression {})
     .on('visitColumnExpression', function* (ctx): ReturnExpression {
       yield ctx.node;
@@ -53,8 +53,13 @@ export const list = (
       for (const command of ctx.visitCommands()) {
         yield* command;
       }
-    })
-    .visitQuery(ast);
+    });
+
+  if (ast.type === 'command') {
+    return visitor.visitCommand(ast);
+  } else {
+    return visitor.visitQuery(ast);
+  }
 };
 
 /**

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

@@ -8,14 +8,16 @@
  */
 
 import { Builder } from '../../../builder';
-import { ESQLAstQueryExpression, ESQLSource } from '../../../types';
+import { ESQLAstQueryExpression, ESQLCommand, ESQLSource } from '../../../types';
 import { Visitor } from '../../../visitor';
 import * as generic from '../../generic';
 import * as util from '../../util';
 import type { Predicate } from '../../types';
 
-export const list = (ast: ESQLAstQueryExpression): IterableIterator<ESQLSource> => {
-  return new Visitor()
+export const list = (
+  ast: ESQLAstQueryExpression | ESQLCommand<'from'>
+): IterableIterator<ESQLSource> => {
+  const visitor = new Visitor()
     .on('visitFromCommand', function* (ctx): IterableIterator<ESQLSource> {
       for (const argument of ctx.arguments()) {
         if (argument.type === 'source') {
@@ -28,8 +30,13 @@ export const list = (ast: ESQLAstQueryExpression): IterableIterator<ESQLSource>
       for (const command of ctx.visitCommands()) {
         yield* command;
       }
-    })
-    .visitQuery(ast);
+    });
+
+  if (ast.type === 'command') {
+    return visitor.visitCommand(ast);
+  } else {
+    return visitor.visitQuery(ast);
+  }
 };
 
 export const findByPredicate = (

src/platform/packages/shared/kbn-esql-ast/src/synth/__tests__/expr_template.test.ts

@@ -39,3 +39,12 @@ test('can compose nodes into templated string', () => {
 
   expect(text).toBe('a.b.c = FN(1, a.b.c)');
 });
+
+test('creates a list of nodes separated by command, if array passed in', () => {
+  const arg1 = expr`1`;
+  const arg2 = expr`a.b.c`;
+  const value = expr`fn(${[arg1, arg2]})`;
+  const text = BasicPrettyPrinter.expression(value);
+
+  expect(text).toBe('FN(1, a.b.c)');
+});

src/platform/packages/shared/kbn-esql-ast/src/synth/helpers.ts

@@ -61,7 +61,17 @@ export const createSynthMethod = <N extends ESQLProperNode>(
         if (i < params.length) {
           const param = params[i];
           if (typeof param === 'string') src += param;
-          else src += serialize(param);
+          else if (Array.isArray(param)) {
+            let list: string = '';
+
+            for (const item of param) {
+              const serialized = typeof item === 'string' ? item : serialize(item);
+
+              list += (list ? ', ' : '') + serialized;
+            }
+
+            src += list;
+          } else src += serialize(param);
         }
       }
       return generator(src, opts);

src/platform/packages/shared/kbn-esql-ast/src/synth/types.ts

@@ -14,7 +14,7 @@ export type SynthGenerator<N extends ESQLProperNode> = (src: string, opts?: Pars
 
 export type SynthTaggedTemplate<N extends ESQLProperNode> = (
   template: TemplateStringsArray,
-  ...params: Array<ESQLAstExpression | string>
+  ...params: Array<ESQLAstExpression | ESQLAstExpression[] | string | []>
 ) => N;
 
 export type SynthTaggedTemplateWithOpts<N extends ESQLProperNode> = (

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

@@ -12,10 +12,12 @@ import type {
   ESQLAstItem,
   ESQLAstMetricsCommand,
   ESQLAstQueryExpression,
+  ESQLColumn,
   ESQLMessage,
   ESQLSingleAstItem,
+  ESQLSource,
 } from '@kbn/esql-ast';
-import { mutate } from '@kbn/esql-ast';
+import { mutate, synth } from '@kbn/esql-ast';
 import { FunctionDefinition } from '../definitions/types';
 import { getAllArrayTypes, getAllArrayValues } from '../shared/helpers';
 import { getMessageFromId } from './errors';
@@ -25,32 +27,43 @@ export function buildQueryForFieldsFromSource(queryString: string, ast: ESQLAst)
   const firstCommand = ast[0];
   if (!firstCommand) return '';
 
-  let query = '';
+  const sources: ESQLSource[] = [];
+  const metadataFields: ESQLColumn[] = [];
 
   if (firstCommand.name === 'metrics') {
     const metrics = firstCommand as ESQLAstMetricsCommand;
-    query = `FROM ${metrics.sources.map((source) => source.name).join(', ')}`;
-  } else {
-    query = queryString.substring(0, firstCommand.location.max + 1);
+
+    sources.push(...metrics.sources);
+  } 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: { name },
-      },
-    }) => name
-  );
+  const joinIndices = joinSummary.map(({ target: { index } }) => index);
 
   if (joinIndices.length > 0) {
-    query += `, ${joinIndices.join(', ')}`;
+    sources.push(...joinIndices);
   }
 
-  return query;
+  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[]) {