diff --git a/src/parser/common/basicSQL.ts b/src/parser/common/basicSQL.ts index 8d863890..bbd8bb8b 100644 --- a/src/parser/common/basicSQL.ts +++ b/src/parser/common/basicSQL.ts @@ -83,6 +83,60 @@ export abstract class BasicSQL< caretTokenIndex: number ): Suggestions; + protected getCandidateTokenRanges( + candidates: CandidatesCollection, + candidateStartTokenIndex: number, + allTokens: Token[], + caretTokenIndex: number + ): Token[] { + // antlr4-c3 可能同时返回实体和别名候选,使用后续最近候选的起始位置作为边界 + const endTokenIndex = Array.from(candidates.rules.values()).reduce( + (nearestStartTokenIndex, candidateRule) => { + if (candidateRule.startTokenIndex <= candidateStartTokenIndex) { + return nearestStartTokenIndex; + } + return Math.min(nearestStartTokenIndex, candidateRule.startTokenIndex); + }, + caretTokenIndex + 1 + ); + const previousVisibleToken = allTokens + .slice(candidateStartTokenIndex, endTokenIndex) + .reverse() + .find((token) => token.channel === Token.DEFAULT_CHANNEL); + const visibleTokenIndexes = allTokens + .slice(endTokenIndex, caretTokenIndex + 1) + .reduce((indexes, token, offset) => { + if (token.channel === Token.DEFAULT_CHANNEL) { + indexes.push(endTokenIndex + offset); + } + return indexes; + }, []); + const firstVisibleToken = allTokens[visibleTokenIndexes[0]]; + let rangeEndTokenIndex = endTokenIndex; + + // 候选边界可能落在多级限定名中间,需要沿标识符与点号链继续扩展 + if (previousVisibleToken?.text === '.' || firstVisibleToken?.text === '.') { + let visibleTokenOffset = 0; + if (previousVisibleToken?.text === '.' && firstVisibleToken) { + rangeEndTokenIndex = visibleTokenIndexes[visibleTokenOffset] + 1; + visibleTokenOffset += 1; + } + + while (allTokens[visibleTokenIndexes[visibleTokenOffset]]?.text === '.') { + rangeEndTokenIndex = visibleTokenIndexes[visibleTokenOffset] + 1; + visibleTokenOffset += 1; + if (visibleTokenOffset < visibleTokenIndexes.length) { + rangeEndTokenIndex = visibleTokenIndexes[visibleTokenOffset] + 1; + visibleTokenOffset += 1; + } + } + } + + return allTokens + .slice(candidateStartTokenIndex, rangeEndTokenIndex) + .filter((token) => token.channel === Token.DEFAULT_CHANNEL); + } + /** * Get a new splitListener instance. */ diff --git a/src/parser/flink/index.ts b/src/parser/flink/index.ts index f6247420..a5fc7771 100644 --- a/src/parser/flink/index.ts +++ b/src/parser/flink/index.ts @@ -93,7 +93,12 @@ export class FlinkSQL extends BasicSQL { for (const candidate of candidates.rules) { const [ruleType, candidateRule] = candidate; - const tokenRanges = allTokens.slice(candidateRule.startTokenIndex, caretTokenIndex + 1); + const tokenRanges = this.getCandidateTokenRanges( + candidates, + candidateRule.startTokenIndex, + allTokens, + caretTokenIndex + ); let syntaxContextType: EntityContextType | StmtContextType | undefined = void 0; switch (ruleType) { diff --git a/src/parser/postgresql/index.ts b/src/parser/postgresql/index.ts index 02944e39..4d5410e8 100644 --- a/src/parser/postgresql/index.ts +++ b/src/parser/postgresql/index.ts @@ -90,7 +90,12 @@ export class PostgreSQL extends BasicSQL; + +const parserFactories: Array<[string, () => SuggestionParser]> = [ + ['MySQL', () => new MySQL()], + ['FlinkSQL', () => new FlinkSQL()], + ['SparkSQL', () => new SparkSQL()], + ['HiveSQL', () => new HiveSQL()], + ['PostgreSQL', () => new PostgreSQL()], + ['TrinoSQL', () => new TrinoSQL()], + ['ImpalaSQL', () => new ImpalaSQL()], + ['GenericSQL', () => new GenericSQL()], +]; + +class TestableSparkSQL extends SparkSQL { + public getCandidateTokenRangesForTest( + candidates: CandidatesCollection, + candidateStartTokenIndex: number, + allTokens: Token[], + caretTokenIndex: number + ): Token[] { + return this.getCandidateTokenRanges( + candidates, + candidateStartTokenIndex, + allTokens, + caretTokenIndex + ); + } +} + +const scenarios = [ + { + name: 'exclude trailing whitespace from table word ranges', + sql: 'SELECT * FROM current_catalog_schema1 ', + expected: ['current_catalog_schema1'], + }, + { + name: 'exclude AS from table word ranges', + sql: 'SELECT * FROM current_catalog_schema1 as', + expected: ['current_catalog_schema1'], + }, + { + name: 'exclude alias from table word ranges', + sql: 'SELECT * FROM current_catalog_schema1 alias', + expected: ['current_catalog_schema1'], + }, + { + name: 'preserve an incomplete qualified table name', + sql: 'SELECT * FROM db.', + expected: ['db', '.'], + }, +]; + +describe.each(parserFactories)('%s syntax suggestion word ranges', (_name, createParser) => { + test.each(scenarios)('$name', ({ sql, expected }) => { + const tableSuggestion = createParser() + .getSuggestionAtCaretPosition(sql, { + lineNumber: 1, + column: sql.length + 1, + }) + ?.syntax.find((suggestion) => suggestion.syntaxContextType === EntityContextType.TABLE); + + expect(tableSuggestion).toBeDefined(); + expect(tableSuggestion?.wordRanges.map((wordRange) => wordRange.text)).toEqual(expected); + }); +}); + +test('SparkSQL preserves a qualified table name separated by hidden tokens', () => { + const sql = 'SELECT * FROM db. table'; + const tableSuggestion = new SparkSQL() + .getSuggestionAtCaretPosition(sql, { + lineNumber: 1, + column: sql.length + 1, + }) + ?.syntax.find((suggestion) => suggestion.syntaxContextType === EntityContextType.TABLE); + + expect(tableSuggestion).toBeDefined(); + expect(tableSuggestion?.wordRanges.map((wordRange) => wordRange.text)).toEqual([ + 'db', + '.', + 'table', + ]); +}); + +test('SparkSQL preserves a qualified table name separated by a comment', () => { + const sql = 'SELECT * FROM db./* comment */table'; + const tableSuggestion = new SparkSQL() + .getSuggestionAtCaretPosition(sql, { + lineNumber: 1, + column: sql.length + 1, + }) + ?.syntax.find((suggestion) => suggestion.syntaxContextType === EntityContextType.TABLE); + + expect(tableSuggestion).toBeDefined(); + expect(tableSuggestion?.wordRanges.map((wordRange) => wordRange.text)).toEqual([ + 'db', + '.', + 'table', + ]); +}); + +test('preserves a multi-level qualified table name when another candidate starts in the middle', () => { + const parser = new TestableSparkSQL(); + const allTokens = parser.getAllTokens('catalog.schema.table'); + const candidates = new CandidatesCollection(); + candidates.rules.set(0, { startTokenIndex: 0, ruleList: [] }); + candidates.rules.set(1, { startTokenIndex: 2, ruleList: [] }); + + const wordRanges = parser.getCandidateTokenRangesForTest(candidates, 0, allTokens, 4); + + expect(wordRanges.map((wordRange) => wordRange.text)).toEqual([ + 'catalog', + '.', + 'schema', + '.', + 'table', + ]); +});