diff --git a/src/autocomplete/databases/yql/helpers.ts b/src/autocomplete/databases/yql/helpers.ts index 4418d0ed..dc984127 100644 --- a/src/autocomplete/databases/yql/helpers.ts +++ b/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,6 +516,7 @@ export function getGranularSuggestions( const suggestAggregateFunctions = getAggregateFunctionsSuggestions(props); const shouldSuggestTableHints = checkShouldSuggestTableHints(props); const suggestEntitySettings = getEntitySettingsSuggestions(props); + const shouldSuggestVariables = checkShouldSuggestVariables(props); return { suggestWindowFunctions, @@ -511,6 +524,7 @@ export function getGranularSuggestions( shouldSuggestColumns, shouldSuggestAllColumns, shouldSuggestColumnAliases: shouldSuggestColumns, + shouldSuggestVariables, suggestSimpleTypes, suggestPragmas, suggestUdfs, diff --git a/src/autocomplete/databases/yql/tests/yql/select/select.test.ts b/src/autocomplete/databases/yql/tests/yql/select/select.test.ts index 865276b8..6a733495 100644 --- a/src/autocomplete/databases/yql/tests/yql/select/select.test.ts +++ b/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); +}); diff --git a/src/autocomplete/databases/yql/types.ts b/src/autocomplete/databases/yql/types.ts index c0e901ac..bfaabc28 100644 --- a/src/autocomplete/databases/yql/types.ts +++ b/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 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( + 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) { shouldSuggestAllColumns, shouldSuggestColumnAliases, shouldSuggestTableIndexes, + shouldSuggestVariables, ...suggestionsFromRules } = processVisitedRules(rules, cursorTokenIndex, tokenStream); const suggestTemplates = shouldSuggestTemplates(query, cursor); @@ -313,6 +458,22 @@ function getEnrichAutocompleteResult(parseTreeGetter: GetParseTree) { 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( diff --git a/src/autocomplete/shared/autocomplete-types.ts b/src/autocomplete/shared/autocomplete-types.ts index c10bebc0..db8046e0 100644 --- a/src/autocomplete/shared/autocomplete-types.ts +++ b/src/autocomplete/shared/autocomplete-types.ts @@ -95,6 +95,7 @@ export type ProcessVisitedRulesResult = Partia shouldSuggestColumnAliases?: boolean; shouldSuggestConstraints?: boolean; shouldSuggestTableIndexes?: boolean; + shouldSuggestVariables?: boolean; }; export type ProcessVisitedRules = ( diff --git a/src/autocomplete/shared/compute-token-position.ts b/src/autocomplete/shared/compute-token-position.ts new file mode 100644 index 00000000..eee75caa --- /dev/null +++ b/src/autocomplete/shared/compute-token-position.ts @@ -0,0 +1,117 @@ +import {ParseTree, ParserRuleContext, TerminalNode, Token, TokenStream} from 'antlr4ng'; + +import {CursorPosition} from './autocomplete-types'; + +type TokenContext = {index: number; context: ParseTree; text: string}; + +export function computeTokenContext( + parseTree: ParseTree, + tokenStream: TokenStream, + cursorPosition: CursorPosition, + identifierTokenTypes: number[] = [], +): TokenContext | undefined { + if (parseTree instanceof TerminalNode) { + return computeTokenContextOfTerminalNode(parseTree, cursorPosition, identifierTokenTypes); + } + return computeTokenContextOfChildNode( + parseTree as ParserRuleContext, + tokenStream, + cursorPosition, + identifierTokenTypes, + ); +} + +function getTokenContext( + token: Token, + text: string, + cursorPosition: CursorPosition, + identifierTokenTypes: number[], + parseTree: ParseTree, +): TokenContext | undefined { + const start = token.column; + const stop = token.column + text.length; + + // It means token is eof + if (token.start > token.stop) { + return { + index: token.tokenIndex, + context: parseTree, + text: text.substring(0, cursorPosition.column), + }; + } + + if ( + token.line === cursorPosition.line && + start <= cursorPosition.column && + stop >= cursorPosition.column + ) { + let index = token.tokenIndex; + if (identifierTokenTypes.includes(token.type)) { + index--; + } + return { + index, + context: parseTree, + text: text.substring(0, cursorPosition.column - start), + }; + } + + return undefined; +} + +function computeTokenContextOfTerminalNode( + parseTree: TerminalNode, + cursorPosition: CursorPosition, + identifierTokenTypes: number[], +): TokenContext | undefined { + const token = parseTree.symbol; + const text = parseTree.getText(); + return getTokenContext(token, text, cursorPosition, identifierTokenTypes, parseTree); +} + +function computeTokenContextOfChildNode( + parseTree: ParserRuleContext, + tokenStream: TokenStream, + cursorPosition: CursorPosition, + identifierTokenTypes: number[], +): TokenContext | undefined { + if ( + !parseTree.start || + !parseTree.stop || + parseTree.start.line > cursorPosition.line || + parseTree.stop.line < cursorPosition.line + ) { + return undefined; + } + + for (let i = 0; i < parseTree.getChildCount(); i++) { + const child = parseTree.getChild(i); + if (!child) { + continue; + } + const tokenContext = computeTokenContext( + child, + tokenStream, + cursorPosition, + identifierTokenTypes, + ); + if (tokenContext) { + return tokenContext; + } + } + + for (let i = parseTree.start.tokenIndex; i <= parseTree.stop.tokenIndex; i++) { + const tokenContext = getTokenContext( + tokenStream.get(i), + tokenStream.get(i).text ?? '', + cursorPosition, + identifierTokenTypes, + parseTree, + ); + if (tokenContext) { + return tokenContext; + } + } + + return undefined; +} diff --git a/src/autocomplete/shared/symbol-table.ts b/src/autocomplete/shared/symbol-table.ts index f7208bad..8f018e83 100644 --- a/src/autocomplete/shared/symbol-table.ts +++ b/src/autocomplete/shared/symbol-table.ts @@ -1,5 +1,6 @@ import * as c3 from 'antlr4-c3'; import {ColumnAliasSuggestion, SymbolTableVisitor, Table} from './autocomplete-types'; +import {ParseTree} from 'antlr4ng'; export class TableSymbol extends c3.TypedSymbol { name: string; @@ -64,3 +65,17 @@ export function getColumnAliasesFromSymbolTable( .getNestedSymbolsOfTypeSync(ColumnAliasSymbol) .map(({name}) => ({name})); } + +export function getScope( + context: ParseTree, + symbolTable: c3.SymbolTable, +): c3.BaseSymbol | undefined { + const scope = symbolTable.symbolWithContextSync(context); + if (scope) { + return scope; + } + if (!context.parent) { + return undefined; + } + return getScope(context.parent, symbolTable); +} diff --git a/src/autocomplete/shared/variables.ts b/src/autocomplete/shared/variables.ts new file mode 100644 index 00000000..bc33cad6 --- /dev/null +++ b/src/autocomplete/shared/variables.ts @@ -0,0 +1,54 @@ +import * as c3 from 'antlr4-c3'; +import {Lexer as LexerType, ParseTree, Parser as ParserType, TokenStream} from 'antlr4ng'; +import {createParser} from './query'; +import { + CursorPosition, + GetParseTree, + LexerConstructor, + ParserConstructor, + SymbolTableVisitor, +} from './autocomplete-types'; +import {getScope} from './symbol-table'; +import {computeTokenContext} from './compute-token-position'; + +export function getVariableSuggestions( + Lexer: LexerConstructor, + Parser: ParserConstructor

, + symbolVariableVisitor: SymbolTableVisitor, + getParseTree: GetParseTree

, + tokenStream: TokenStream, + cursor: CursorPosition, + query: string, +): string[] { + const parser = createParser(Lexer, Parser, query); + const parseTree = getParseTree(parser); + + const tokenPosition = computeTokenContext(parseTree, tokenStream, cursor); + + if (!tokenPosition) { + throw new Error(`Could not find tokenPosition at Ln ${cursor.line}, Col ${cursor.column}`); + } + + symbolVariableVisitor.visit(parseTree); + + return suggestVariables(symbolVariableVisitor.symbolTable, tokenPosition.context); +} + +function suggestVariables(symbolTable: c3.SymbolTable, context: ParseTree): string[] { + const scope = getScope(context, symbolTable); + let symbols: c3.VariableSymbol[] = []; + + // Local scope + if (scope instanceof c3.ScopedSymbol) { + symbols = scope.getNestedSymbolsOfTypeSync(c3.VariableSymbol); + + // Global scope + } else if (symbolTable) { + symbols = symbolTable + .getNestedSymbolsOfTypeSync(c3.VariableSymbol) + // 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); + } + + return symbols.map((symbol) => symbol.name); +}