fix: #239 correct suggestion word ranges

src/parser/common/basicSQL.ts

@@ -83,6 +83,38 @@ 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);
+        // look past hidden tokens to detect dot, preserving dot and following identifier in incomplete qualified names
+        const rangeEndTokenIndex =
+            endTokenIndex <= caretTokenIndex &&
+            (allTokens[endTokenIndex]?.text === '.' || previousVisibleToken?.text === '.')
+                ? endTokenIndex + 1
+                : endTokenIndex;
+
+        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,95 @@
+import {
+    FlinkSQL,
+    GenericSQL,
+    HiveSQL,
+    ImpalaSQL,
+    MySQL,
+    PostgreSQL,
+    SparkSQL,
+    TrinoSQL,
+} from 'src/index';
+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()],
+];
+
+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',
+    ]);
+});