From 5409bde13dfbf42e0d4cdee04a0f37429e9102cc Mon Sep 17 00:00:00 2001 From: Elena Makarova Date: Thu, 12 Dec 2024 11:16:25 +0300 Subject: [PATCH] feat(YQL): support variables and scopes --- src/autocomplete/databases/yql/helpers.ts | 14 ++ .../yql/tests/yql/select/select.test.ts | 59 +++++- src/autocomplete/databases/yql/types.ts | 2 + .../databases/yql/yql-autocomplete.ts | 188 ++++++++++++++++++ src/autocomplete/shared/autocomplete-types.ts | 1 + .../shared/compute-token-position.ts | 122 ++++++++++++ src/autocomplete/shared/symbol-table.ts | 16 ++ src/autocomplete/shared/variables.ts | 53 +++++ 8 files changed, 454 insertions(+), 1 deletion(-) create mode 100644 src/autocomplete/shared/compute-token-position.ts create mode 100644 src/autocomplete/shared/variables.ts 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); + } + + 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): {} => { + try { + let index: number | null = 0; + while (index !== null) { + const variable = context + .opt_bind_parameter(index) + ?.bind_parameter() + ?.an_id_or_type() + ?.getText(); + if (variable) { + this.symbolTable.addNewSymbolOfType( + c3.VariableSymbol, + this.scope, + variable, + undefined, + ); + index++; + } else { + index = null; + } + } + } catch (error) { + if (!(error instanceof c3.DuplicateSymbolError)) { + throw error; + } + } + + return this.visitChildren(context) as {}; + }; + visitNamed_nodes_stmt = (context: Named_nodes_stmtContext): {} => { + try { + let index: number | null = 0; + while (index !== null) { + const variable = context + .bind_parameter_list() + ?.bind_parameter(index) + ?.an_id_or_type() + ?.getText(); + if (variable) { + this.symbolTable.addNewSymbolOfType( + c3.VariableSymbol, + this.scope, + variable, + undefined, + ); + index++; + } else { + index = null; + } + } + } catch (error) { + if (!(error instanceof c3.DuplicateSymbolError)) { + throw error; + } + } + + 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 callback = (): {} => { + try { + const lambdaArgs = context.smart_parenthesis()?.named_expr_list(); + + let index: number | null = 0; + while (index !== null) { + const variable = lambdaArgs?.named_expr(index)?.expr()?.getText(); + if (variable) { + if (variable.startsWith('$')) { + this.symbolTable.addNewSymbolOfType( + c3.VariableSymbol, + this.scope, + variable.slice(1), + undefined, + ); + } + + index++; + } else { + index = null; + } + } + } catch (error) { + if (!(error instanceof c3.DuplicateSymbolError)) { + throw error; + } + } + return this.visitChildren(context) as {}; + }; + + return this.withScope(context, c3.RoutineSymbol, [context.getText()], callback) ?? {}; + }; + + withScope( + tree: ParseTree, + type: new (...args: any[]) => c3.ScopedSymbol, + 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 +474,7 @@ function getEnrichAutocompleteResult(parseTreeGetter: GetParseTree) { shouldSuggestAllColumns, shouldSuggestColumnAliases, shouldSuggestTableIndexes, + shouldSuggestVariables, ...suggestionsFromRules } = processVisitedRules(rules, cursorTokenIndex, tokenStream); const suggestTemplates = shouldSuggestTemplates(query, cursor); @@ -313,6 +486,21 @@ function getEnrichAutocompleteResult(parseTreeGetter: GetParseTree) { const contextSuggestionsNeeded = shouldSuggestColumns || shouldSuggestColumnAliases || shouldSuggestTableIndexes; + if (shouldSuggestVariables) { + const visitor = new YQLSymbolTableVisitor2(); + const data = getVariablesSuggestions( + 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..e7facbdc --- /dev/null +++ b/src/autocomplete/shared/compute-token-position.ts @@ -0,0 +1,122 @@ +import {ParseTree, ParserRuleContext, TerminalNode, Token, TokenStream} from 'antlr4ng'; + +import {CursorPosition} from './autocomplete-types'; + +type TokenPosition = {index: number; context: ParseTree; text: string}; + +export function tokenPositionComputer(identifierTokenTypes: number[] = []) { + return (parseTree: ParseTree, tokens: TokenStream, caretPosition: CursorPosition) => + computeTokenPosition(parseTree, tokens, caretPosition, identifierTokenTypes); +} + +export function computeTokenPosition( + parseTree: ParseTree | null, + tokens: TokenStream, + caretPosition: CursorPosition, + identifierTokenTypes: number[] = [], +): TokenPosition | undefined { + if (!parseTree) { + return; + } + if (parseTree instanceof TerminalNode) { + return computeTokenPositionOfTerminal( + parseTree, + tokens, + caretPosition, + identifierTokenTypes, + ); + } else { + return computeTokenPositionOfChildNode( + parseTree as ParserRuleContext, + tokens, + caretPosition, + identifierTokenTypes, + ); + } +} + +function positionOfToken( + token: Token, + text: string, + caretPosition: CursorPosition, + identifierTokenTypes: number[], + parseTree: ParseTree, +): TokenPosition | undefined { + const start = token.column; + const stop = token.column + text.length; + if ( + token.line === caretPosition.line && + start <= caretPosition.column && + stop >= caretPosition.column + ) { + let index = token.tokenIndex; + if (identifierTokenTypes.includes(token.type)) { + index--; + } + return { + index, + context: parseTree, + text: text.substring(0, caretPosition.column - start), + }; + //it means token is eof + } else if (token.start > token.stop) { + return { + index: token.tokenIndex, + context: parseTree, + text: text.substring(0, caretPosition.column), + }; + } else { + return undefined; + } +} + +function computeTokenPositionOfTerminal( + parseTree: TerminalNode, + _tokens: TokenStream, + caretPosition: CursorPosition, + identifierTokenTypes: number[], +): TokenPosition | undefined { + const token = parseTree.symbol; + const text = parseTree.getText(); + return positionOfToken(token, text, caretPosition, identifierTokenTypes, parseTree); +} + +function computeTokenPositionOfChildNode( + parseTree: ParserRuleContext, + tokens: TokenStream, + caretPosition: CursorPosition, + identifierTokenTypes: number[], +): TokenPosition | undefined { + if ( + (parseTree.start && parseTree.start.line > caretPosition.line) || + (parseTree.stop && parseTree.stop.line < caretPosition.line) + ) { + return undefined; + } + for (let i = 0; i < parseTree.getChildCount(); i++) { + const position = computeTokenPosition( + parseTree.getChild(i), + tokens, + caretPosition, + identifierTokenTypes, + ); + if (position !== undefined) { + return position; + } + } + if (parseTree.start && parseTree.stop) { + for (let i = parseTree.start.tokenIndex; i <= parseTree.stop.tokenIndex; i++) { + const pos = positionOfToken( + tokens.get(i), + tokens.get(i).text ?? '', + caretPosition, + identifierTokenTypes, + parseTree, + ); + if (pos) { + return pos; + } + } + } + return undefined; +} diff --git a/src/autocomplete/shared/symbol-table.ts b/src/autocomplete/shared/symbol-table.ts index f7208bad..d1b3b225 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,18 @@ export function getColumnAliasesFromSymbolTable( .getNestedSymbolsOfTypeSync(ColumnAliasSymbol) .map(({name}) => ({name})); } + +export function getScope( + context: ParseTree | null | undefined, + symbolTable: c3.SymbolTable | null, +): c3.BaseSymbol | undefined { + if (!context || !symbolTable) { + return undefined; + } + const scope = symbolTable.symbolWithContextSync(context); + if (scope) { + return scope; + } else { + 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..d6da6877 --- /dev/null +++ b/src/autocomplete/shared/variables.ts @@ -0,0 +1,53 @@ +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 {computeTokenPosition} from './compute-token-position'; + +export function getVariablesSuggestions( + 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 tokenPosion = computeTokenPosition(parseTree, tokenStream, cursor); + + if (!tokenPosion) { + throw new Error(`Could not find tokenPosion at Ln ${cursor.line}, Col ${cursor.column}`); + } + + const symbolTable: c3.SymbolTable | null = symbolVariableVisitor.visit(parseTree) as any; + + const variables = suggestVariables(symbolTable, tokenPosion.context); + + return variables; +} + +function suggestVariables(symbolTable: c3.SymbolTable | null, context?: ParseTree): string[] { + const scope = getScope(context, symbolTable); + let symbols: c3.VariableSymbol[] = []; + if (scope instanceof c3.ScopedSymbol) { + //Local scope + symbols = scope.getNestedSymbolsOfTypeSync(c3.VariableSymbol); + } else if (symbolTable) { + //Global scope + 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((sym) => !sym.parent?.context); + } + return symbols.map((s) => s.name); +}