Skip to content

Commit

Permalink
feat!: add table context to constraint suggestions (#159)
Browse files Browse the repository at this point in the history
* feat!: add table context to constraint suggestions

* refactor: create enhanceAutocompleteResult function and split out logic

* refactor: naming
  • Loading branch information
roberthovsepyan authored Feb 28, 2024
1 parent ce056c5 commit 23ee8dd
Show file tree
Hide file tree
Showing 18 changed files with 516 additions and 244 deletions.
42 changes: 30 additions & 12 deletions src/autocomplete/autocomplete-types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -27,13 +27,13 @@ export interface AutocompleteResultBase {
export interface MySqlAutocompleteResult extends AutocompleteResultBase {
suggestIndexes?: boolean;
suggestTriggers?: boolean;
suggestConstraints?: boolean;
suggestConstraints?: ConstraintSuggestion;
}

export interface PostgreSqlAutocompleteResult extends AutocompleteResultBase {
suggestIndexes?: boolean;
suggestTriggers?: boolean;
suggestConstraints?: boolean;
suggestConstraints?: ConstraintSuggestion;
suggestSequences?: boolean;
suggestSchemas?: boolean;
}
Expand All @@ -50,10 +50,19 @@ export interface KeywordSuggestion {
value: string;
}

export interface ColumnSuggestion {
tables?: {name: string; alias?: string}[];
export interface Table {
name: string;
alias?: string;
}

export interface TableContextSuggestion {
tables?: Table[];
}

export type ColumnSuggestion = TableContextSuggestion;

export type ConstraintSuggestion = TableContextSuggestion;

export enum TableOrViewSuggestion {
ALL = 'ALL',
TABLES = 'TABLES',
Expand All @@ -80,36 +89,45 @@ export interface ISymbolTableVisitor {
scope: c3.ScopedSymbol;
}

export type SymbolTableVisitor = ISymbolTableVisitor & AbstractParseTreeVisitor<{}>;

export type GetParseTree<P> = (
parser: P,
type?: TableQueryPosition['type'] | 'select',
) => ParseTree;

export type GenerateSuggestionsFromRulesResult<A extends AutocompleteResultBase> = Partial<A> & {
export type ProcessVisitedRulesResult<A extends AutocompleteResultBase> = Partial<A> & {
shouldSuggestColumns?: boolean;
shouldSuggestColumnAliases?: boolean;
shouldSuggestConstraints?: boolean;
};
export type GenerateSuggestionsFromRules<A extends AutocompleteResultBase> = (
export type ProcessVisitedRules<A extends AutocompleteResultBase> = (
rules: c3.CandidatesCollection['rules'],
cursorTokenIndex: number,
tokenStream: TokenStream,
) => GenerateSuggestionsFromRulesResult<A>;
) => ProcessVisitedRulesResult<A>;

export type EnrichAutocompleteResult<A extends AutocompleteResultBase> = (
result: AutocompleteResultBase,
rules: c3.CandidatesCollection['rules'],
tokenStream: TokenStream,
cursorTokenIndex: number,
cursor: CursorPosition,
query: string,
) => A;

export interface AutocompleteData<
A extends AutocompleteResultBase,
L extends LexerType,
P extends ParserType,
S extends ISymbolTableVisitor & AbstractParseTreeVisitor<{}>,
> {
Lexer: LexerConstructor<L>;
Parser: ParserConstructor<P>;
SymbolTableVisitor: SymbolTableVisitorConstructor<S>;
getParseTree: GetParseTree<P>;
tokenDictionary: TokenDictionary;
generateSuggestionsFromRules: GenerateSuggestionsFromRules<A>;
ignoredTokens: Set<number>;
preferredRules: Set<number>;
explicitlyParseJoin: boolean;
rulesToVisit: Set<number>;
enrichAutocompleteResult: EnrichAutocompleteResult<A>;
}

export interface CursorPosition {
Expand Down
199 changes: 24 additions & 175 deletions src/autocomplete/autocomplete.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,33 +2,23 @@ import {
AutocompleteResultBase,
ClickHouseAutocompleteResult,
CursorPosition,
GenerateSuggestionsFromRules,
EnrichAutocompleteResult,
GetParseTree,
ISymbolTableVisitor,
KeywordSuggestion,
LexerConstructor,
MySqlAutocompleteResult,
ParserConstructor,
PostgreSqlAutocompleteResult,
SymbolTableVisitorConstructor,
} from './autocomplete-types';
import {postgreSqlAutocompleteData} from './databases/postgresql/postgresql-autocomplete';
import {mySqlAutocompleteData} from './databases/mysql/mysql-autocomplete';
import {
AbstractParseTreeVisitor,
CharStreams,
CommonTokenStream,
Lexer as LexerType,
Parser as ParserType,
TokenStream,
} from 'antlr4ng';
import {TokenDictionary, getTableQueryPosition} from './shared/tables';
import {Lexer as LexerType, Parser as ParserType} from 'antlr4ng';
import {TokenDictionary} from './shared/tables';
import {createParser} from './shared/query';
import {SqlErrorListener} from './shared/sql-error-listener';
import * as c3 from 'antlr4-c3';
import {findCursorTokenIndex, getCursorIndex} from './shared/cursor';
import {getCurrentStatement, shouldSuggestTemplates} from './shared/query';
import {findCursorTokenIndex} from './shared/cursor';
import {clickHouseAutocompleteData} from './databases/clickhouse/clickhouse-autocomplete';
import {ColumnAliasSymbol, TableSymbol} from './shared/symbol-table';

function parseQueryWithoutCursor<L extends LexerType, P extends ParserType>(
Lexer: LexerConstructor<L>,
Expand All @@ -37,10 +27,7 @@ function parseQueryWithoutCursor<L extends LexerType, P extends ParserType>(
getParseTree: GetParseTree<P>,
query: string,
): Pick<AutocompleteResultBase, 'errors'> {
const inputStream = CharStreams.fromString(query);
const lexer = new Lexer(inputStream);
const tokenStream = new CommonTokenStream(lexer);
const parser = new Parser(tokenStream);
const parser = createParser(Lexer, Parser, query);
const errorListener = new SqlErrorListener(tokenDictionary.SPACE);

parser.removeErrorListeners();
Expand All @@ -50,130 +37,25 @@ function parseQueryWithoutCursor<L extends LexerType, P extends ParserType>(
return {errors: errorListener.errors};
}

function getColumnSuggestions<
L extends LexerType,
P extends ParserType,
S extends ISymbolTableVisitor & AbstractParseTreeVisitor<{}>,
>(
Lexer: LexerConstructor<L>,
Parser: ParserConstructor<P>,
SymbolTableVisitor: SymbolTableVisitorConstructor<S>,
tokenDictionary: TokenDictionary,
explicitlyParseJoin: boolean,
getParseTree: GetParseTree<P>,
initialTokenStream: TokenStream,
cursor: CursorPosition,
initialQuery: string,
shouldSuggestColumnAliases?: boolean,
): Pick<AutocompleteResultBase, 'suggestColumns' | 'suggestColumnAliases'> {
// Here we need the actual token index, without special logic for spaces
const realCursorTokenIndex = findCursorTokenIndex(
initialTokenStream,
cursor,
tokenDictionary.SPACE,
true,
);

if (!realCursorTokenIndex) {
throw new Error(
`Could not find realCursorTokenIndex at Ln ${cursor.line}, Col ${cursor.column}`,
);
}

const tableQueryPosition = getTableQueryPosition(
initialTokenStream,
realCursorTokenIndex,
tokenDictionary,
);
const result: Pick<AutocompleteResultBase, 'suggestColumns' | 'suggestColumnAliases'> = {};

if (tableQueryPosition) {
const query = initialQuery.slice(tableQueryPosition.start, tableQueryPosition.end);

const inputStream = CharStreams.fromString(query);
const lexer = new Lexer(inputStream);
const tokenStream = new CommonTokenStream(lexer);
const parser = new Parser(tokenStream);

parser.removeErrorListeners();
const parseTree = getParseTree(parser, tableQueryPosition.type);
const visitor = new SymbolTableVisitor();

visitor.visit(parseTree);

if (explicitlyParseJoin && tableQueryPosition.joinTableQueryPosition) {
const joinTableQuery = initialQuery.slice(
tableQueryPosition.joinTableQueryPosition.start,
tableQueryPosition.joinTableQueryPosition.end,
);
const joinInputStream = CharStreams.fromString(joinTableQuery);
const joinLexer = new Lexer(joinInputStream);
const joinTokenStream = new CommonTokenStream(joinLexer);
const joinParser = new Parser(joinTokenStream);

joinParser.removeErrorListeners();
const joinParseTree = getParseTree(joinParser, 'from');
visitor.visit(joinParseTree);
}

if (shouldSuggestColumnAliases && tableQueryPosition.selectTableQueryPosition) {
const selectTableQuery = initialQuery.slice(
tableQueryPosition.selectTableQueryPosition.start,
tableQueryPosition.selectTableQueryPosition.end,
);
const selectInputStream = CharStreams.fromString(selectTableQuery);
const selectLexer = new Lexer(selectInputStream);
const selectTokenStream = new CommonTokenStream(selectLexer);
const selectParser = new Parser(selectTokenStream);

selectParser.removeErrorListeners();
const selectParseTree = getParseTree(selectParser, 'select');
visitor.visit(selectParseTree);
}

const tables = visitor.symbolTable.getNestedSymbolsOfTypeSync(TableSymbol);
if (tables.length) {
result.suggestColumns = {
tables: tables.map((tableSymbol) => ({
name: tableSymbol.name,
alias: tableSymbol.alias,
})),
};
}

const columnAliases = visitor.symbolTable.getNestedSymbolsOfTypeSync(ColumnAliasSymbol);
if (columnAliases.length) {
result.suggestColumnAliases = columnAliases.map(({name}) => ({name}));
}
}

return result;
}

const quotesRegex = /^'(.*)'$/;

export function parseQuery<
A extends AutocompleteResultBase,
L extends LexerType,
P extends ParserType,
S extends ISymbolTableVisitor & AbstractParseTreeVisitor<{}>,
>(
Lexer: LexerConstructor<L>,
Parser: ParserConstructor<P>,
SymbolTableVisitor: SymbolTableVisitorConstructor<S>,
tokenDictionary: TokenDictionary,
ignoredTokens: Set<number>,
preferredRules: Set<number>,
explicitlyParseJoin: boolean,
rulesToVisit: Set<number>,
getParseTree: GetParseTree<P>,
generateSuggestionsFromRules: GenerateSuggestionsFromRules<A>,
enrichAutocompleteResult: EnrichAutocompleteResult<A>,
query: string,
cursor: CursorPosition,
): AutocompleteResultBase {
const inputStream = CharStreams.fromString(query);
const lexer = new Lexer(inputStream);
const tokenStream = new CommonTokenStream(lexer);
const parser = new Parser(tokenStream);
): A {
const parser = createParser(Lexer, Parser, query);
const {tokenStream} = parser;
const errorListener = new SqlErrorListener(tokenDictionary.SPACE);

parser.removeErrorListeners();
Expand All @@ -182,23 +64,17 @@ export function parseQuery<

const core = new c3.CodeCompletionCore(parser);
core.ignoredTokens = ignoredTokens;
core.preferredRules = preferredRules;
core.preferredRules = rulesToVisit;
const cursorTokenIndex = findCursorTokenIndex(tokenStream, cursor, tokenDictionary.SPACE);
const suggestKeywords: KeywordSuggestion[] = [];
let result: AutocompleteResultBase = {
errors: errorListener.errors,
};

if (cursorTokenIndex === undefined) {
throw new Error(
`Could not find cursor token index for line: ${cursor.line}, column: ${cursor.column}`,
);
}

const suggestKeywords: KeywordSuggestion[] = [];
const {tokens, rules} = core.collectCandidates(cursorTokenIndex);
const {shouldSuggestColumns, shouldSuggestColumnAliases, ...suggestionsFromRules} =
generateSuggestionsFromRules(rules, cursorTokenIndex, tokenStream);
result = {...result, ...suggestionsFromRules};
tokens.forEach((_, tokenType) => {
// Literal keyword names are quoted
const literalName = parser.vocabulary.getLiteralName(tokenType)?.replace(quotesRegex, '$1');
Expand All @@ -214,33 +90,12 @@ export function parseQuery<
});
});

const cursorIndex = getCursorIndex(query, cursor);
// We can get this by token instead of splitting the string
const currentStatement = getCurrentStatement(query, cursorIndex);

if (shouldSuggestColumns) {
const {suggestColumns, suggestColumnAliases} = getColumnSuggestions(
Lexer,
Parser,
SymbolTableVisitor,
tokenDictionary,
explicitlyParseJoin,
getParseTree,
tokenStream,
cursor,
query,
shouldSuggestColumnAliases,
);
result.suggestColumns = suggestColumns;
result.suggestColumnAliases = suggestColumnAliases;
}
const result: AutocompleteResultBase = {
errors: errorListener.errors,
suggestKeywords,
};

result.suggestTemplates = shouldSuggestTemplates(
currentStatement.statement,
currentStatement.cursorIndex,
);
result.suggestKeywords = suggestKeywords;
return result;
return enrichAutocompleteResult(result, rules, tokenStream, cursorTokenIndex, cursor, query);
}

export function parseMySqlQueryWithoutCursor(
Expand All @@ -259,13 +114,11 @@ export function parseMySqlQuery(query: string, cursor: CursorPosition): MySqlAut
return parseQuery(
mySqlAutocompleteData.Lexer,
mySqlAutocompleteData.Parser,
mySqlAutocompleteData.SymbolTableVisitor,
mySqlAutocompleteData.tokenDictionary,
mySqlAutocompleteData.ignoredTokens,
mySqlAutocompleteData.preferredRules,
mySqlAutocompleteData.explicitlyParseJoin,
mySqlAutocompleteData.rulesToVisit,
mySqlAutocompleteData.getParseTree,
mySqlAutocompleteData.generateSuggestionsFromRules,
mySqlAutocompleteData.enrichAutocompleteResult,
query,
cursor,
);
Expand All @@ -290,13 +143,11 @@ export function parsePostgreSqlQuery(
return parseQuery(
postgreSqlAutocompleteData.Lexer,
postgreSqlAutocompleteData.Parser,
postgreSqlAutocompleteData.SymbolTableVisitor,
postgreSqlAutocompleteData.tokenDictionary,
postgreSqlAutocompleteData.ignoredTokens,
postgreSqlAutocompleteData.preferredRules,
postgreSqlAutocompleteData.explicitlyParseJoin,
postgreSqlAutocompleteData.rulesToVisit,
postgreSqlAutocompleteData.getParseTree,
postgreSqlAutocompleteData.generateSuggestionsFromRules,
postgreSqlAutocompleteData.enrichAutocompleteResult,
query,
cursor,
);
Expand All @@ -321,13 +172,11 @@ export function parseClickHouseQuery(
return parseQuery(
clickHouseAutocompleteData.Lexer,
clickHouseAutocompleteData.Parser,
clickHouseAutocompleteData.SymbolTableVisitor,
clickHouseAutocompleteData.tokenDictionary,
clickHouseAutocompleteData.ignoredTokens,
clickHouseAutocompleteData.preferredRules,
clickHouseAutocompleteData.explicitlyParseJoin,
clickHouseAutocompleteData.rulesToVisit,
clickHouseAutocompleteData.getParseTree,
clickHouseAutocompleteData.generateSuggestionsFromRules,
clickHouseAutocompleteData.enrichAutocompleteResult,
query,
cursor,
);
Expand Down
Loading

0 comments on commit 23ee8dd

Please sign in to comment.