diff --git a/src/autocomplete/autocomplete-types.ts b/src/autocomplete/autocomplete-types.ts index fca53558..5dadca19 100644 --- a/src/autocomplete/autocomplete-types.ts +++ b/src/autocomplete/autocomplete-types.ts @@ -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; } @@ -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', @@ -80,36 +89,45 @@ export interface ISymbolTableVisitor { scope: c3.ScopedSymbol; } +export type SymbolTableVisitor = ISymbolTableVisitor & AbstractParseTreeVisitor<{}>; + export type GetParseTree

= ( parser: P, type?: TableQueryPosition['type'] | 'select', ) => ParseTree; -export type GenerateSuggestionsFromRulesResult = Partial & { +export type ProcessVisitedRulesResult = Partial & { shouldSuggestColumns?: boolean; shouldSuggestColumnAliases?: boolean; + shouldSuggestConstraints?: boolean; }; -export type GenerateSuggestionsFromRules = ( +export type ProcessVisitedRules = ( rules: c3.CandidatesCollection['rules'], cursorTokenIndex: number, tokenStream: TokenStream, -) => GenerateSuggestionsFromRulesResult; +) => ProcessVisitedRulesResult; + +export type EnrichAutocompleteResult = ( + 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; Parser: ParserConstructor

; - SymbolTableVisitor: SymbolTableVisitorConstructor; getParseTree: GetParseTree

; tokenDictionary: TokenDictionary; - generateSuggestionsFromRules: GenerateSuggestionsFromRules; ignoredTokens: Set; - preferredRules: Set; - explicitlyParseJoin: boolean; + rulesToVisit: Set; + enrichAutocompleteResult: EnrichAutocompleteResult; } export interface CursorPosition { diff --git a/src/autocomplete/autocomplete.ts b/src/autocomplete/autocomplete.ts index 5b46b60e..cb7ccc3d 100644 --- a/src/autocomplete/autocomplete.ts +++ b/src/autocomplete/autocomplete.ts @@ -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( Lexer: LexerConstructor, @@ -37,10 +27,7 @@ function parseQueryWithoutCursor( getParseTree: GetParseTree

, query: string, ): Pick { - 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(); @@ -50,130 +37,25 @@ function parseQueryWithoutCursor( return {errors: errorListener.errors}; } -function getColumnSuggestions< - L extends LexerType, - P extends ParserType, - S extends ISymbolTableVisitor & AbstractParseTreeVisitor<{}>, ->( - Lexer: LexerConstructor, - Parser: ParserConstructor

, - SymbolTableVisitor: SymbolTableVisitorConstructor, - tokenDictionary: TokenDictionary, - explicitlyParseJoin: boolean, - getParseTree: GetParseTree

, - initialTokenStream: TokenStream, - cursor: CursorPosition, - initialQuery: string, - shouldSuggestColumnAliases?: boolean, -): Pick { - // 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 = {}; - - 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, Parser: ParserConstructor

, - SymbolTableVisitor: SymbolTableVisitorConstructor, tokenDictionary: TokenDictionary, ignoredTokens: Set, - preferredRules: Set, - explicitlyParseJoin: boolean, + rulesToVisit: Set, getParseTree: GetParseTree

, - generateSuggestionsFromRules: GenerateSuggestionsFromRules, + enrichAutocompleteResult: EnrichAutocompleteResult, 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(); @@ -182,12 +64,8 @@ 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( @@ -195,10 +73,8 @@ export function parseQuery< ); } + 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'); @@ -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( @@ -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, ); @@ -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, ); @@ -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, ); diff --git a/src/autocomplete/databases/clickhouse/clickhouse-autocomplete.ts b/src/autocomplete/databases/clickhouse/clickhouse-autocomplete.ts index ae475455..553b0b3c 100644 --- a/src/autocomplete/databases/clickhouse/clickhouse-autocomplete.ts +++ b/src/autocomplete/databases/clickhouse/clickhouse-autocomplete.ts @@ -4,9 +4,11 @@ import * as c3 from 'antlr4-c3'; import {ColumnAliasSymbol, TableSymbol} from '../../shared/symbol-table.js'; import { AutocompleteData, + AutocompleteResultBase, ClickHouseAutocompleteResult, - GenerateSuggestionsFromRulesResult, + CursorPosition, ISymbolTableVisitor, + ProcessVisitedRulesResult, TableOrViewSuggestion, } from '../../autocomplete-types.js'; import {ClickHouseLexer} from './generated/ClickHouseLexer.js'; @@ -17,8 +19,14 @@ import { TableIdentifierContext, } from './generated/ClickHouseParser.js'; import {ClickHouseParserVisitor} from './generated/ClickHouseParserVisitor.js'; -import {TableQueryPosition, TokenDictionary, getPreviousToken} from '../../shared/tables.js'; +import { + TableQueryPosition, + TokenDictionary, + getContextSuggestions, + getPreviousToken, +} from '../../shared/tables.js'; import {isStartingToWriteRule} from '../../shared/cursor'; +import {shouldSuggestTemplates} from '../../shared/query.js'; const engines = ['Null', 'Set', 'Log', 'Memory', 'TinyLog', 'StripeLog']; @@ -83,7 +91,7 @@ function getIgnoredTokens(): number[] { const ignoredTokens = new Set(getIgnoredTokens()); -const preferredRules = new Set([ +const rulesToVisit = new Set([ ClickHouseParser.RULE_databaseIdentifier, ClickHouseParser.RULE_tableIdentifier, ClickHouseParser.RULE_identifier, @@ -151,11 +159,11 @@ class ClickHouseSymbolTableVisitor }; } -function generateSuggestionsFromRules( +function processVisitedRules( rules: c3.CandidatesCollection['rules'], cursorTokenIndex: number, tokenStream: TokenStream, -): GenerateSuggestionsFromRulesResult { +): ProcessVisitedRulesResult { let suggestViewsOrTables: ClickHouseAutocompleteResult['suggestViewsOrTables']; let suggestAggregateFunctions = false; let suggestFunctions = false; @@ -268,19 +276,58 @@ function getParseTree( } } +function enrichAutocompleteResult( + baseResult: AutocompleteResultBase, + rules: c3.CandidatesCollection['rules'], + tokenStream: TokenStream, + cursorTokenIndex: number, + cursor: CursorPosition, + query: string, +): ClickHouseAutocompleteResult { + const {shouldSuggestColumns, shouldSuggestColumnAliases, ...suggestionsFromRules} = + processVisitedRules(rules, cursorTokenIndex, tokenStream); + const suggestTemplates = shouldSuggestTemplates(query, cursor); + const result: ClickHouseAutocompleteResult = { + ...baseResult, + ...suggestionsFromRules, + suggestTemplates, + }; + const contextSuggestionsNeeded = shouldSuggestColumns || shouldSuggestColumnAliases; + + if (contextSuggestionsNeeded) { + const visitor = new ClickHouseSymbolTableVisitor(); + const {tableContextSuggestion, suggestColumnAliases} = getContextSuggestions( + ClickHouseLexer, + ClickHouseParser, + visitor, + tokenDictionary, + getParseTree, + tokenStream, + cursor, + query, + ); + + if (shouldSuggestColumns && tableContextSuggestion) { + result.suggestColumns = tableContextSuggestion; + } + if (shouldSuggestColumnAliases && suggestColumnAliases) { + result.suggestColumnAliases = suggestColumnAliases; + } + } + + return result; +} + export const clickHouseAutocompleteData: AutocompleteData< ClickHouseAutocompleteResult, ClickHouseLexer, - ClickHouseParser, - ClickHouseSymbolTableVisitor + ClickHouseParser > = { Lexer: ClickHouseLexer, Parser: ClickHouseParser, - SymbolTableVisitor: ClickHouseSymbolTableVisitor, tokenDictionary, ignoredTokens, - preferredRules, - explicitlyParseJoin: false, + rulesToVisit, getParseTree, - generateSuggestionsFromRules, + enrichAutocompleteResult, }; diff --git a/src/autocomplete/databases/mysql/grammar/MySqlParser.g4 b/src/autocomplete/databases/mysql/grammar/MySqlParser.g4 index a1c8568d..5d402246 100644 --- a/src/autocomplete/databases/mysql/grammar/MySqlParser.g4 +++ b/src/autocomplete/databases/mysql/grammar/MySqlParser.g4 @@ -539,7 +539,7 @@ partitionOption // Alter statements alterDatabase - : ALTER dbFormat = (DATABASE | SCHEMA) databaseName createDatabaseOption+ # alterSimpleDatabase + : ALTER dbFormat = (DATABASE | SCHEMA) databaseName createDatabaseOption+ # alterSimpleDatabase | ALTER dbFormat = (DATABASE | SCHEMA) databaseName UPGRADE DATA DIRECTORY NAME # alterUpgradeName ; @@ -1592,7 +1592,7 @@ showStatement | SHOW logFormat = (BINLOG | RELAYLOG) EVENTS (IN filename = STRING_LITERAL)? ( FROM fromPosition = decimalLiteral)? (LIMIT (offset = decimalLiteral COMMA)? rowCount = decimalLiteral)? # showLogEvents | SHOW showCommonEntity showFilter? # showObjectFilter | SHOW FULL? columnsFormat = (COLUMNS | FIELDS) tableFormat = (FROM | IN) tableName ( schemaFormat = (FROM | IN) uid)? showFilter? # showColumns - | SHOW CREATE schemaFormat = (DATABASE | SCHEMA) ifNotExists? databaseName # showCreateDb + | SHOW CREATE schemaFormat = (DATABASE | SCHEMA) ifNotExists? databaseName # showCreateDb | SHOW CREATE namedEntity = (EVENT | FUNCTION | PROCEDURE) fullId # showCreateFullIdObject | SHOW CREATE (TABLE | VIEW) tableName # showCreateTableOrView | SHOW CREATE TRIGGER triggerName # showCreateTrigger diff --git a/src/autocomplete/databases/mysql/mysql-autocomplete.ts b/src/autocomplete/databases/mysql/mysql-autocomplete.ts index 5fa0f5df..c369371e 100644 --- a/src/autocomplete/databases/mysql/mysql-autocomplete.ts +++ b/src/autocomplete/databases/mysql/mysql-autocomplete.ts @@ -4,9 +4,11 @@ import * as c3 from 'antlr4-c3'; import {ColumnAliasSymbol, TableSymbol} from '../../shared/symbol-table.js'; import { AutocompleteData, - GenerateSuggestionsFromRulesResult, + AutocompleteResultBase, + CursorPosition, ISymbolTableVisitor, MySqlAutocompleteResult, + ProcessVisitedRulesResult, TableOrViewSuggestion, } from '../../autocomplete-types.js'; import {MySqlLexer} from './generated/MySqlLexer.js'; @@ -17,8 +19,14 @@ import { TableNameContext, } from './generated/MySqlParser.js'; import {MySqlParserVisitor} from './generated/MySqlParserVisitor.js'; -import {TableQueryPosition, TokenDictionary, getPreviousToken} from '../../shared/tables.js'; +import { + TableQueryPosition, + TokenDictionary, + getContextSuggestions, + getPreviousToken, +} from '../../shared/tables.js'; import {isStartingToWriteRule} from '../../shared/cursor'; +import {shouldSuggestTemplates} from '../../shared/query.js'; const tokenDictionary: TokenDictionary = { SPACE: MySqlParser.SPACE, @@ -75,7 +83,7 @@ function getIgnoredTokens(): number[] { const ignoredTokens = new Set(getIgnoredTokens()); -const preferredRules = new Set([ +const rulesToVisit = new Set([ // We don't need to go inside of it, we already know that this is a database name MySqlParser.RULE_databaseName, // We don't need to go inside of it, we already know that this is a constraint name @@ -165,18 +173,18 @@ class MySqlSymbolTableVisitor extends MySqlParserVisitor<{}> implements ISymbolT }; } -function generateSuggestionsFromRules( +function processVisitedRules( rules: c3.CandidatesCollection['rules'], cursorTokenIndex: number, tokenStream: TokenStream, -): GenerateSuggestionsFromRulesResult { +): ProcessVisitedRulesResult { let suggestViewsOrTables: MySqlAutocompleteResult['suggestViewsOrTables']; let suggestAggregateFunctions = false; let suggestFunctions = false; let suggestIndexes = false; let suggestTriggers = false; - let suggestConstraints = false; let suggestDatabases = false; + let shouldSuggestConstraints = false; let shouldSuggestColumns = false; let shouldSuggestColumnAliases = false; @@ -246,7 +254,7 @@ function generateSuggestionsFromRules( break; } case MySqlParser.RULE_constraintName: { - suggestConstraints = true; + shouldSuggestConstraints = true; break; } case MySqlParser.RULE_databaseName: { @@ -289,8 +297,8 @@ function generateSuggestionsFromRules( suggestFunctions, suggestIndexes, suggestTriggers, - suggestConstraints, suggestDatabases, + shouldSuggestConstraints, shouldSuggestColumns, shouldSuggestColumnAliases, }; @@ -318,19 +326,66 @@ function getParseTree( } } +function enrichAutocompleteResult( + baseResult: AutocompleteResultBase, + rules: c3.CandidatesCollection['rules'], + tokenStream: TokenStream, + cursorTokenIndex: number, + cursor: CursorPosition, + query: string, +): MySqlAutocompleteResult { + const { + shouldSuggestColumns, + shouldSuggestColumnAliases, + shouldSuggestConstraints, + ...suggestionsFromRules + } = processVisitedRules(rules, cursorTokenIndex, tokenStream); + const suggestTemplates = shouldSuggestTemplates(query, cursor); + const result: MySqlAutocompleteResult = { + ...baseResult, + ...suggestionsFromRules, + suggestTemplates, + }; + const contextSuggestionsNeeded = + shouldSuggestColumns || shouldSuggestConstraints || shouldSuggestColumnAliases; + + if (contextSuggestionsNeeded) { + const visitor = new MySqlSymbolTableVisitor(); + const {tableContextSuggestion, suggestColumnAliases} = getContextSuggestions( + MySqlLexer, + MySqlParser, + visitor, + tokenDictionary, + getParseTree, + tokenStream, + cursor, + query, + ); + + if (shouldSuggestColumns && tableContextSuggestion) { + result.suggestColumns = tableContextSuggestion; + } + if (shouldSuggestConstraints && tableContextSuggestion) { + result.suggestConstraints = tableContextSuggestion; + } + if (shouldSuggestColumnAliases && suggestColumnAliases) { + result.suggestColumnAliases = suggestColumnAliases; + } + } + + return result; +} + export const mySqlAutocompleteData: AutocompleteData< MySqlAutocompleteResult, MySqlLexer, - MySqlParser, - MySqlSymbolTableVisitor + MySqlParser > = { Lexer: MySqlLexer, Parser: MySqlParser, - SymbolTableVisitor: MySqlSymbolTableVisitor, tokenDictionary, ignoredTokens, - preferredRules, - explicitlyParseJoin: false, + rulesToVisit, getParseTree, - generateSuggestionsFromRules, + enrichAutocompleteResult, }; diff --git a/src/autocomplete/databases/mysql/tests/alter/alter-constraint.test.ts b/src/autocomplete/databases/mysql/tests/alter/alter-constraint.test.ts index c77aeb22..786b226f 100644 --- a/src/autocomplete/databases/mysql/tests/alter/alter-constraint.test.ts +++ b/src/autocomplete/databases/mysql/tests/alter/alter-constraint.test.ts @@ -1,5 +1,5 @@ import {parseMySqlQueryWithCursor} from '../../../../shared/parse-query-with-cursor'; -import {KeywordSuggestion} from '../../../../autocomplete-types'; +import {ConstraintSuggestion, KeywordSuggestion} from '../../../../autocomplete-types'; import {parseMySqlQueryWithoutCursor} from '../../../../autocomplete'; test('should suggest table name after ALTER CONSTRAINT', () => { @@ -10,7 +10,24 @@ test('should suggest table name after ALTER CONSTRAINT', () => { const keywordSuggestion: KeywordSuggestion[] = [{value: 'CHECK'}]; expect(autocompleteResult.suggestKeywords).toEqual(keywordSuggestion); - expect(autocompleteResult.suggestConstraints).toEqual(true); + const constraintSuggestion: ConstraintSuggestion = { + tables: [{name: 'test_table'}], + }; + expect(autocompleteResult.suggestConstraints).toEqual(constraintSuggestion); +}); + +test('should suggest table name after ALTER CONSTRAINT between statements', () => { + const autocompleteResult = parseMySqlQueryWithCursor( + 'ALTER TABLE before_table DROP COLUMN id; ALTER TABLE test_table ALTER CONSTRAINT | ; ALTER TABLE after_table DROP COLUMN id;', + ); + + const keywordSuggestion: KeywordSuggestion[] = [{value: 'CHECK'}]; + expect(autocompleteResult.suggestKeywords).toEqual(keywordSuggestion); + + const constraintSuggestion: ConstraintSuggestion = { + tables: [{name: 'test_table'}], + }; + expect(autocompleteResult.suggestConstraints).toEqual(constraintSuggestion); }); test('should not report errors on a full statement', () => { diff --git a/src/autocomplete/databases/mysql/tests/alter/drop-constraint.test.ts b/src/autocomplete/databases/mysql/tests/alter/drop-constraint.test.ts index cfa13fcb..0bedc712 100644 --- a/src/autocomplete/databases/mysql/tests/alter/drop-constraint.test.ts +++ b/src/autocomplete/databases/mysql/tests/alter/drop-constraint.test.ts @@ -1,13 +1,29 @@ import {parseMySqlQueryWithCursor} from '../../../../shared/parse-query-with-cursor'; import {parseMySqlQueryWithoutCursor} from '../../../../autocomplete'; +import {ConstraintSuggestion} from '../../../../autocomplete-types'; test('should suggest table name after DROP CONSTRAINT', () => { const autocompleteResult = parseMySqlQueryWithCursor( 'ALTER TABLE test_table DROP CONSTRAINT |', ); + const constraintSuggestion: ConstraintSuggestion = { + tables: [{name: 'test_table'}], + }; expect(autocompleteResult.suggestKeywords).toEqual([]); - expect(autocompleteResult.suggestConstraints).toEqual(true); + expect(autocompleteResult.suggestConstraints).toEqual(constraintSuggestion); +}); + +test('should suggest table name after DROP CONSTRAINT between statements', () => { + const autocompleteResult = parseMySqlQueryWithCursor( + 'ALTER TABLE before_table DROP COLUMN id; ALTER TABLE test_table DROP CONSTRAINT | ; ALTER TABLE after_table DROP COLUMN id;', + ); + const constraintSuggestion: ConstraintSuggestion = { + tables: [{name: 'test_table'}], + }; + + expect(autocompleteResult.suggestKeywords).toEqual([]); + expect(autocompleteResult.suggestConstraints).toEqual(constraintSuggestion); }); test('should not report errors on a full statement', () => { diff --git a/src/autocomplete/databases/postgresql/postgresql-autocomplete.ts b/src/autocomplete/databases/postgresql/postgresql-autocomplete.ts index d6b32d9c..2b56aa18 100644 --- a/src/autocomplete/databases/postgresql/postgresql-autocomplete.ts +++ b/src/autocomplete/databases/postgresql/postgresql-autocomplete.ts @@ -4,9 +4,11 @@ import * as c3 from 'antlr4-c3'; import {ColumnAliasSymbol, TableSymbol} from '../../shared/symbol-table.js'; import { AutocompleteData, - GenerateSuggestionsFromRulesResult, + AutocompleteResultBase, + CursorPosition, ISymbolTableVisitor, PostgreSqlAutocompleteResult, + ProcessVisitedRulesResult, TableOrViewSuggestion, } from '../../autocomplete-types.js'; import {PostgreSqlLexer} from './generated/PostgreSqlLexer.js'; @@ -19,8 +21,14 @@ import { ViewNameContext, } from './generated/PostgreSqlParser.js'; import {PostgreSqlParserVisitor} from './generated/PostgreSqlParserVisitor.js'; -import {TableQueryPosition, TokenDictionary, getPreviousToken} from '../../shared/tables.js'; +import { + TableQueryPosition, + TokenDictionary, + getContextSuggestions, + getPreviousToken, +} from '../../shared/tables.js'; import {isStartingToWriteRule} from '../../shared/cursor'; +import {shouldSuggestTemplates} from '../../shared/query.js'; const tokenDictionary: TokenDictionary = { SPACE: PostgreSqlParser.Whitespace, @@ -62,7 +70,7 @@ function getIgnoredTokens(): number[] { const ignoredTokens = new Set(getIgnoredTokens()); -const preferredRules = new Set([ +const rulesToVisit = new Set([ PostgreSqlParser.RULE_columnId, PostgreSqlParser.RULE_functionName, PostgreSqlParser.RULE_functionExpressionCommonSubexpr, @@ -176,20 +184,20 @@ class PostgreSqlSymbolTableVisitor }; } -function generateSuggestionsFromRules( +function processVisitedRules( rules: c3.CandidatesCollection['rules'], cursorTokenIndex: number, tokenStream: TokenStream, -): GenerateSuggestionsFromRulesResult { +): ProcessVisitedRulesResult { let suggestViewsOrTables: PostgreSqlAutocompleteResult['suggestViewsOrTables']; let suggestAggregateFunctions = false; let suggestFunctions = false; let suggestIndexes = false; let suggestTriggers = false; - let suggestConstraints = false; let suggestSequences = false; let suggestSchemas = false; let suggestDatabases = false; + let shouldSuggestConstraints = false; let shouldSuggestColumns = false; let shouldSuggestColumnAliases = false; @@ -277,7 +285,7 @@ function generateSuggestionsFromRules( break; } case PostgreSqlParser.RULE_constraintName: { - suggestConstraints = true; + shouldSuggestConstraints = true; break; } case PostgreSqlParser.RULE_sequenceName: { @@ -301,7 +309,7 @@ function generateSuggestionsFromRules( suggestFunctions, suggestIndexes, suggestTriggers, - suggestConstraints, + shouldSuggestConstraints, suggestSequences, suggestSchemas, suggestDatabases, @@ -332,19 +340,67 @@ function getParseTree( } } +function enrichAutocompleteResult( + baseResult: AutocompleteResultBase, + rules: c3.CandidatesCollection['rules'], + tokenStream: TokenStream, + cursorTokenIndex: number, + cursor: CursorPosition, + query: string, +): PostgreSqlAutocompleteResult { + const { + shouldSuggestColumns, + shouldSuggestColumnAliases, + shouldSuggestConstraints, + ...suggestionsFromRules + } = processVisitedRules(rules, cursorTokenIndex, tokenStream); + const suggestTemplates = shouldSuggestTemplates(query, cursor); + const result: PostgreSqlAutocompleteResult = { + ...baseResult, + ...suggestionsFromRules, + suggestTemplates, + }; + const contextSuggestionsNeeded = + shouldSuggestColumns || shouldSuggestConstraints || shouldSuggestColumnAliases; + + if (contextSuggestionsNeeded) { + const visitor = new PostgreSqlSymbolTableVisitor(); + const {tableContextSuggestion, suggestColumnAliases} = getContextSuggestions( + PostgreSqlLexer, + PostgreSqlParser, + visitor, + tokenDictionary, + getParseTree, + tokenStream, + cursor, + query, + true, + ); + + if (shouldSuggestColumns && tableContextSuggestion) { + result.suggestColumns = tableContextSuggestion; + } + if (shouldSuggestConstraints && tableContextSuggestion) { + result.suggestConstraints = tableContextSuggestion; + } + if (shouldSuggestColumnAliases && suggestColumnAliases) { + result.suggestColumnAliases = suggestColumnAliases; + } + } + + return result; +} + export const postgreSqlAutocompleteData: AutocompleteData< PostgreSqlAutocompleteResult, PostgreSqlLexer, - PostgreSqlParser, - PostgreSqlSymbolTableVisitor + PostgreSqlParser > = { Lexer: PostgreSqlLexer, Parser: PostgreSqlParser, - SymbolTableVisitor: PostgreSqlSymbolTableVisitor, tokenDictionary, ignoredTokens, - preferredRules, - explicitlyParseJoin: true, + rulesToVisit, getParseTree, - generateSuggestionsFromRules, + enrichAutocompleteResult, }; diff --git a/src/autocomplete/databases/postgresql/tests/alter/alter-constraint.test.ts b/src/autocomplete/databases/postgresql/tests/alter/alter-constraint.test.ts index aff7f281..8122fc32 100644 --- a/src/autocomplete/databases/postgresql/tests/alter/alter-constraint.test.ts +++ b/src/autocomplete/databases/postgresql/tests/alter/alter-constraint.test.ts @@ -1,5 +1,5 @@ import {parsePostgreSqlQueryWithCursor} from '../../../../shared/parse-query-with-cursor'; -import {KeywordSuggestion} from '../../../../autocomplete-types'; +import {ConstraintSuggestion, KeywordSuggestion} from '../../../../autocomplete-types'; import {parsePostgreSqlQueryWithoutCursor} from '../../../../autocomplete'; test('should suggest view name after ALTER CONSTRAINT', () => { @@ -7,7 +7,32 @@ test('should suggest view name after ALTER CONSTRAINT', () => { 'ALTER TABLE test_table ALTER CONSTRAINT |', ); - expect(autocompleteResult.suggestConstraints).toEqual(true); + const constraintSuggestion: ConstraintSuggestion = { + tables: [{name: 'test_table'}], + }; + expect(autocompleteResult.suggestConstraints).toEqual(constraintSuggestion); + + const keywordSuggestion: KeywordSuggestion[] = [ + {value: 'OPTIONS'}, + {value: 'SET'}, + {value: 'TYPE'}, + {value: 'DROP'}, + {value: 'RESTART'}, + {value: 'ADD'}, + {value: 'RESET'}, + ]; + expect(autocompleteResult.suggestKeywords).toEqual(keywordSuggestion); +}); + +test('should suggest view name after ALTER CONSTRAINT between statements', () => { + const autocompleteResult = parsePostgreSqlQueryWithCursor( + 'ALTER TABLE before_table DROP COLUMN id; ALTER TABLE test_table ALTER CONSTRAINT | ; ALTER TABLE after_table DROP COLUMN id;', + ); + + const constraintSuggestion: ConstraintSuggestion = { + tables: [{name: 'test_table'}], + }; + expect(autocompleteResult.suggestConstraints).toEqual(constraintSuggestion); const keywordSuggestion: KeywordSuggestion[] = [ {value: 'OPTIONS'}, diff --git a/src/autocomplete/databases/postgresql/tests/alter/drop-constraint.test.ts b/src/autocomplete/databases/postgresql/tests/alter/drop-constraint.test.ts index 9e087747..fa2374ca 100644 --- a/src/autocomplete/databases/postgresql/tests/alter/drop-constraint.test.ts +++ b/src/autocomplete/databases/postgresql/tests/alter/drop-constraint.test.ts @@ -1,5 +1,5 @@ import {parsePostgreSqlQueryWithCursor} from '../../../../shared/parse-query-with-cursor'; -import {KeywordSuggestion} from '../../../../autocomplete-types'; +import {ConstraintSuggestion, KeywordSuggestion} from '../../../../autocomplete-types'; import {parsePostgreSqlQueryWithoutCursor} from '../../../../autocomplete'; test('should suggest view name after DROP CONSTRAINT', () => { @@ -7,7 +7,28 @@ test('should suggest view name after DROP CONSTRAINT', () => { 'ALTER TABLE test_table DROP CONSTRAINT |', ); - expect(autocompleteResult.suggestConstraints).toEqual(true); + const constraintSuggestion: ConstraintSuggestion = { + tables: [{name: 'test_table'}], + }; + expect(autocompleteResult.suggestConstraints).toEqual(constraintSuggestion); + + const keywordSuggestion: KeywordSuggestion[] = [ + {value: 'IF'}, + {value: 'CASCADE'}, + {value: 'RESTRICT'}, + ]; + expect(autocompleteResult.suggestKeywords).toEqual(keywordSuggestion); +}); + +test('should suggest view name after DROP CONSTRAINT between statements', () => { + const autocompleteResult = parsePostgreSqlQueryWithCursor( + 'ALTER TABLE before_table DROP COLUMN id; ALTER TABLE test_table DROP CONSTRAINT | ; ALTER TABLE after_table DROP COLUMN id;', + ); + + const constraintSuggestion: ConstraintSuggestion = { + tables: [{name: 'test_table'}], + }; + expect(autocompleteResult.suggestConstraints).toEqual(constraintSuggestion); const keywordSuggestion: KeywordSuggestion[] = [ {value: 'IF'}, diff --git a/src/autocomplete/databases/postgresql/tests/alter/rename-constraint.test.ts b/src/autocomplete/databases/postgresql/tests/alter/rename-constraint.test.ts index a0972f94..405030f9 100644 --- a/src/autocomplete/databases/postgresql/tests/alter/rename-constraint.test.ts +++ b/src/autocomplete/databases/postgresql/tests/alter/rename-constraint.test.ts @@ -1,4 +1,4 @@ -import {KeywordSuggestion} from '../../../../autocomplete-types'; +import {ConstraintSuggestion, KeywordSuggestion} from '../../../../autocomplete-types'; import {parsePostgreSqlQueryWithCursor} from '../../../../shared/parse-query-with-cursor'; import {parsePostgreSqlQueryWithoutCursor} from '../../../../autocomplete'; @@ -10,7 +10,24 @@ test('should suggest view name after RENAME CONSTRAINT', () => { const keywordSuggestion: KeywordSuggestion[] = [{value: 'TO'}]; expect(autocompleteResult.suggestKeywords).toEqual(keywordSuggestion); - expect(autocompleteResult.suggestConstraints).toEqual(true); + const constraintSuggestion: ConstraintSuggestion = { + tables: [{name: 'test_table'}], + }; + expect(autocompleteResult.suggestConstraints).toEqual(constraintSuggestion); +}); + +test('should suggest view name after RENAME CONSTRAINT between statements', () => { + const autocompleteResult = parsePostgreSqlQueryWithCursor( + 'ALTER TABLE before_table DROP COLUMN id; ALTER TABLE test_table RENAME CONSTRAINT | ; ALTER TABLE after_table DROP COLUMN id;', + ); + + const keywordSuggestion: KeywordSuggestion[] = [{value: 'TO'}]; + expect(autocompleteResult.suggestKeywords).toEqual(keywordSuggestion); + + const constraintSuggestion: ConstraintSuggestion = { + tables: [{name: 'test_table'}], + }; + expect(autocompleteResult.suggestConstraints).toEqual(constraintSuggestion); }); test('should not report errors', () => { diff --git a/src/autocomplete/databases/postgresql/tests/alter/validate-constraint.test.ts b/src/autocomplete/databases/postgresql/tests/alter/validate-constraint.test.ts index 3b5eeabc..f81576f6 100644 --- a/src/autocomplete/databases/postgresql/tests/alter/validate-constraint.test.ts +++ b/src/autocomplete/databases/postgresql/tests/alter/validate-constraint.test.ts @@ -1,12 +1,28 @@ import {parsePostgreSqlQueryWithCursor} from '../../../../shared/parse-query-with-cursor'; import {parsePostgreSqlQueryWithoutCursor} from '../../../../autocomplete'; +import {ConstraintSuggestion} from '../../../../autocomplete-types'; test('should suggest view name after VALIDATE CONSTRAINT', () => { const autocompleteResult = parsePostgreSqlQueryWithCursor( 'ALTER TABLE test_table VALIDATE CONSTRAINT |', ); - expect(autocompleteResult.suggestConstraints).toEqual(true); + const constraintSuggestion: ConstraintSuggestion = { + tables: [{name: 'test_table'}], + }; + expect(autocompleteResult.suggestConstraints).toEqual(constraintSuggestion); + expect(autocompleteResult.suggestKeywords).toEqual([]); +}); + +test('should suggest view name after VALIDATE CONSTRAINT between statements', () => { + const autocompleteResult = parsePostgreSqlQueryWithCursor( + 'ALTER TABLE before_table DROP COLUMN id; ALTER TABLE test_table VALIDATE CONSTRAINT | ; ALTER TABLE after_table DROP COLUMN id;', + ); + + const constraintSuggestion: ConstraintSuggestion = { + tables: [{name: 'test_table'}], + }; + expect(autocompleteResult.suggestConstraints).toEqual(constraintSuggestion); expect(autocompleteResult.suggestKeywords).toEqual([]); }); diff --git a/src/autocomplete/databases/postgresql/tests/comment/comment-constraint.test.ts b/src/autocomplete/databases/postgresql/tests/comment/comment-constraint.test.ts index fc3e3e2e..88d574fa 100644 --- a/src/autocomplete/databases/postgresql/tests/comment/comment-constraint.test.ts +++ b/src/autocomplete/databases/postgresql/tests/comment/comment-constraint.test.ts @@ -1,7 +1,8 @@ import {parsePostgreSqlQueryWithCursor} from '../../../../shared/parse-query-with-cursor'; import {parsePostgreSqlQueryWithoutCursor} from '../../../../autocomplete'; -test('should suggest properly after COMMENT ON CONSTRAINT', () => { +// TODO Get context of table in COMMENT statement +test.skip('should suggest properly after COMMENT ON CONSTRAINT', () => { const autocompleteResult = parsePostgreSqlQueryWithCursor('COMMENT ON CONSTRAINT |'); expect(autocompleteResult.suggestKeywords).toEqual([]); diff --git a/src/autocomplete/shared/cursor.ts b/src/autocomplete/shared/cursor.ts index a22c6329..42072dfc 100644 --- a/src/autocomplete/shared/cursor.ts +++ b/src/autocomplete/shared/cursor.ts @@ -28,7 +28,7 @@ export function findCursorTokenIndex( tokenStream: TokenStream, cursor: CursorPosition, whitespaceToken: number, - realIndex?: boolean, + actualIndex?: boolean, ): number | undefined { // Cursor position is 1-based, while token's charPositionInLine is 0-based const cursorCol = cursor.column - 1; @@ -42,7 +42,7 @@ export function findCursorTokenIndex( // endColumn makes sense only if startLine === endLine if (endLine > cursor.line || (startLine === cursor.line && endColumn > cursorCol)) { - if (realIndex) { + if (actualIndex) { return i; } diff --git a/src/autocomplete/shared/query.ts b/src/autocomplete/shared/query.ts index 9ae757b9..a9cd13b1 100644 --- a/src/autocomplete/shared/query.ts +++ b/src/autocomplete/shared/query.ts @@ -1,3 +1,8 @@ +import {CharStreams, CommonTokenStream, Lexer as LexerType, Parser as ParserType} from 'antlr4ng'; + +import {CursorPosition, LexerConstructor, ParserConstructor} from '../autocomplete-types'; +import {getCursorIndex} from './cursor'; + export function getCurrentStatement( query: string, cursorIndex: number, @@ -22,8 +27,13 @@ const spaceSymbols = '(\\s|\r\n|\n|\r)+'; const explainRegex = new RegExp(`^(${spaceSymbols})?explain${spaceSymbols}$`); const multipleKeywordsRegex = new RegExp(`^(${spaceSymbols})?\\S+${spaceSymbols}`); -export function shouldSuggestTemplates(statement: string, cursorIndex: number): boolean { - const currentStatementBeforeCursor = statement.slice(0, cursorIndex).toLowerCase(); +// TODO Find a better way to suggestTemplates +export function shouldSuggestTemplates(query: string, cursor: CursorPosition): boolean { + const cursorIndex = getCursorIndex(query, cursor); + const currentStatement = getCurrentStatement(query, cursorIndex); + const currentStatementBeforeCursor = currentStatement.statement + .slice(0, currentStatement.cursorIndex) + .toLowerCase(); return Boolean( cursorIndex === 0 || @@ -33,3 +43,18 @@ export function shouldSuggestTemplates(statement: string, cursorIndex: number): currentStatementBeforeCursor.match(explainRegex), ); } + +export function createParser( + Lexer: LexerConstructor, + Parser: ParserConstructor

, + query: string, +): P { + const inputStream = CharStreams.fromString(query); + const lexer = new Lexer(inputStream); + const tokenStream = new CommonTokenStream(lexer); + const parser = new Parser(tokenStream); + + parser.removeErrorListeners(); + + return parser; +} diff --git a/src/autocomplete/shared/symbol-table.ts b/src/autocomplete/shared/symbol-table.ts index 677c288a..fd113802 100644 --- a/src/autocomplete/shared/symbol-table.ts +++ b/src/autocomplete/shared/symbol-table.ts @@ -1,4 +1,5 @@ import * as c3 from 'antlr4-c3'; +import {ColumnAliasSuggestion, SymbolTableVisitor, Table} from '../autocomplete-types'; export class TableSymbol extends c3.TypedSymbol { name: string; @@ -12,6 +13,13 @@ export class TableSymbol extends c3.TypedSymbol { } } +export function getTablesFromSymbolTable(visitor: SymbolTableVisitor): Table[] { + return visitor.symbolTable.getNestedSymbolsOfTypeSync(TableSymbol).map((tableSymbol) => ({ + name: tableSymbol.name, + alias: tableSymbol.alias, + })); +} + export class ColumnAliasSymbol extends c3.TypedSymbol { name: string; @@ -21,3 +29,11 @@ export class ColumnAliasSymbol extends c3.TypedSymbol { this.name = name; } } + +export function getColumnAliasesFromSymbolTable( + visitor: SymbolTableVisitor, +): ColumnAliasSuggestion[] { + return visitor.symbolTable + .getNestedSymbolsOfTypeSync(ColumnAliasSymbol) + .map(({name}) => ({name})); +} diff --git a/src/autocomplete/shared/tables.ts b/src/autocomplete/shared/tables.ts index c7484ed3..f867f14d 100644 --- a/src/autocomplete/shared/tables.ts +++ b/src/autocomplete/shared/tables.ts @@ -1,4 +1,16 @@ -import {Token, TokenStream} from 'antlr4ng'; +import {Lexer as LexerType, Parser as ParserType, Token, TokenStream} from 'antlr4ng'; +import {findCursorTokenIndex} from './cursor'; +import {createParser} from './query'; +import {getColumnAliasesFromSymbolTable, getTablesFromSymbolTable} from './symbol-table'; +import { + AutocompleteResultBase, + CursorPosition, + GetParseTree, + LexerConstructor, + ParserConstructor, + SymbolTableVisitor, + TableContextSuggestion, +} from '../autocomplete-types'; interface TableQueryPositionBase { start: number; @@ -224,3 +236,82 @@ export function getPreviousToken( return undefined; } + +interface ContextSuggestions { + tableContextSuggestion?: TableContextSuggestion; + suggestColumnAliases?: AutocompleteResultBase['suggestColumnAliases']; +} + +export function getContextSuggestions( + Lexer: LexerConstructor, + Parser: ParserConstructor

, + symbolTableVisitor: SymbolTableVisitor, + tokenDictionary: TokenDictionary, + getParseTree: GetParseTree

, + tokenStream: TokenStream, + cursor: CursorPosition, + query: string, + explicitlyParseJoin?: boolean, +): ContextSuggestions { + // The actual token index, without special logic for spaces + const actualCursorTokenIndex = findCursorTokenIndex( + tokenStream, + cursor, + tokenDictionary.SPACE, + true, + ); + if (!actualCursorTokenIndex) { + throw new Error( + `Could not find actualCursorTokenIndex at Ln ${cursor.line}, Col ${cursor.column}`, + ); + } + + const contextSuggestions: ContextSuggestions = {}; + const tableQueryPosition = getTableQueryPosition( + tokenStream, + actualCursorTokenIndex, + tokenDictionary, + ); + + if (tableQueryPosition) { + const tableQuery = query.slice(tableQueryPosition.start, tableQueryPosition.end); + const parser = createParser(Lexer, Parser, tableQuery); + const parseTree = getParseTree(parser, tableQueryPosition.type); + + symbolTableVisitor.visit(parseTree); + + if (explicitlyParseJoin && tableQueryPosition.joinTableQueryPosition) { + const joinTableQuery = query.slice( + tableQueryPosition.joinTableQueryPosition.start, + tableQueryPosition.joinTableQueryPosition.end, + ); + const joinParser = createParser(Lexer, Parser, joinTableQuery); + const joinParseTree = getParseTree(joinParser, 'from'); + symbolTableVisitor.visit(joinParseTree); + } + + if (tableQueryPosition.selectTableQueryPosition) { + const selectTableQuery = query.slice( + tableQueryPosition.selectTableQueryPosition.start, + tableQueryPosition.selectTableQueryPosition.end, + ); + const selectParser = createParser(Lexer, Parser, selectTableQuery); + const selectParseTree = getParseTree(selectParser, 'select'); + symbolTableVisitor.visit(selectParseTree); + } + + const tables = getTablesFromSymbolTable(symbolTableVisitor); + if (tables.length) { + contextSuggestions.tableContextSuggestion = { + tables, + }; + } + + const columnAliases = getColumnAliasesFromSymbolTable(symbolTableVisitor); + if (columnAliases.length) { + contextSuggestions.suggestColumnAliases = columnAliases.map(({name}) => ({name})); + } + } + + return contextSuggestions; +} diff --git a/src/index.ts b/src/index.ts index 76750462..21655a18 100644 --- a/src/index.ts +++ b/src/index.ts @@ -19,4 +19,6 @@ export { ColumnAliasSuggestion, EngineSuggestion, CursorPosition, + ConstraintSuggestion, + Table, } from './autocomplete/autocomplete-types';