feat(YQL): support variables and scopes

src/autocomplete/databases/yql/helpers.ts

@@ -344,6 +344,18 @@ function checkShouldSuggestAllColumns(props: GetParticularSuggestionProps): bool
     return false;
 }
 
+function checkShouldSuggestVariables({
+    anyRuleInList,
+}: GetParticularSuggestionProps): boolean | undefined {
+    return anyRuleInList([
+        YQLParser.RULE_expr,
+        YQLParser.RULE_table_ref,
+        YQLParser.RULE_simple_table_ref_core,
+        YQLParser.RULE_pure_column_or_named,
+        YQLParser.RULE_using_call_expr,
+    ]);
+}
+
 function getSimpleTypesSuggestions({
     anyRuleInList,
     allRulesInList,
@@ -504,13 +516,15 @@ export function getGranularSuggestions(
     const suggestAggregateFunctions = getAggregateFunctionsSuggestions(props);
     const shouldSuggestTableHints = checkShouldSuggestTableHints(props);
     const suggestEntitySettings = getEntitySettingsSuggestions(props);
+    const shouldSuggestVariables = checkShouldSuggestVariables(props);
 
     return {
         suggestWindowFunctions,
         shouldSuggestTableIndexes,
         shouldSuggestColumns,
         shouldSuggestAllColumns,
         shouldSuggestColumnAliases: shouldSuggestColumns,
+        shouldSuggestVariables,
         suggestSimpleTypes,
         suggestPragmas,
         suggestUdfs,

src/autocomplete/databases/yql/tests/yql/select/select.test.ts

@@ -311,9 +311,10 @@ test('should suggest properly after LIMIT', () => {
 
 test('should suggest tables with inline comment', () => {
     const autocompleteResult = parseYqlQueryWithCursor(
-        'SELECT * FROM | --SELECT * FROM test_table',
+        '$foo = "bar"; SELECT * FROM | --SELECT * FROM test_table',
     );
     expect(autocompleteResult.suggestEntity).toEqual(['table', 'view', 'externalTable']);
+    expect(autocompleteResult.suggestVariables).toEqual(['foo']);
 });
 
 test('should suggest tables with multiline comment', () => {
@@ -329,3 +330,59 @@ test('should not report errors', () => {
 
     expect(autocompleteResult.errors).toHaveLength(0);
 });
+
+test('should suggest variables name for column', () => {
+    const autocompleteResult = parseYqlQueryWithCursor(
+        'DECLARE $prefix AS String; SELECT | FROM test_table',
+    );
+    const variablesSuggestions = ['prefix'];
+
+    expect(autocompleteResult.suggestVariables).toEqual(variablesSuggestions);
+});
+test('should suggest variables name for table name', () => {
+    const autocompleteResult = parseYqlQueryWithCursor(
+        'DECLARE $prefix AS String; SELECT * FROM |',
+    );
+    const variablesSuggestions = ['prefix'];
+
+    expect(autocompleteResult.suggestVariables).toEqual(variablesSuggestions);
+});
+test('should suggest variables name as columns', () => {
+    const autocompleteResult = parseYqlQueryWithCursor('$prefix, $foo = (2, 3); SELECT |');
+    const variablesSuggestions = ['prefix', 'foo'];
+
+    expect(autocompleteResult.suggestVariables).toEqual(variablesSuggestions);
+});
+test('should suggest variables name in global scope', () => {
+    const autocompleteResult = parseYqlQueryWithCursor(
+        '$test = 1; DEFINE SUBQUERY $foo($name) AS $baz = 1;\n select * from test_table where bar == $name AND baz == $baz END DEFINE; $baz2 = 2; select |',
+    );
+    const variablesSuggestions = ['test', 'foo', 'baz2'];
+
+    expect(autocompleteResult.suggestVariables).toEqual(variablesSuggestions);
+});
+test('should suggest variables name in local scope', () => {
+    const autocompleteResult = parseYqlQueryWithCursor(
+        '$test = 1; DEFINE SUBQUERY $foo($name) AS $baz = 1;\n select | from test_table where bar == $name AND baz == $baz END DEFINE; $baz2 = 2; select ',
+    );
+    const variablesSuggestions = ['name', 'baz'];
+
+    expect(autocompleteResult.suggestVariables).toEqual(variablesSuggestions);
+});
+test('should suggest variables inside lambda', () => {
+    const autocompleteResult = parseYqlQueryWithCursor(
+        '$f = ($y, $z) -> { $prefix = "x"; RETURN | ;}; select ',
+    );
+    const variablesSuggestions = ['y', 'z', 'prefix'];
+
+    expect(autocompleteResult.suggestVariables).toEqual(variablesSuggestions);
+});
+
+test('should suggest variables outside lambda', () => {
+    const autocompleteResult = parseYqlQueryWithCursor(
+        '$foo = "a"; \n$f = ($y) -> { $prefix = "x"; RETURN $prefix;}; select |',
+    );
+    const variablesSuggestions = ['foo', 'f'];
+
+    expect(autocompleteResult.suggestVariables).toEqual(variablesSuggestions);
+});

src/autocomplete/databases/yql/types.ts

@@ -37,6 +37,7 @@ export interface InternalSuggestions
     shouldSuggestColumns?: boolean;
     shouldSuggestAllColumns?: boolean;
     shouldSuggestColumnAliases?: boolean;
+    shouldSuggestVariables?: boolean;
 }
 
 export type YQLEntity =
@@ -66,6 +67,7 @@ export interface YqlAutocompleteResult extends Omit<SqlAutocompleteResult, 'sugg
     suggestTableHints?: string;
     suggestEntitySettings?: YQLEntity;
     suggestColumns?: YQLColumnsSuggestion;
+    suggestVariables?: string[];
 }
 
 export interface YqlTokenizeResult extends TokenizeResult {}

src/autocomplete/databases/yql/yql-autocomplete.ts

@@ -3,9 +3,14 @@ import * as c3 from 'antlr4-c3';
 
 import {YQLLexer} from './generated/YQLLexer';
 import {
+    Action_or_subquery_argsContext,
     Alter_table_store_stmtContext,
+    Declare_stmtContext,
+    Define_action_or_subquery_stmtContext,
+    LambdaContext,
     Named_columnContext,
     Named_exprContext,
+    Named_nodes_stmtContext,
     Named_single_sourceContext,
     Result_columnContext,
     Simple_table_ref_coreContext,
@@ -29,6 +34,7 @@ import {isStartingToWriteRule} from '../../shared/cursor.js';
 import {shouldSuggestTemplates} from '../../shared/query.js';
 import {EntitySuggestionToYqlEntity, getGranularSuggestions, tokenDictionary} from './helpers';
 import {EntitySuggestion, InternalSuggestions, YqlAutocompleteResult} from './types';
+import {getVariableSuggestions} from '../../shared/variables';
 
 // These are keywords that we do not want to show in autocomplete
 function getIgnoredTokens(): number[] {
@@ -86,6 +92,144 @@ const rulesToVisit = new Set([
     YQLParser.RULE_id_as_compat,
 ]);
 
+class YQLVariableSymbolTableVisitor extends YQLVisitor<{}> implements ISymbolTableVisitor {
+    symbolTable: c3.SymbolTable;
+    scope: c3.ScopedSymbol;
+
+    constructor() {
+        super();
+        this.symbolTable = new c3.SymbolTable('', {allowDuplicateSymbols: true});
+        this.scope = this.symbolTable.addNewSymbolOfType(c3.ScopedSymbol, undefined);
+    }
+
+    addVariableSymbol = (getVariable: (index: number) => string | undefined): void => {
+        try {
+            let index: number | null = 0;
+            while (index !== null) {
+                const variable = getVariable(index);
+                if (variable) {
+                    this.symbolTable.addNewSymbolOfType(
+                        c3.VariableSymbol,
+                        this.scope,
+                        variable,
+                        undefined,
+                    );
+                    index++;
+                } else {
+                    index = null;
+                }
+            }
+        } catch (error) {
+            if (!(error instanceof c3.DuplicateSymbolError)) {
+                throw error;
+            }
+        }
+    };
+
+    visitDeclare_stmt = (context: Declare_stmtContext): {} => {
+        try {
+            const variable = context.bind_parameter()?.an_id_or_type()?.getText();
+            if (variable) {
+                const value = context.literal_value()?.getText();
+
+                this.symbolTable.addNewSymbolOfType(c3.VariableSymbol, this.scope, variable, value);
+            }
+        } catch (error) {
+            if (!(error instanceof c3.DuplicateSymbolError)) {
+                throw error;
+            }
+        }
+
+        return this.visitChildren(context) as {};
+    };
+
+    visitAction_or_subquery_args = (context: Action_or_subquery_argsContext): {} => {
+        this.addVariableSymbol(
+            (index: number) =>
+                context.opt_bind_parameter(index)?.bind_parameter()?.an_id_or_type()?.getText(),
+        );
+        return this.visitChildren(context) as {};
+    };
+
+    visitNamed_nodes_stmt = (context: Named_nodes_stmtContext): {} => {
+        this.addVariableSymbol(
+            (index: number) =>
+                context.bind_parameter_list()?.bind_parameter(index)?.an_id_or_type()?.getText(),
+        );
+
+        return this.visitChildren(context) as {};
+    };
+
+    visitDefine_action_or_subquery_stmt = (context: Define_action_or_subquery_stmtContext): {} => {
+        try {
+            // this variable should be in global scope
+            const variable = context.bind_parameter()?.an_id_or_type()?.getText();
+            if (variable) {
+                this.symbolTable.addNewSymbolOfType(
+                    c3.VariableSymbol,
+                    this.scope,
+                    variable,
+                    undefined,
+                );
+            }
+        } catch (error) {
+            if (!(error instanceof c3.DuplicateSymbolError)) {
+                throw error;
+            }
+        }
+
+        return (
+            this.withScope(
+                context,
+                c3.RoutineSymbol,
+                [context.bind_parameter()?.an_id_or_type()?.getText()],
+                () => this.visitChildren(context),
+            ) ?? {}
+        );
+    };
+
+    visitLambda = (context: LambdaContext): {} => {
+        // this variable should be in local scope, so it should be extracted inside withScope callback
+        const addVariables = (): {} => {
+            const lambdaArgs = context.smart_parenthesis()?.named_expr_list();
+            this.addVariableSymbol((index: number) => {
+                const variable = lambdaArgs?.named_expr(index)?.expr()?.getText();
+                if (variable) {
+                    if (variable.startsWith('$')) {
+                        return variable.slice(1);
+                    }
+                }
+                return variable;
+            });
+            return this.visitChildren(context) as {};
+        };
+
+        return this.withScope(context, c3.RoutineSymbol, [context.getText()], addVariables) ?? {};
+    };
+
+    withScope<T>(
+        tree: ParseTree,
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        type: new (...args: any[]) => c3.ScopedSymbol,
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        args: any[],
+        action: () => T,
+    ): T {
+        const scope = this.symbolTable.addNewSymbolOfType(type, this.scope, ...args);
+        scope.context = tree;
+        this.scope = scope;
+        try {
+            return action();
+        } finally {
+            this.scope = scope.parent as c3.ScopedSymbol;
+        }
+    }
+
+    protected defaultResult(): c3.SymbolTable {
+        return this.symbolTable;
+    }
+}
+
 class YQLSymbolTableVisitor extends YQLVisitor<{}> implements ISymbolTableVisitor {
     symbolTable: c3.SymbolTable;
     scope: c3.ScopedSymbol;
@@ -302,6 +446,7 @@ function getEnrichAutocompleteResult(parseTreeGetter: GetParseTree<YQLParser>) {
             shouldSuggestAllColumns,
             shouldSuggestColumnAliases,
             shouldSuggestTableIndexes,
+            shouldSuggestVariables,
             ...suggestionsFromRules
         } = processVisitedRules(rules, cursorTokenIndex, tokenStream);
         const suggestTemplates = shouldSuggestTemplates(query, cursor);
@@ -313,6 +458,22 @@ function getEnrichAutocompleteResult(parseTreeGetter: GetParseTree<YQLParser>) {
         const contextSuggestionsNeeded =
             shouldSuggestColumns || shouldSuggestColumnAliases || shouldSuggestTableIndexes;
 
+        if (shouldSuggestVariables) {
+            const visitor = new YQLVariableSymbolTableVisitor();
+            const data = getVariableSuggestions(
+                YQLLexer,
+                YQLParser,
+                visitor,
+                parseTreeGetter,
+                tokenStream,
+                cursor,
+                query,
+            );
+            if (data.length) {
+                result.suggestVariables = data;
+            }
+        }
+
         if (contextSuggestionsNeeded) {
             const visitor = new YQLSymbolTableVisitor();
             const {tableContextSuggestion, suggestColumnAliases} = getContextSuggestions(

src/autocomplete/shared/autocomplete-types.ts

@@ -95,6 +95,7 @@ export type ProcessVisitedRulesResult<A extends AutocompleteResultBase> = Partia
     shouldSuggestColumnAliases?: boolean;
     shouldSuggestConstraints?: boolean;
     shouldSuggestTableIndexes?: boolean;
+    shouldSuggestVariables?: boolean;
 };
 
 export type ProcessVisitedRules<A extends AutocompleteResultBase> = (