fix: #239 correct suggestion word ranges

src/parser/common/basicSQL.ts

@@ -83,6 +83,60 @@ export abstract class BasicSQL<
         caretTokenIndex: number
     ): Suggestions<Token>;
 
+    protected getCandidateTokenRanges(
+        candidates: CandidatesCollection,
+        candidateStartTokenIndex: number,
+        allTokens: Token[],
+        caretTokenIndex: number
+    ): Token[] {
+        // antlr4-c3 may return both entity and alias candidates; use the nearest candidate's start index as boundary
+        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<number[]>((indexes, token, offset) => {
+                if (token.channel === Token.DEFAULT_CHANNEL) {
+                    indexes.push(endTokenIndex + offset);
+                }
+                return indexes;
+            }, []);
+        const firstVisibleToken = allTokens[visibleTokenIndexes[0]];
+        let rangeEndTokenIndex = endTokenIndex;
+
+        // candidate boundary may fall inside a multi-level qualified name; extend along the identifier and dot chain
+        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.
      */

src/parser/flink/index.ts

@@ -93,7 +93,12 @@ export class FlinkSQL extends BasicSQL<FlinkSqlLexer, ProgramContext, FlinkSqlPa
 
         for (let 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) {

src/parser/generic/index.ts

@@ -91,7 +91,12 @@ export class GenericSQL extends BasicSQL<GenericSqlLexer, ProgramContext, Generi
 
         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) {

src/parser/hive/index.ts

@@ -81,7 +81,12 @@ export class HiveSQL extends BasicSQL<HiveSqlLexer, ProgramContext, HiveSqlParse
         const keywords: string[] = [];
         for (let 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) {

src/parser/impala/index.ts

@@ -79,7 +79,12 @@ export class ImpalaSQL extends BasicSQL<ImpalaSqlLexer, ProgramContext, ImpalaSq
         const keywords: string[] = [];
         for (let 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) {

src/parser/mysql/index.ts

@@ -80,7 +80,12 @@ export class MySQL extends BasicSQL<MySqlLexer, ProgramContext, MySqlParser> {
 
         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) {

src/parser/postgresql/index.ts

@@ -90,7 +90,12 @@ export class PostgreSQL extends BasicSQL<PostgreSqlLexer, ProgramContext, Postgr
         const keywords: string[] = [];
         for (let 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) {

src/parser/spark/index.ts

@@ -85,7 +85,12 @@ export class SparkSQL extends BasicSQL<SparkSqlLexer, ProgramContext, SparkSqlPa
 
         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) {

src/parser/trino/index.ts

@@ -82,7 +82,12 @@ export class TrinoSQL extends BasicSQL<TrinoSqlLexer, ProgramContext, TrinoSqlPa
 
         for (let 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) {

test/parser/syntaxSuggestionWordRange.test.ts

@@ -0,0 +1,131 @@
+import {
+    FlinkSQL,
+    GenericSQL,
+    HiveSQL,
+    ImpalaSQL,
+    MySQL,
+    PostgreSQL,
+    SparkSQL,
+    TrinoSQL,
+} from 'src/index';
+import { CandidatesCollection } from 'antlr4-c3';
+import { Token } from 'antlr4ng';
+import { EntityContextType } from 'src/parser/common/types';
+
+type SuggestionParser = Pick<MySQL, 'getSuggestionAtCaretPosition'>;
+
+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',
+    ]);
+});