feat(YQL): handle multi spaces before cursor (#274)

* feat(YQL): handle multi spaces before cursor * fix: function name * fix: remove identifierTokenTypes from computeTokenContext
gravity-ui · Jan 9, 2025 · 223111b · 223111b
1 parent 92abdea
commit 223111b
Show file tree

Hide file tree

Showing 5 changed files with 95 additions and 20 deletions.
diff --git a/src/autocomplete/databases/yql/tests/yql/select/select-scopes.test.ts b/src/autocomplete/databases/yql/tests/yql/select/select-scopes.test.ts
@@ -68,3 +68,25 @@ test('should suggest columns from first subquery', () => {
     };
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);
 });
+
+test('should suggest columns from subquery with multi spaces before', () => {
+    const autocompleteResult = parseYqlQueryWithCursor(
+        '$a = select foo from bar; select   | FROM $a ',
+    );
+    const columnSuggestions: YQLColumnsSuggestion = {
+        tables: [{name: '$a'}],
+        all: true,
+    };
+    expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);
+});
+
+test('should suggest columns from subquery with multi spaces before if curson is on EOF', () => {
+    const autocompleteResult = parseYqlQueryWithCursor(
+        '$a = select foo from bar; FROM $a select   | ',
+    );
+    const columnSuggestions: YQLColumnsSuggestion = {
+        tables: [{name: '$a'}],
+        all: true,
+    };
+    expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);
+});
diff --git a/src/autocomplete/databases/yql/yql-autocomplete.ts b/src/autocomplete/databases/yql/yql-autocomplete.ts
@@ -535,6 +535,7 @@ function getEnrichAutocompleteResult(parseTreeGetter: GetParseTree<YQLParser>) {
                 tokenStream,
                 cursor,
                 query,
+                [YQLParser.WS],
             );
             if (data.length) {
                 result.suggestVariables = data;
@@ -551,6 +552,7 @@ function getEnrichAutocompleteResult(parseTreeGetter: GetParseTree<YQLParser>) {
                 tokenStream,
                 cursor,
                 query,
+                [YQLParser.WS],
             );
 
             if (shouldSuggestColumns && tableContextSuggestion) {

diff --git a/src/autocomplete/shared/compute-token-position.ts b/src/autocomplete/shared/compute-token-position.ts
@@ -4,33 +4,86 @@ import {CursorPosition} from './autocomplete-types';
 
 type TokenContext = {index: number; context: ParseTree; text: string};
 
+function getWhitespaceBeforeCursor(
+    tokenStream: TokenStream,
+    cursorPosition: CursorPosition,
+    whitespaceTokenTypes: number[],
+): number {
+    for (let i = 0; i < tokenStream.size; i++) {
+        const token = tokenStream.get(i);
+        const tokenStart = token.column;
+        const tokenLength = token.text?.length ?? 0;
+        const tokenStop = token.column + tokenLength;
+        // If it is token before cursor
+        if (
+            token.line === cursorPosition.line &&
+            tokenStart <= cursorPosition.column - 1 &&
+            tokenStop >= cursorPosition.column - 1 &&
+            whitespaceTokenTypes.includes(token.type)
+        ) {
+            let j = i - 1;
+            while (j >= 0) {
+                const prevToken = tokenStream.get(j);
+                if (whitespaceTokenTypes.includes(prevToken.type)) {
+                    j -= 1;
+                } else {
+                    return i + 1 - j;
+                }
+            }
+        }
+    }
+    return 0;
+}
+
 export function computeTokenContext(
     parseTree: ParseTree,
     tokenStream: TokenStream,
     cursorPosition: CursorPosition,
-    identifierTokenTypes: number[] = [],
+    whitespaceTokenTypes: number[] = [],
+): TokenContext | undefined {
+    const whitespacesBeforeCursor = getWhitespaceBeforeCursor(
+        tokenStream,
+        cursorPosition,
+        whitespaceTokenTypes,
+    );
+    return computeTokenContextInternal(
+        parseTree,
+        tokenStream,
+        cursorPosition,
+        whitespacesBeforeCursor,
+    );
+}
+
+function computeTokenContextInternal(
+    parseTree: ParseTree,
+    tokenStream: TokenStream,
+    cursorPosition: CursorPosition,
+    whitespacesBeforeCursor: number,
 ): TokenContext | undefined {
     if (parseTree instanceof TerminalNode) {
-        return computeTokenContextOfTerminalNode(parseTree, cursorPosition, identifierTokenTypes);
+        return computeTokenContextOfTerminalNode(
+            parseTree,
+            cursorPosition,
+            whitespacesBeforeCursor,
+        );
     }
     return computeTokenContextOfChildNode(
         parseTree as ParserRuleContext,
         tokenStream,
         cursorPosition,
-        identifierTokenTypes,
+        whitespacesBeforeCursor,
     );
 }
 
 function getTokenContext(
     token: Token,
     text: string,
     cursorPosition: CursorPosition,
-    identifierTokenTypes: number[],
     parseTree: ParseTree,
+    whitespacesBeforeCursor: number,
 ): TokenContext | undefined {
     const start = token.column;
     const stop = token.column + text.length;
-
     // It means token is eof
     if (token.start > token.stop) {
         return {
@@ -44,15 +97,11 @@ function getTokenContext(
         token.line === cursorPosition.line &&
         // To get context correctly we need to take previous token i.e.
         // "select | from" - need to get match with `select` token
-        start <= cursorPosition.column - 2 &&
-        stop >= cursorPosition.column - 2
+        start <= cursorPosition.column - 1 - whitespacesBeforeCursor &&
+        stop >= cursorPosition.column - 1 - whitespacesBeforeCursor
     ) {
-        let index = token.tokenIndex;
-        if (identifierTokenTypes.includes(token.type)) {
-            index--;
-        }
         return {
-            index,
+            index: token.tokenIndex,
             context: parseTree,
             text: text.substring(0, cursorPosition.column - start),
         };
@@ -64,18 +113,18 @@ function getTokenContext(
 function computeTokenContextOfTerminalNode(
     parseTree: TerminalNode,
     cursorPosition: CursorPosition,
-    identifierTokenTypes: number[],
+    whitespacesBeforeCursor: number,
 ): TokenContext | undefined {
     const token = parseTree.symbol;
     const text = parseTree.getText();
-    return getTokenContext(token, text, cursorPosition, identifierTokenTypes, parseTree);
+    return getTokenContext(token, text, cursorPosition, parseTree, whitespacesBeforeCursor);
 }
 
 function computeTokenContextOfChildNode(
     parseTree: ParserRuleContext,
     tokenStream: TokenStream,
     cursorPosition: CursorPosition,
-    identifierTokenTypes: number[],
+    whitespacesBeforeCursor: number,
 ): TokenContext | undefined {
     if (
         !parseTree.start ||
@@ -91,11 +140,11 @@ function computeTokenContextOfChildNode(
         if (!child) {
             continue;
         }
-        const tokenContext = computeTokenContext(
+        const tokenContext = computeTokenContextInternal(
             child,
             tokenStream,
             cursorPosition,
-            identifierTokenTypes,
+            whitespacesBeforeCursor,
         );
         if (tokenContext) {
             return tokenContext;
@@ -107,8 +156,8 @@ function computeTokenContextOfChildNode(
             tokenStream.get(i),
             tokenStream.get(i).text ?? '',
             cursorPosition,
-            identifierTokenTypes,
             parseTree,
+            whitespacesBeforeCursor,
         );
         if (tokenContext) {
             return tokenContext;

diff --git a/src/autocomplete/shared/extended-tables.ts b/src/autocomplete/shared/extended-tables.ts
@@ -28,11 +28,12 @@ export function getExtendedTableSuggestions<L extends LexerType, P extends Parse
     tokenStream: TokenStream,
     cursor: CursorPosition,
     query: string,
+    whitespaceTokenTypes: number[] = [],
 ): ContextSuggestions {
     const result: ContextSuggestions = {};
     const parser = createParser(Lexer, Parser, query);
     const parseTree = getParseTree(parser);
-    const tokenContext = computeTokenContext(parseTree, tokenStream, cursor);
+    const tokenContext = computeTokenContext(parseTree, tokenStream, cursor, whitespaceTokenTypes);
 
     if (!tokenContext) {
         throw new Error(`Could not find tokenContext at Ln ${cursor.line}, Col ${cursor.column}`);

diff --git a/src/autocomplete/shared/variables.ts b/src/autocomplete/shared/variables.ts
@@ -20,11 +20,12 @@ export function getVariableSuggestions<L extends LexerType, P extends ParserType
     tokenStream: TokenStream,
     cursor: CursorPosition,
     query: string,
+    whitespaceTokenTypes: number[] = [],
 ): VariableSuggestion[] {
     const parser = createParser(Lexer, Parser, query);
     const parseTree = getParseTree(parser);
 
-    const tokenContext = computeTokenContext(parseTree, tokenStream, cursor);
+    const tokenContext = computeTokenContext(parseTree, tokenStream, cursor, whitespaceTokenTypes);
 
     if (!tokenContext) {
         throw new Error(`Could not find tokenContext at Ln ${cursor.line}, Col ${cursor.column}`);