Skip to content

Commit

Permalink
feat(YQL): support variables and scopes (#258)
Browse files Browse the repository at this point in the history 
* feat(YQL): support variables and scopes

* fix: review

* fix: spaces

* fix: review

* fix: review
  • Loading branch information
@Raubzeug
Raubzeug authored Dec 13, 2024
1 parent e5caf9c commit 2406c4d
Show file tree
Hide file tree
Showing 8 changed files with 422 additions and 1 deletion.
14 changes: 14 additions & 0 deletions src/autocomplete/databases/yql/helpers.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
59 changes: 58 additions & 1 deletion src/autocomplete/databases/yql/tests/yql/select/select.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand All @@ -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);
});
2 changes: 2 additions & 0 deletions src/autocomplete/databases/yql/types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ export interface InternalSuggestions
shouldSuggestColumns?: boolean;
shouldSuggestAllColumns?: boolean;
shouldSuggestColumnAliases?: boolean;
shouldSuggestVariables?: boolean;
}

export type YQLEntity =
Expand Down Expand Up @@ -66,6 +67,7 @@ export interface YqlAutocompleteResult extends Omit<SqlAutocompleteResult, 'sugg
suggestTableHints?: string;
suggestEntitySettings?: YQLEntity;
suggestColumns?: YQLColumnsSuggestion;
suggestVariables?: string[];
}

export interface YqlTokenizeResult extends TokenizeResult {}
161 changes: 161 additions & 0 deletions src/autocomplete/databases/yql/yql-autocomplete.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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[] {
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -302,6 +446,7 @@ function getEnrichAutocompleteResult(parseTreeGetter: GetParseTree<YQLParser>) {
shouldSuggestAllColumns,
shouldSuggestColumnAliases,
shouldSuggestTableIndexes,
shouldSuggestVariables,
...suggestionsFromRules
} = processVisitedRules(rules, cursorTokenIndex, tokenStream);
const suggestTemplates = shouldSuggestTemplates(query, cursor);
Expand All @@ -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(
Expand Down
1 change: 1 addition & 0 deletions src/autocomplete/shared/autocomplete-types.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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> = (
Expand Down
Loading

0 comments on commit 2406c4d

Please sign in to comment.