feat(YQL): improve variable symbol value types

gravity-ui · Dec 25, 2024 · 467e2c3 · 467e2c3
1 parent b8037da
commit 467e2c3
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 9 deletions.
diff --git a/src/autocomplete/databases/yql/yql-autocomplete.ts b/src/autocomplete/databases/yql/yql-autocomplete.ts
@@ -32,7 +32,7 @@ import {
     ISymbolTableVisitor,
     ProcessVisitedRulesResult,
 } from '../../shared/autocomplete-types';
-import {ColumnAliasSymbol, TableSymbol} from '../../shared/symbol-table.js';
+import {ColumnAliasSymbol, ExtendedVariable, TableSymbol} from '../../shared/symbol-table.js';
 import {TableQueryPosition} from '../../shared/tables';
 import {isStartingToWriteRule} from '../../shared/cursor.js';
 import {shouldSuggestTemplates} from '../../shared/query.js';
@@ -165,7 +165,7 @@ class YQLVariableSymbolTableVisitor extends YQLSymbolTableVisitor {
                 const variable = getVariable(index);
                 if (variable) {
                     this.symbolTable.addNewSymbolOfType(
-                        c3.VariableSymbol,
+                        ExtendedVariable,
                         this.scope,
                         variable,
                         variableValue,
@@ -188,7 +188,7 @@ class YQLVariableSymbolTableVisitor extends YQLSymbolTableVisitor {
             if (variable) {
                 const value = context.literal_value()?.getText();
 
-                this.symbolTable.addNewSymbolOfType(c3.VariableSymbol, this.scope, variable, value);
+                this.symbolTable.addNewSymbolOfType(ExtendedVariable, this.scope, variable, value);
             }
         } catch (error) {
             if (!(error instanceof c3.DuplicateSymbolError)) {
@@ -228,7 +228,7 @@ class YQLVariableSymbolTableVisitor extends YQLSymbolTableVisitor {
             const variable = context.bind_parameter()?.an_id_or_type()?.getText();
             if (variable) {
                 this.symbolTable.addNewSymbolOfType(
-                    c3.VariableSymbol,
+                    ExtendedVariable,
                     this.scope,
                     variable,
                     undefined,

diff --git a/src/autocomplete/shared/autocomplete-types.ts b/src/autocomplete/shared/autocomplete-types.ts
@@ -74,7 +74,7 @@ export interface ColumnAliasSuggestion {
 
 export interface VariableSuggestion {
     name: string;
-    value: unknown;
+    value?: {columns?: string[]};
 }
 
 export type LexerConstructor<T> = new (input: CharStream) => T;

diff --git a/src/autocomplete/shared/symbol-table.ts b/src/autocomplete/shared/symbol-table.ts
@@ -16,6 +16,10 @@ export class TableSymbol extends c3.TypedSymbol {
     }
 }
 
+export class ExtendedVariable extends c3.VariableSymbol {
+    value: {columns?: string[]} | undefined;
+}
+
 export function getUniqueTableSuggestions(symbols: TableSymbol[] = []): Table[] {
     const suggestionsMap = symbols.reduce(
         (acc, table) => {

diff --git a/src/autocomplete/shared/variables.ts b/src/autocomplete/shared/variables.ts
@@ -9,7 +9,7 @@ import {
     SymbolTableVisitor,
     VariableSuggestion,
 } from './autocomplete-types';
-import {getScope} from './symbol-table';
+import {ExtendedVariable, getScope} from './symbol-table';
 import {computeTokenContext} from './compute-token-position';
 
 export function getVariableSuggestions<L extends LexerType, P extends ParserType>(
@@ -37,16 +37,16 @@ export function getVariableSuggestions<L extends LexerType, P extends ParserType
 
 function suggestVariables(symbolTable: c3.SymbolTable, context: ParseTree): VariableSuggestion[] {
     const scope = getScope(context, symbolTable);
-    let symbols: c3.VariableSymbol[] = [];
+    let symbols: ExtendedVariable[] = [];
 
     // Local scope
     if (scope instanceof c3.ScopedSymbol) {
-        symbols = scope.getNestedSymbolsOfTypeSync(c3.VariableSymbol);
+        symbols = scope.getNestedSymbolsOfTypeSync(ExtendedVariable);
 
         // Global scope
     } else if (symbolTable) {
         symbols = symbolTable
-            .getNestedSymbolsOfTypeSync(c3.VariableSymbol)
+            .getNestedSymbolsOfTypeSync(ExtendedVariable)
             // If symbol's parent has context it means it is local scoped, so no need to suggest it in global scope
             .filter((symbol) => !symbol.parent?.context);
     }