feat(autocomplete): suggest columns from subqueries

Author: Raubzeug

package-lock.json

@@ -24,7 +24,7 @@
     "@gravity-ui/react-data-table": "^2.1.1",
     "@gravity-ui/table": "^1.7.0",
     "@gravity-ui/uikit": "^6.33.0",
-    "@gravity-ui/websql-autocomplete": "^12.7.0",
+    "@gravity-ui/websql-autocomplete": "^13.1.0",
     "@hookform/resolvers": "^3.9.0",
     "@reduxjs/toolkit": "^2.2.3",
     "@tanstack/react-table": "^8.19.3",

package.json

@@ -1,6 +1,7 @@
 import type {
     ColumnAliasSuggestion,
     KeywordSuggestion,
+    VariableSuggestion,
 } from '@gravity-ui/websql-autocomplete/shared';
 import type {YQLEntity, YqlAutocompleteResult} from '@gravity-ui/websql-autocomplete/yql';
 import * as monaco from 'monaco-editor/esm/vs/editor/editor.api';
@@ -107,6 +108,10 @@ function removeBackticks(value: string) {
     return value.slice(sliceStart, sliceEnd);
 }
 
+function isVariable(value: string) {
+    return value.startsWith('$');
+}
+
 function removeStartSlash(value: string) {
     if (value.startsWith('/')) {
         return value.slice(1);
@@ -199,7 +204,8 @@ function getColumnDetails(col: AutocompleteColumn) {
 
 export async function generateColumnsSuggestion(
     rangeToInsertSuggestion: monaco.IRange,
-    suggestColumns: YqlAutocompleteResult['suggestColumns'] | undefined,
+    suggestColumns: YqlAutocompleteResult['suggestColumns'],
+    suggestVariables: YqlAutocompleteResult['suggestVariables'],
     database: string,
 ): Promise<monaco.languages.CompletionItem[]> {
     if (!suggestColumns?.tables) {
@@ -212,7 +218,7 @@ export async function generateColumnsSuggestion(
     const normalizedTableNames =
         suggestColumns.tables?.map((entity) => {
             let normalizedEntityName = removeBackticks(entity.name);
-            if (!normalizedEntityName.endsWith('/')) {
+            if (!normalizedEntityName.endsWith('/') && !isVariable(normalizedEntityName)) {
                 normalizedEntityName = `${normalizedEntityName}/`;
             }
             return normalizeEntityPrefix(normalizedEntityName, database);
@@ -221,13 +227,31 @@ export async function generateColumnsSuggestion(
     // remove duplicates if any
     const filteredTableNames = Array.from(new Set(normalizedTableNames));
 
-    const autocompleteResponse = await window.api.viewer.autocomplete({
-        database,
-        table: filteredTableNames,
-        limit: 1000,
-    });
-    if (!autocompleteResponse.Success) {
-        return [];
+    const tableSources = filteredTableNames.filter((name) => !isVariable(name));
+
+    let autocompleteEntities: TAutocompleteEntity[] = [];
+    if (tableSources.length) {
+        const autocompleteResponse = await window.api.viewer.autocomplete({
+            database,
+            table: tableSources,
+            limit: 1000,
+        });
+        if (autocompleteResponse.Success) {
+            autocompleteEntities = autocompleteResponse.Result.Entities ?? [];
+        }
+    }
+
+    const variableSources = filteredTableNames.filter(isVariable);
+    let additionalColumns: TAutocompleteEntity[] = [];
+    if (variableSources.length) {
+        variableSources.forEach((source) => {
+            additionalColumns =
+                suggestVariables
+                    // Variable name from suggestions doesn't include $ sign
+                    ?.find((variable) => source.slice(1) === variable.name)
+                    ?.value?.columns?.map((col) => ({Name: col, Type: 'column', Parent: source})) ??
+                [];
+        });
     }
 
     const tableNameToAliasMap = suggestColumns.tables?.reduce(
@@ -246,7 +270,7 @@ export async function generateColumnsSuggestion(
         {} as Record<string, string[]>,
     );
 
-    autocompleteResponse.Result.Entities?.forEach((col) => {
+    autocompleteEntities.concat(additionalColumns).forEach((col) => {
         if (!isAutocompleteColumn(col)) {
             return;
         }
@@ -293,7 +317,7 @@ export async function generateColumnsSuggestion(
             normalizedColumns?.push(columnNameSuggestion);
         }
     });
-    if (normalizedColumns && normalizedColumns.length > 0) {
+    if (normalizedColumns && normalizedColumns.length > 1) {
         const allColumns = normalizedColumns.join(', ');
         suggestions.push({
             label: allColumns,
@@ -341,13 +365,13 @@ export function generateKeywordsSuggestion(
 
 export function generateVariableSuggestion(
     rangeToInsertSuggestion: monaco.IRange,
-    suggestVariables?: string[],
+    suggestVariables?: VariableSuggestion[],
 ) {
     if (!suggestVariables) {
         return [];
     }
-    return suggestVariables.map((rawVariable) => {
-        const variable = '$' + rawVariable;
+    return suggestVariables.map(({name}) => {
+        const variable = '$' + name;
         return {
             label: variable,
             insertText: variable,

src/utils/monaco/yql/generateSuggestions.ts

@@ -128,6 +128,7 @@ async function getSuggestions(
     const columnsSuggestions = await generateColumnsSuggestion(
         rangeToInsertSuggestion,
         parseResult.suggestColumns,
+        parseResult.suggestVariables,
         database,
     );