feat(YQL): support UPSERT and all columns suggestions (#234)

* feat(YQL): support UPSERT statement * feat(YQL): suggest all columns in some cases * fix: review
gravity-ui · Nov 7, 2024 · 037ef37 · 037ef37
1 parent 3b8e2e1
commit 037ef37
Show file tree

Hide file tree

Showing 12 changed files with 245 additions and 61 deletions.
diff --git a/src/autocomplete/databases/yql/helpers.ts b/src/autocomplete/databases/yql/helpers.ts
@@ -11,6 +11,7 @@ export const tokenDictionary: TokenDictionary = {
     CLOSING_BRACKET: YQLParser.RPAREN,
     ALTER: YQLParser.ALTER,
     INSERT: YQLParser.INSERT,
+    UPSERT: YQLParser.UPSERT,
     UPDATE: YQLParser.UPDATE,
     JOIN: YQLParser.JOIN,
     SEMICOLON: YQLParser.SEMICOLON,
@@ -20,6 +21,26 @@ export const tokenDictionary: TokenDictionary = {
 type AnyRuleInList = (rules: number | number[]) => boolean;
 type AllRulesInList = (rules: number[]) => boolean;
 
+function isFirstPreviousTokenOfType(
+    tokenStream: TokenStream,
+    dictionary: TokenDictionary,
+    tokenIndex: number,
+    tokenType: number,
+): boolean {
+    let currentIndex = tokenIndex - 1;
+    let token;
+
+    do {
+        token = tokenStream.get(currentIndex);
+        if (token?.type === tokenType) {
+            return true;
+        }
+        currentIndex--;
+    } while (token?.type === dictionary.SPACE);
+
+    return false;
+}
+
 function getRuleCheckHelpers(ruleList: c3.RuleList): {
     anyRuleInList: AnyRuleInList;
     allRulesInList: AllRulesInList;
@@ -243,11 +264,9 @@ function getExternalDatasourceSuggestions({
     ]);
 }
 
-function getTableIndexesSuggestions({
-    anyRuleInList,
-}: GetParticularSuggestionProps): boolean | undefined {
+function checkShouldSuggestTableIndexes({anyRuleInList}: GetParticularSuggestionProps): boolean {
     if (!anyRuleInList(YQLParser.RULE_an_id)) {
-        return;
+        return false;
     }
 
     return anyRuleInList([
@@ -256,11 +275,11 @@ function getTableIndexesSuggestions({
     ]);
 }
 
-function getColumnsSuggestions({
+function checkShouldSuggestColumns({
     anyRuleInList,
     tokenStream,
     cursorTokenIndex,
-}: GetParticularSuggestionProps): boolean | undefined {
+}: GetParticularSuggestionProps): boolean {
     if (
         !anyRuleInList([YQLParser.RULE_an_id, YQLParser.RULE_id_expr]) ||
         anyRuleInList([
@@ -270,7 +289,7 @@ function getColumnsSuggestions({
             YQLParser.RULE_lambda_stmt,
         ])
     ) {
-        return;
+        return false;
     }
 
     const existingColumnInSelect =
@@ -295,6 +314,36 @@ function getColumnsSuggestions({
     );
 }
 
+function checkShouldSuggestAllColumns(props: GetParticularSuggestionProps): boolean {
+    const shouldSuggestColumns = checkShouldSuggestColumns(props);
+    if (!shouldSuggestColumns) {
+        return false;
+    }
+    const {tokenStream, cursorTokenIndex, anyRuleInList, allRulesInList} = props;
+    const isIntoTable =
+        anyRuleInList([YQLParser.RULE_into_table_stmt, YQLParser.RULE_into_table_stmt_yq]) &&
+        anyRuleInList(YQLParser.RULE_into_values_source);
+    if (isIntoTable) {
+        return isFirstPreviousTokenOfType(
+            tokenStream,
+            tokenDictionary,
+            cursorTokenIndex,
+            YQLParser.LPAREN,
+        );
+    }
+
+    const isSelect = allRulesInList([YQLParser.RULE_select_stmt, YQLParser.RULE_result_column]);
+    if (isSelect) {
+        return isFirstPreviousTokenOfType(
+            tokenStream,
+            tokenDictionary,
+            cursorTokenIndex,
+            YQLParser.SELECT,
+        );
+    }
+    return false;
+}
+
 function getSimpleTypesSuggestions({
     anyRuleInList,
     allRulesInList,
@@ -354,9 +403,7 @@ function getAggregateFunctionsSuggestions({
     return;
 }
 
-function getTableHintsSuggestions({
-    allRulesInList,
-}: GetParticularSuggestionProps): boolean | undefined {
+function checkShouldSuggestTableHints({allRulesInList}: GetParticularSuggestionProps): boolean {
     return allRulesInList([YQLParser.RULE_an_id_hint, YQLParser.RULE_table_hint]);
 }
 
@@ -446,21 +493,23 @@ export function getGranularSuggestions(
     const suggestReplication = getReplicationSuggestions(props);
     const suggestExternalTable = getExternalTableSuggestions(props);
     const suggestExternalDatasource = getExternalDatasourceSuggestions(props);
-    const shouldSuggestTableIndexes = getTableIndexesSuggestions(props);
-    const shouldSuggestColumns = getColumnsSuggestions(props);
+    const shouldSuggestTableIndexes = checkShouldSuggestTableIndexes(props);
+    const shouldSuggestColumns = checkShouldSuggestColumns(props);
+    const shouldSuggestAllColumns = checkShouldSuggestAllColumns(props);
     const suggestSimpleTypes = getSimpleTypesSuggestions(props);
     const suggestPragmas = getPragmasSuggestions(props);
     const suggestUdfs = getUdfsSuggestions(props);
     const suggestTableFunctions = getTableFunctionsSuggestions(props);
     const suggestFunctions = getFunctionsSuggestions(props);
     const suggestAggregateFunctions = getAggregateFunctionsSuggestions(props);
-    const shouldSuggestTableHints = getTableHintsSuggestions(props);
+    const shouldSuggestTableHints = checkShouldSuggestTableHints(props);
     const suggestEntitySettings = getEntitySettingsSuggestions(props);
 
     return {
         suggestWindowFunctions,
         shouldSuggestTableIndexes,
         shouldSuggestColumns,
+        shouldSuggestAllColumns,
         shouldSuggestColumnAliases: shouldSuggestColumns,
         suggestSimpleTypes,
         suggestPragmas,

diff --git a/src/autocomplete/databases/yql/tests/yq/insert/insert.test.ts b/src/autocomplete/databases/yql/tests/yq/insert/insert.test.ts
@@ -1,5 +1,6 @@
 import {parseYqQueryWithCursor} from '../../../index';
-import {ColumnSuggestion, KeywordSuggestion} from '../../../../../shared/autocomplete-types';
+import {KeywordSuggestion} from '../../../../../shared/autocomplete-types';
+import {YQLColumnsSuggestion} from '../../../types';
 
 test('should suggest properly after INSERT', () => {
     const autocompleteResult = parseYqQueryWithCursor('INSERT |');
@@ -50,7 +51,7 @@ test('should suggest properly after table name with a bracket', () => {
         {value: 'FROM'},
         {value: 'SELECT'},
     ];
-    const columnSuggestion: ColumnSuggestion = {tables: [{name: 'test_table'}]};
+    const columnSuggestion: YQLColumnsSuggestion = {tables: [{name: 'test_table'}], all: true};
     expect(autocompleteResult.suggestKeywords).toEqual(keywordsSuggestion);
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestion);
 });

diff --git a/src/autocomplete/databases/yql/tests/yq/select/nested-statements.test.ts b/src/autocomplete/databases/yql/tests/yq/select/nested-statements.test.ts
@@ -1,5 +1,6 @@
 import {parseYqQueryWithCursor, parseYqQueryWithoutCursor} from '../../../index';
-import {ColumnSuggestion, KeywordSuggestion} from '../../../../../shared/autocomplete-types';
+import {KeywordSuggestion} from '../../../../../shared/autocomplete-types';
+import {YQLColumnsSuggestion} from '../../../types';
 
 test('should suggest nested SELECT', () => {
     const autocompleteResult = parseYqQueryWithCursor('SELECT * FROM (|');
@@ -10,7 +11,7 @@ test('should suggest nested SELECT', () => {
 
 test('should suggest table name for nested SELECT column', () => {
     const autocompleteResult = parseYqQueryWithCursor('SELECT * FROM (SELECT | FROM test_table');
-    const columnSuggestion: ColumnSuggestion = {tables: [{name: 'test_table'}]};
+    const columnSuggestion: YQLColumnsSuggestion = {tables: [{name: 'test_table'}], all: true};
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestion);
 });
@@ -19,7 +20,7 @@ test('should suggest table name for nested SELECT column between statements', ()
     const autocompleteResult = parseYqQueryWithCursor(
         'SELECT * FROM before_table; SELECT * FROM (SELECT | FROM test_table ; SELECT * FROM after_table;',
     );
-    const columnSuggestion: ColumnSuggestion = {tables: [{name: 'test_table'}]};
+    const columnSuggestion: YQLColumnsSuggestion = {tables: [{name: 'test_table'}], all: true};
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestion);
 });
@@ -28,7 +29,7 @@ test('should suggest table name for nested WHERE condition', () => {
     const autocompleteResult = parseYqQueryWithCursor(
         'SELECT * FROM (SELECT * FROM test_table WHERE |',
     );
-    const columnSuggestion: ColumnSuggestion = {tables: [{name: 'test_table'}]};
+    const columnSuggestion: YQLColumnsSuggestion = {tables: [{name: 'test_table'}]};
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestion);
 });
@@ -37,7 +38,7 @@ test('should suggest table name for nested JOIN condition', () => {
     const autocompleteResult = parseYqQueryWithCursor(
         'SELECT * FROM (SELECT * FROM test_table_1 t1 JOIN test_table_2 t2 ON |',
     );
-    const columnSuggestion: ColumnSuggestion = {
+    const columnSuggestion: YQLColumnsSuggestion = {
         tables: [
             {name: 'test_table_1', alias: 't1'},
             {name: 'test_table_2', alias: 't2'},
@@ -58,7 +59,7 @@ test('should suggest table name for double nested SELECT column', () => {
     const autocompleteResult = parseYqQueryWithCursor(
         'SELECT * FROM (SELECT * FROM (SELECT | FROM test_table',
     );
-    const columnSuggestion: ColumnSuggestion = {tables: [{name: 'test_table'}]};
+    const columnSuggestion: YQLColumnsSuggestion = {tables: [{name: 'test_table'}], all: true};
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestion);
 });
@@ -67,7 +68,7 @@ test('should suggest table name for double nested SELECT column between statemen
     const autocompleteResult = parseYqQueryWithCursor(
         'SELECT * FROM before_table; SELECT * FROM (SELECT * FROM (SELECT | FROM test_table ; SELECT * FROM after_table;',
     );
-    const columnSuggestion: ColumnSuggestion = {tables: [{name: 'test_table'}]};
+    const columnSuggestion: YQLColumnsSuggestion = {tables: [{name: 'test_table'}], all: true};
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestion);
 });
@@ -76,7 +77,7 @@ test('should suggest table name for double nested WHERE condition', () => {
     const autocompleteResult = parseYqQueryWithCursor(
         'SELECT * FROM (SELECT * FROM (SELECT * FROM test_table WHERE |',
     );
-    const columnSuggestion: ColumnSuggestion = {tables: [{name: 'test_table'}]};
+    const columnSuggestion: YQLColumnsSuggestion = {tables: [{name: 'test_table'}]};
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestion);
 });
@@ -85,7 +86,7 @@ test('should suggest table name for double nested JOIN condition', () => {
     const autocompleteResult = parseYqQueryWithCursor(
         'SELECT * FROM (SELECT * FROM (SELECT * FROM test_table_1 t1 JOIN test_table_2 t2 ON |',
     );
-    const columnSuggestion: ColumnSuggestion = {
+    const columnSuggestion: YQLColumnsSuggestion = {
         tables: [
             {name: 'test_table_1', alias: 't1'},
             {name: 'test_table_2', alias: 't2'},

diff --git a/src/autocomplete/databases/yql/tests/yq/select/select.test.ts b/src/autocomplete/databases/yql/tests/yq/select/select.test.ts
@@ -1,5 +1,6 @@
 import {parseYqQueryWithCursor, parseYqQueryWithoutCursor} from '../../../index';
-import {ColumnSuggestion, KeywordSuggestion} from '../../../../../shared/autocomplete-types';
+import {KeywordSuggestion} from '../../../../../shared/autocomplete-types';
+import {YQLColumnsSuggestion} from '../../../types';
 
 test('should suggest properly after SELECT', () => {
     const autocompleteResult = parseYqQueryWithCursor('SELECT |');
@@ -137,7 +138,7 @@ test('should suggest table hints after WITH', () => {
 
 test('should suggest table name for column', () => {
     const autocompleteResult = parseYqQueryWithCursor('SELECT | FROM test_table');
-    const columnSuggestions: ColumnSuggestion = {tables: [{name: 'test_table'}]};
+    const columnSuggestions: YQLColumnsSuggestion = {tables: [{name: 'test_table'}], all: true};
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);
 });
@@ -146,43 +147,50 @@ test('should suggest table name for column between statements', () => {
     const autocompleteResult = parseYqQueryWithCursor(
         'SELECT * FROM before_table; SELECT | FROM test_table ; SELECT * FROM after_table;',
     );
-    const columnSuggestions: ColumnSuggestion = {tables: [{name: 'test_table'}]};
+    const columnSuggestions: YQLColumnsSuggestion = {tables: [{name: 'test_table'}], all: true};
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);
 });
 
 test('should suggest table name and alias for column', () => {
     const autocompleteResult = parseYqQueryWithCursor('SELECT | FROM test_table t');
-    const columnSuggestions: ColumnSuggestion = {tables: [{name: 'test_table', alias: 't'}]};
+    const columnSuggestions: YQLColumnsSuggestion = {
+        tables: [{name: 'test_table', alias: 't'}],
+        all: true,
+    };
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);
 });
 
 test('should suggest table name and alias (with AS) for column', () => {
     const autocompleteResult = parseYqQueryWithCursor('SELECT | FROM test_table AS t');
-    const columnSuggestions: ColumnSuggestion = {tables: [{name: 'test_table', alias: 't'}]};
+    const columnSuggestions: YQLColumnsSuggestion = {
+        tables: [{name: 'test_table', alias: 't'}],
+        all: true,
+    };
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);
 });
 
 test('should suggest table name and alias for second column', () => {
     const autocompleteResult = parseYqQueryWithCursor('SELECT id, | FROM test_table AS t');
-    const columnSuggestions: ColumnSuggestion = {tables: [{name: 'test_table', alias: 't'}]};
+    const columnSuggestions: YQLColumnsSuggestion = {tables: [{name: 'test_table', alias: 't'}]};
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);
 });
 
 test('should suggest table name and alias for second column if first equals to keyword', () => {
     const autocompleteResult = parseYqQueryWithCursor('SELECT key, | FROM test_table AS t');
-    const columnSuggestions: ColumnSuggestion = {tables: [{name: 'test_table', alias: 't'}]};
+    const columnSuggestions: YQLColumnsSuggestion = {tables: [{name: 'test_table', alias: 't'}]};
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);
 });
 
 test('should suggest multiple table names for column', () => {
     const autocompleteResult = parseYqQueryWithCursor('SELECT | FROM test_table_1, test_table_2');
-    const columnSuggestions: ColumnSuggestion = {
+    const columnSuggestions: YQLColumnsSuggestion = {
         tables: [{name: 'test_table_1'}, {name: 'test_table_2'}],
+        all: true,
     };
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);
@@ -192,8 +200,9 @@ test('should suggest multiple table names for column between statements', () =>
     const autocompleteResult = parseYqQueryWithCursor(
         'SELECT * FROM before_table; SELECT | FROM test_table_1, test_table_2 ; SELECT * FROM after_table;',
     );
-    const columnSuggestions: ColumnSuggestion = {
+    const columnSuggestions: YQLColumnsSuggestion = {
         tables: [{name: 'test_table_1'}, {name: 'test_table_2'}],
+        all: true,
     };
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);
@@ -203,11 +212,12 @@ test('should suggest multiple table names and aliases for column', () => {
     const autocompleteResult = parseYqQueryWithCursor(
         'SELECT | FROM test_table_1 t1, test_table_2 t2',
     );
-    const columnSuggestions: ColumnSuggestion = {
+    const columnSuggestions: YQLColumnsSuggestion = {
         tables: [
             {name: 'test_table_1', alias: 't1'},
             {name: 'test_table_2', alias: 't2'},
         ],
+        all: true,
     };
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);
@@ -217,11 +227,12 @@ test('should suggest multiple table names and aliases (with AS) for column', ()
     const autocompleteResult = parseYqQueryWithCursor(
         'SELECT | FROM test_table_1 AS t1, test_table_2 AS t2',
     );
-    const columnSuggestions: ColumnSuggestion = {
+    const columnSuggestions: YQLColumnsSuggestion = {
         tables: [
             {name: 'test_table_1', alias: 't1'},
             {name: 'test_table_2', alias: 't2'},
         ],
+        all: true,
     };
 
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestions);

diff --git a/src/autocomplete/databases/yql/tests/yql/insert/insert.test.ts b/src/autocomplete/databases/yql/tests/yql/insert/insert.test.ts
@@ -1,5 +1,6 @@
 import {parseYqlQueryWithCursor} from '../../../index';
-import {ColumnSuggestion, KeywordSuggestion} from '../../../../../shared/autocomplete-types';
+import {KeywordSuggestion} from '../../../../../shared/autocomplete-types';
+import {YQLColumnsSuggestion} from '../../../types';
 
 test('should suggest properly after INSERT', () => {
     const autocompleteResult = parseYqlQueryWithCursor('INSERT |');
@@ -50,7 +51,15 @@ test('should suggest properly after table name with a bracket', () => {
         {value: 'FROM'},
         {value: 'SELECT'},
     ];
-    const columnSuggestion: ColumnSuggestion = {tables: [{name: 'test_table'}]};
+    const columnSuggestion: YQLColumnsSuggestion = {tables: [{name: 'test_table'}], all: true};
+    expect(autocompleteResult.suggestKeywords).toEqual(keywordsSuggestion);
+    expect(autocompleteResult.suggestColumns).toEqual(columnSuggestion);
+});
+
+test('should suggest properly after table name with a bracket and column', () => {
+    const autocompleteResult = parseYqlQueryWithCursor('INSERT INTO test_table( col1, |');
+    const keywordsSuggestion: KeywordSuggestion[] = [];
+    const columnSuggestion: YQLColumnsSuggestion = {tables: [{name: 'test_table'}]};
     expect(autocompleteResult.suggestKeywords).toEqual(keywordsSuggestion);
     expect(autocompleteResult.suggestColumns).toEqual(columnSuggestion);
 });