support row and ts in from e where in subqueries

Author: bartoval

src/platform/packages/shared/kbn-esql-language/src/commands/registry/complete_items.ts

@@ -375,18 +375,14 @@ function buildSubqueryCompleteItem(sourceCommand: string): ISuggestionItem {
   });
 }
 
-export function buildSubqueryCompleteItems(sourceCommands?: string[]): ISuggestionItem[] {
-  const allowedSourceCommands = sourceCommands
-    ? new Set(sourceCommands.map((command) => command.toLowerCase()))
-    : undefined;
-
+export function buildSubqueryCompleteItems(): ISuggestionItem[] {
   return esqlCommandRegistry
     .getAllCommands()
     .filter(
-      ({ name, metadata }) =>
+      ({ metadata }) =>
         metadata.subquerySource === true &&
         metadata.hidden !== true &&
-        (!allowedSourceCommands || allowedSourceCommands.has(name))
+        !metadata.subquerySourceHidden
     )
     .map(({ name }) => buildSubqueryCompleteItem(name));
 }

src/platform/packages/shared/kbn-esql-language/src/commands/registry/from/autocomplete.test.ts

@@ -37,6 +37,8 @@ const fromExpectSuggestions = (
   );
 };
 
+const subquerySuggestions = ['(FROM $0)'];
+
 const visibleIndices =
   mockContext.sources?.filter((source) => !source.hidden).map((source) => source.name) || [];
 
@@ -92,8 +94,16 @@ describe('FROM Autocomplete', () => {
     });
 
     test('suggests visible indices on space', async () => {
-      await fromExpectSuggestions('from /', [...visibleIndices, '(FROM $0)'], mockCallbacks);
-      await fromExpectSuggestions('FROM /', [...visibleIndices, '(FROM $0)'], mockCallbacks);
+      await fromExpectSuggestions(
+        'from /',
+        [...visibleIndices, ...subquerySuggestions],
+        mockCallbacks
+      );
+      await fromExpectSuggestions(
+        'FROM /',
+        [...visibleIndices, ...subquerySuggestions],
+        mockCallbacks
+      );
       await fromExpectSuggestions('from /index', visibleIndices, mockCallbacks);
     });
 
@@ -102,13 +112,17 @@ describe('FROM Autocomplete', () => {
     });
 
     test('does create suggestions after a closed quote', async () => {
-      await fromExpectSuggestions('FROM "lolz", ', [...visibleIndices, '(FROM $0)'], mockCallbacks);
+      await fromExpectSuggestions(
+        'FROM "lolz", ',
+        [...visibleIndices, ...subquerySuggestions],
+        mockCallbacks
+      );
     });
 
     test('doesnt suggest indices twice', async () => {
       await fromExpectSuggestions(
         'from index, ',
-        [...visibleIndices.filter((i) => i !== 'index'), '(FROM $0)'],
+        [...visibleIndices.filter((i) => i !== 'index'), ...subquerySuggestions],
         mockCallbacks
       );
     });
@@ -134,16 +148,24 @@ describe('FROM Autocomplete', () => {
       const expectedSuggestions = visibleDataSources.map((source) => source.name);
       mockContext.sources = visibleDataSources;
 
-      await fromExpectSuggestions('from ', [...expectedSuggestions, '(FROM $0)'], mockCallbacks);
-      await fromExpectSuggestions('FROM ', [...expectedSuggestions, '(FROM $0)'], mockCallbacks);
+      await fromExpectSuggestions(
+        'from ',
+        [...expectedSuggestions, ...subquerySuggestions],
+        mockCallbacks
+      );
+      await fromExpectSuggestions(
+        'FROM ',
+        [...expectedSuggestions, ...subquerySuggestions],
+        mockCallbacks
+      );
       await fromExpectSuggestions(
         'FROM a,/',
         expectedSuggestions.filter((i) => i !== 'a'),
         mockCallbacks
       );
       await fromExpectSuggestions(
         'from a, /',
-        [...expectedSuggestions.filter((i) => i !== 'a'), '(FROM $0)'],
+        [...expectedSuggestions.filter((i) => i !== 'a'), ...subquerySuggestions],
         mockCallbacks
       );
       await fromExpectSuggestions('from *,/', expectedSuggestions, mockCallbacks);
@@ -161,7 +183,7 @@ describe('FROM Autocomplete', () => {
       const expectedFromViews = ['my_saved_view', 'my-view'];
       await fromExpectSuggestions(
         'from ',
-        [...expectedFromSources, ...expectedFromViews, '(FROM $0)'],
+        [...expectedFromSources, ...expectedFromViews, ...subquerySuggestions],
         mockCallbacks,
         contextWithViews
       );
@@ -192,7 +214,7 @@ describe('FROM Autocomplete', () => {
 
       await fromExpectSuggestions(
         'from ',
-        [...expectedFromSources, 'my_saved_view', 'my_dataset', '(FROM $0)'],
+        [...expectedFromSources, 'my_saved_view', 'my_dataset', ...subquerySuggestions],
         mockCallbacks,
         contextWithDatasets
       );
@@ -294,13 +316,17 @@ describe('FROM Autocomplete', () => {
     ].sort();
 
     test('suggests subquery on space after FROM', async () => {
-      await fromExpectSuggestions('from /', [...visibleIndices, '(FROM $0)'], mockCallbacks);
+      await fromExpectSuggestions(
+        'from /',
+        [...visibleIndices, ...subquerySuggestions],
+        mockCallbacks
+      );
     });
 
     test('suggests subquery after comma', async () => {
       await fromExpectSuggestions(
         'from index, /',
-        [...visibleIndices.filter((i) => i !== 'index'), '(FROM $0)'],
+        [...visibleIndices.filter((i) => i !== 'index'), ...subquerySuggestions],
         mockCallbacks
       );
     });

src/platform/packages/shared/kbn-esql-language/src/commands/registry/from/autocomplete.ts

@@ -71,10 +71,8 @@ async function handleFromAutocomplete(
   // Use commandText for pattern matching (e.g., /METADATA\s+$/ and trailing whitespace)
   // checks need to operate on the current command only, not the entire query
   const commandText = query.substring(command.location.min, cursorPos);
-  const subquerySuggestion =
-    commandText.length > command.name.length
-      ? buildSubqueryCompleteItems([command.name]).at(0)
-      : undefined;
+  const subquerySuggestions =
+    commandText.length > command.name.length ? buildSubqueryCompleteItems() : undefined;
   const indicesBrowserSuggestion = await getIndicesBrowserSuggestion({
     callbacks,
     context,
@@ -99,7 +97,7 @@ async function handleFromAutocomplete(
       context,
       innerText,
       command.location.min,
-      subquerySuggestion
+      subquerySuggestions
     );
     if (indicesBrowserSuggestion) {
       suggestions.unshift(indicesBrowserSuggestion);
@@ -122,7 +120,7 @@ async function handleFromAutocomplete(
     context,
     callbacks,
     indexes,
-    subquerySuggestion,
+    subquerySuggestions,
     {
       textBeforeCursor: innerText,
       commandStart: command.location.min, // Full-query offset of this FROM command.
@@ -141,7 +139,7 @@ function suggestInitialSources(
   context: ICommandContext | undefined,
   innerText: string,
   commandStart: number,
-  subquerySuggestion?: ISuggestionItem
+  subquerySuggestions?: ISuggestionItem[]
 ): ISuggestionItem[] {
   let sources = context?.sources ?? [];
 
@@ -157,8 +155,8 @@ function suggestInitialSources(
   const datasetSuggestions = buildDatasetsDefinitions(context?.datasets ?? [], []);
   const suggestions = [...sourceSuggestions, ...viewSuggestions, ...datasetSuggestions];
 
-  if (subquerySuggestion && shouldSuggestSubquery(context)) {
-    suggestions.push(subquerySuggestion);
+  if (subquerySuggestions && shouldSuggestSubquery(context)) {
+    suggestions.push(...subquerySuggestions);
   }
 
   return suggestions;
@@ -189,7 +187,7 @@ async function suggestAdditionalSources(
   context: ICommandContext | undefined,
   callbacks: ICommandCallbacks | undefined,
   indexes: ReturnType<typeof getSourcesFromCommands>,
-  subquerySuggestion?: ISuggestionItem,
+  subquerySuggestions?: ISuggestionItem[],
   sourceReplacementContext?: {
     textBeforeCursor: string;
     commandStart: number;
@@ -227,8 +225,8 @@ async function suggestAdditionalSources(
     canRewriteFromToTs ? sourceReplacementContext : undefined
   );
 
-  if (subquerySuggestion && isRestartingExpression(innerText) && shouldSuggestSubquery(context)) {
-    suggestions.push(subquerySuggestion);
+  if (subquerySuggestions && isRestartingExpression(innerText) && shouldSuggestSubquery(context)) {
+    suggestions.push(...subquerySuggestions);
   }
 
   return suggestions;

src/platform/packages/shared/kbn-esql-language/src/commands/registry/registry.test.ts

@@ -261,12 +261,16 @@ describe('CommandRegistry', () => {
       expect(commandNames).toContain('hiddenOutsideCommand');
     });
 
-    test('should only return FROM when starting a subquery', () => {
-      const fromCommand = createMockCommand('from');
+    test('should only return subquery source commands when starting a subquery', () => {
+      const fromCommand = createMockCommand('from', undefined, { subquerySource: true });
+      const rowCommand = createMockCommand('row', undefined, { subquerySource: true, subquerySourceHidden: true });
+      const tsCommand = createMockCommand('ts', undefined, { subquerySource: true, subquerySourceHidden: true });
       const evalCommand = createMockCommand('eval');
       const whereCommand = createMockCommand('where');
 
       registry.registerCommand(fromCommand);
+      registry.registerCommand(rowCommand);
+      registry.registerCommand(tsCommand);
       registry.registerCommand(evalCommand);
       registry.registerCommand(whereCommand);
 
@@ -276,6 +280,10 @@ describe('CommandRegistry', () => {
 
       const commandNames = commands.map((cmd) => cmd.name);
       expect(commandNames).toEqual(['from']);
+      expect(commandNames).not.toContain('row');
+      expect(commandNames).not.toContain('ts');
+      expect(commandNames).not.toContain('eval');
+      expect(commandNames).not.toContain('where');
     });
   });
 

src/platform/packages/shared/kbn-esql-language/src/commands/registry/registry.ts

@@ -99,6 +99,7 @@ export interface ICommandMetadata {
   hiddenAfterCommands?: string[]; // Optional list of command names; this command is not suggested when any of them appear anywhere in the pipeline
   subquerySupport?: boolean; // Temporary property to indicate if the command supports subquery suggestions.
   subquerySource?: boolean; // Optional property to indicate if the command can start a subquery expression.
+  subquerySourceHidden?: boolean; // Hides the command from subquery source suggestions (e.g. snapshot-only commands).
   subqueryRestrictions?: {
     hideInside: boolean; // Command is hidden inside subqueries
     hideOutside: boolean; // Command is hidden outside subqueries (at root level)
@@ -282,7 +283,7 @@ export class CommandRegistry implements ICommandRegistry {
     const queryContainsSubqueries = options?.queryContainsSubqueries ?? false;
 
     const filtered = isStartingSubquery
-      ? allCommands.filter(({ name }) => name === 'from')
+      ? allCommands.filter(({ metadata }) => !!metadata.subquerySource && !metadata.subquerySourceHidden)
       : allCommands;
 
     // Then apply subquery restrictions

src/platform/packages/shared/kbn-esql-language/src/commands/registry/row/index.ts

@@ -27,6 +27,7 @@ export const rowCommand = {
   metadata: {
     type: 'source' as const,
     subquerySource: true,
+    subquerySourceHidden: true,
     description: i18n.translate('kbn-esql-language.esql.definitions.rowDoc', {
       defaultMessage:
         'Produces a row with one or more columns with values that you specify. This can be useful for testing.',

src/platform/packages/shared/kbn-esql-language/src/commands/registry/timeseries/index.ts

@@ -26,6 +26,8 @@ export const timeseriesCommand = {
   methods: timeseriesCommandMethods,
   metadata: {
     type: 'source' as const,
+    subquerySource: true,
+    subquerySourceHidden: true,
     hidden: false,
     preview: true,
     isTimeseries: true,

src/platform/packages/shared/kbn-esql-language/src/commands/registry/where/autocomplete.test.ts

@@ -305,12 +305,10 @@ describe('WHERE Autocomplete', () => {
       await whereExpectSuggestions('from index | WHERE doubleField in ', [
         '($0)',
         '(FROM $0)',
-        '(ROW $0)',
       ]);
       await whereExpectSuggestions('from index | WHERE doubleField not in ', [
         '($0)',
         '(FROM $0)',
-        '(ROW $0)',
       ]);
 
       await whereExpectSuggestions(