From a2ebf02f0ce5603c95b00dab75c9ae1f3bbbd99d Mon Sep 17 00:00:00 2001
From: Elena Makarova <el-makarova@el-makarova-dev.sas.yp-c.yandex.net>
Date: Tue, 4 Jun 2024 14:49:07 +0300
Subject: [PATCH] fix(YQL): add replication and topicConsumer

---
 src/autocomplete/autocomplete-types.ts          |  3 ++-
 src/autocomplete/databases/yql/helpers.ts       | 14 +++++++++++++-
 .../tests/yql/create/create-replication.test.ts | 17 +++++++++++++++++
 .../yql/tests/yql/create/create-topic.test.ts   | 16 ++++++++++++++++
 4 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/src/autocomplete/autocomplete-types.ts b/src/autocomplete/autocomplete-types.ts
index 2e1b66e9..ad940cb2 100644
--- a/src/autocomplete/autocomplete-types.ts
+++ b/src/autocomplete/autocomplete-types.ts
@@ -59,7 +59,8 @@ export type YQLEntity =
     | 'topic'
     | 'group'
     | 'user'
-    | 'tableIndex';
+    | 'tableIndex'
+    | 'topicConsumer';
 
 export interface YqlAutocompleteResult extends AutocompleteResultBase {
     suggestTableIndexes?: TableIndexSuggestion;
diff --git a/src/autocomplete/databases/yql/helpers.ts b/src/autocomplete/databases/yql/helpers.ts
index 0dc23b34..4d634fc9 100644
--- a/src/autocomplete/databases/yql/helpers.ts
+++ b/src/autocomplete/databases/yql/helpers.ts
@@ -361,7 +361,13 @@ function getEntitySettingsSuggestions({
     allRulesInList,
     anyRuleInList,
 }: GetParticularSuggestionProps): YQLEntity | undefined {
-    if (anyRuleInList([YQLParser.RULE_table_setting_value, YQLParser.RULE_topic_setting_value])) {
+    if (
+        anyRuleInList([
+            YQLParser.RULE_table_setting_value,
+            YQLParser.RULE_topic_setting_value,
+            YQLParser.RULE_topic_consumer_setting_value,
+        ])
+    ) {
         return;
     }
     if (allRulesInList([YQLParser.RULE_with_table_settings, YQLParser.RULE_an_id])) {
@@ -379,6 +385,12 @@ function getEntitySettingsSuggestions({
     if (allRulesInList([YQLParser.RULE_with_topic_settings, YQLParser.RULE_an_id])) {
         return 'topic';
     }
+    if (allRulesInList([YQLParser.RULE_topic_consumer_with_settings, YQLParser.RULE_an_id])) {
+        return 'topicConsumer';
+    }
+    if (allRulesInList([YQLParser.RULE_replication_settings, YQLParser.RULE_an_id])) {
+        return 'replication';
+    }
     return;
 }
 
diff --git a/src/autocomplete/databases/yql/tests/yql/create/create-replication.test.ts b/src/autocomplete/databases/yql/tests/yql/create/create-replication.test.ts
index 223bf138..2eb7d12e 100644
--- a/src/autocomplete/databases/yql/tests/yql/create/create-replication.test.ts
+++ b/src/autocomplete/databases/yql/tests/yql/create/create-replication.test.ts
@@ -29,3 +29,20 @@ test('should suggest properly after FOR', () => {
     expect(autocompleteResult.suggestKeywords).toEqual(keywordsSuggestion);
     expect(autocompleteResult.suggestEntity).toEqual(['table']);
 });
+test('should suggest properly after FOR', () => {
+    const autocompleteResult = parseYqlQueryWithCursor(
+        'CREATE ASYNC REPLICATION test FOR target AS target |',
+    );
+
+    const keywordsSuggestion: KeywordSuggestion[] = [{value: 'WITH'}];
+    expect(autocompleteResult.suggestKeywords).toEqual(keywordsSuggestion);
+});
+test('should suggest properly after WITH', () => {
+    const autocompleteResult = parseYqlQueryWithCursor(
+        'CREATE ASYNC REPLICATION test FOR target AS target WITH (|',
+    );
+
+    const keywordsSuggestion: KeywordSuggestion[] = [];
+    expect(autocompleteResult.suggestKeywords).toEqual(keywordsSuggestion);
+    expect(autocompleteResult.suggestEntitySettings).toEqual('replication');
+});
diff --git a/src/autocomplete/databases/yql/tests/yql/create/create-topic.test.ts b/src/autocomplete/databases/yql/tests/yql/create/create-topic.test.ts
index 4e6240aa..e8749525 100644
--- a/src/autocomplete/databases/yql/tests/yql/create/create-topic.test.ts
+++ b/src/autocomplete/databases/yql/tests/yql/create/create-topic.test.ts
@@ -28,6 +28,22 @@ test('should suggest properly after consumer name', () => {
     const keywordsSuggestion: KeywordSuggestion[] = [{value: 'WITH'}];
     expect(autocompleteResult.suggestKeywords).toEqual(keywordsSuggestion);
 });
+test('should suggest properly after consumer WITH', () => {
+    const autocompleteResult = parseYqlQueryWithCursor(
+        'CREATE TOPIC test (CONSUMER test_consumer WITH (|',
+    );
+
+    const keywordsSuggestion: KeywordSuggestion[] = [];
+    expect(autocompleteResult.suggestKeywords).toEqual(keywordsSuggestion);
+    expect(autocompleteResult.suggestEntitySettings).toEqual('topicConsumer');
+});
+test('should suggest properly after consumer statement', () => {
+    const autocompleteResult = parseYqlQueryWithCursor(
+        'CREATE TOPIC test (CONSUMER test_consumer) |',
+    );
+    const keywordsSuggestion: KeywordSuggestion[] = [{value: 'WITH'}];
+    expect(autocompleteResult.suggestKeywords).toEqual(keywordsSuggestion);
+});
 test('should suggest properly after WITH', () => {
     const autocompleteResult = parseYqlQueryWithCursor(
         'CREATE TOPIC test (CONSUMER test_consumer) WITH (|',