YQL-19747: Support statement blacklist

I leaved the YDB configuration empty as need your feedback.

---

- Related to `YQL-19747`
- Related to https://github.com/vityaman/ydb/issues/39

---

Pull Request resolved: https://github.com/ytsaurus/ytsaurus/pull/1289
commit_hash:408c7e0f101027fc6ca3d86d251b5639fcae07bb

9 files changed, 128 insertions, 11 deletions

diff --git a/yql/essentials/sql/v1/complete/antlr4/c3i.h b/yql/essentials/sql/v1/complete/antlr4/c3i.h
index ed69f3cf29b..89201e12562 100644
--- a/yql/essentials/sql/v1/complete/antlr4/c3i.h
+++ b/yql/essentials/sql/v1/complete/antlr4/c3i.h
@@ -36,6 +36,7 @@ namespace NSQLComplete {
         struct TConfig {
             std::unordered_set<TTokenId> IgnoredTokens;
             std::unordered_set<TRuleId> PreferredRules;
+            std::unordered_set<TRuleId> IgnoredRules;
         };
 
         virtual TC3Candidates Complete(TStringBuf text, size_t caretTokenIndex) = 0;
diff --git a/yql/essentials/sql/v1/complete/antlr4/c3t.h b/yql/essentials/sql/v1/complete/antlr4/c3t.h
index 1f71553e04f..2b49068ac8f 100644
--- a/yql/essentials/sql/v1/complete/antlr4/c3t.h
+++ b/yql/essentials/sql/v1/complete/antlr4/c3t.h
@@ -25,12 +25,17 @@ namespace NSQLComplete {
             , Tokens_(&Lexer_)
             , Parser_(&Tokens_)
             , CompletionCore_(&Parser_)
+            , IgnoredRules_(std::move(config.IgnoredRules))
         {
             Lexer_.removeErrorListeners();
             Parser_.removeErrorListeners();
 
             CompletionCore_.ignoredTokens = std::move(config.IgnoredTokens);
             CompletionCore_.preferredRules = std::move(config.PreferredRules);
+
+            for (TRuleId rule : IgnoredRules_) {
+                CompletionCore_.preferredRules.emplace(rule);
+            }
         }
 
         TC3Candidates Complete(TStringBuf text, size_t caretTokenIndex) override {
@@ -47,23 +52,36 @@ namespace NSQLComplete {
             Tokens_.fill();
         }
 
-        static TC3Candidates Converted(c3::CandidatesCollection candidates) {
+        TC3Candidates Converted(c3::CandidatesCollection candidates) const {
             TC3Candidates converted;
+
             for (auto& [token, following] : candidates.tokens) {
                 converted.Tokens.emplace_back(token, std::move(following));
             }
+
             for (auto& [rule, data] : candidates.rules) {
+                if (IsIgnored(rule, data.ruleList)) {
+                    continue;
+                }
+
                 converted.Rules.emplace_back(rule, std::move(data.ruleList));
                 converted.Rules.back().ParserCallStack.emplace_back(rule);
             }
+
             return converted;
         }
 
+        bool IsIgnored(TRuleId head, const std::vector<TRuleId> tail) const {
+            return IgnoredRules_.contains(head) ||
+                   AnyOf(tail, [this](TRuleId r) { return IgnoredRules_.contains(r); });
+        }
+
         antlr4::ANTLRInputStream Chars_;
         G::TLexer Lexer_;
         antlr4::BufferedTokenStream Tokens_;
         G::TParser Parser_;
         c3::CodeCompletionCore CompletionCore_;
+        std::unordered_set<TRuleId> IgnoredRules_;
     };
 
 } // namespace NSQLComplete
diff --git a/yql/essentials/sql/v1/complete/sql_complete.cpp b/yql/essentials/sql/v1/complete/sql_complete.cpp
index 00d346f0770..99388cb9304 100644
--- a/yql/essentials/sql/v1/complete/sql_complete.cpp
+++ b/yql/essentials/sql/v1/complete/sql_complete.cpp
@@ -1,5 +1,6 @@
 #include "sql_complete.h"
 
+#include <yql/essentials/sql/v1/complete/syntax/grammar.h>
 #include <yql/essentials/sql/v1/complete/text/word.h>
 #include <yql/essentials/sql/v1/complete/name/service/static/name_service.h>
 #include <yql/essentials/sql/v1/complete/syntax/local.h>
@@ -18,7 +19,7 @@ namespace NSQLComplete {
             INameService::TPtr names,
             ISqlCompletionEngine::TConfiguration configuration)
             : Configuration_(std::move(configuration))
-            , SyntaxAnalysis_(MakeLocalSyntaxAnalysis(lexer))
+            , SyntaxAnalysis_(MakeLocalSyntaxAnalysis(lexer, Configuration_.IgnoredRules))
             , GlobalAnalysis_(MakeGlobalAnalysis())
             , Names_(std::move(names))
         {
@@ -213,6 +214,41 @@ namespace NSQLComplete {
         INameService::TPtr Names_;
     };
 
+    ISqlCompletionEngine::TConfiguration MakeYDBConfiguration() {
+        return {};
+    }
+
+    ISqlCompletionEngine::TConfiguration MakeYQLConfiguration() {
+        std::unordered_set<std::string> whitelist = {
+            "lambda_stmt",
+            "sql_stmt",
+            "pragma_stmt",
+            "select_stmt",
+            "named_nodes_stmt",
+            "drop_table_stmt",
+            "use_stmt",
+            "into_table_stmt",
+            "commit_stmt",
+            "declare_stmt",
+            "import_stmt",
+            "export_stmt",
+            "do_stmt",
+            "define_action_or_subquery_stmt",
+            "if_stmt",
+            "for_stmt",
+            "values_stmt",
+        };
+
+        ISqlCompletionEngine::TConfiguration config;
+        for (const std::string& name : GetSqlGrammar().GetAllRules()) {
+            if (name.ends_with("_stmt") && !whitelist.contains(name)) {
+                config.IgnoredRules.emplace(name);
+            }
+        }
+
+        return config;
+    }
+
     ISqlCompletionEngine::TPtr MakeSqlCompletionEngine(
         TLexerSupplier lexer,
         INameService::TPtr names,
diff --git a/yql/essentials/sql/v1/complete/sql_complete.h b/yql/essentials/sql/v1/complete/sql_complete.h
index 1bc2c0ecf4e..a600efcdac9 100644
--- a/yql/essentials/sql/v1/complete/sql_complete.h
+++ b/yql/essentials/sql/v1/complete/sql_complete.h
@@ -8,6 +8,7 @@
 
 #include <util/generic/string.h>
 #include <util/generic/vector.h>
+#include <util/generic/hash_set.h>
 
 namespace NSQLComplete {
 
@@ -46,6 +47,7 @@ namespace NSQLComplete {
 
         struct TConfiguration {
             size_t Limit = 256;
+            THashSet<TString> IgnoredRules;
         };
 
         virtual ~ISqlCompletionEngine() = default;
@@ -55,6 +57,10 @@ namespace NSQLComplete {
 
     using TLexerSupplier = std::function<NSQLTranslation::ILexer::TPtr(bool ansi)>;
 
+    ISqlCompletionEngine::TConfiguration MakeYDBConfiguration();
+
+    ISqlCompletionEngine::TConfiguration MakeYQLConfiguration();
+
     ISqlCompletionEngine::TPtr MakeSqlCompletionEngine(
         TLexerSupplier lexer,
         INameService::TPtr names,
diff --git a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
index d3a39d1c503..dc5690853f5 100644
--- a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
+++ b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
@@ -1,5 +1,6 @@
 #include "sql_complete.h"
 
+#include <yql/essentials/sql/v1/complete/syntax/grammar.h>
 #include <yql/essentials/sql/v1/complete/name/cluster/static/discovery.h>
 #include <yql/essentials/sql/v1/complete/name/object/dispatch/schema.h>
 #include <yql/essentials/sql/v1/complete/name/object/simple/schema.h>
@@ -1176,4 +1177,30 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
         }
     }
 
+    Y_UNIT_TEST(IgnoredRules) {
+        auto lexer = MakePureLexerSupplier();
+
+        TNameSet names;
+        TFrequencyData frequency;
+        auto service = MakeStaticNameService(names, MakeDefaultRanking(frequency));
+
+        {
+            auto engine = MakeSqlCompletionEngine(lexer, service);
+            UNIT_ASSERT_UNEQUAL(Complete(engine, {"UPDA"}).size(), 0);
+            UNIT_ASSERT_UNEQUAL(Complete(engine, {"DELE"}).size(), 0);
+            UNIT_ASSERT_UNEQUAL(Complete(engine, {"ROLL"}).size(), 0);
+            UNIT_ASSERT_UNEQUAL(Complete(engine, {"INSE"}).size(), 0);
+            UNIT_ASSERT_UNEQUAL(Complete(engine, {"SELE"}).size(), 0);
+        }
+
+        auto config = MakeYQLConfiguration();
+        auto engine = MakeSqlCompletionEngine(lexer, std::move(service), config);
+
+        UNIT_ASSERT_EQUAL(Complete(engine, {"UPDA"}).size(), 0);
+        UNIT_ASSERT_EQUAL(Complete(engine, {"DELE"}).size(), 0);
+        UNIT_ASSERT_EQUAL(Complete(engine, {"ROLL"}).size(), 0);
+        UNIT_ASSERT_UNEQUAL(Complete(engine, {"INSE"}).size(), 0);
+        UNIT_ASSERT_UNEQUAL(Complete(engine, {"SELE"}).size(), 0);
+    }
+
 } // Y_UNIT_TEST_SUITE(SqlCompleteTests)
diff --git a/yql/essentials/sql/v1/complete/syntax/grammar.cpp b/yql/essentials/sql/v1/complete/syntax/grammar.cpp
index c080fae5ae4..790f272db86 100644
--- a/yql/essentials/sql/v1/complete/syntax/grammar.cpp
+++ b/yql/essentials/sql/v1/complete/syntax/grammar.cpp
@@ -34,6 +34,18 @@ namespace NSQLComplete {
             return Parser_->getRuleNames().at(rule);
         }
 
+        TRuleId GetRuleId(std::string_view symbolized) const override {
+            TRuleId index = Parser_->getRuleIndex(std::string(symbolized));
+            if (index == INVALID_INDEX) {
+                ythrow yexception() << "Rule \"" << symbolized << "\" not found";
+            }
+            return index;
+        }
+
+        const std::vector<std::string>& GetAllRules() const override {
+            return Parser_->getRuleNames();
+        }
+
     private:
         static THolder<antlr4::Parser> MakeDummyParser() {
             return MakeHolder<NALADefaultAntlr4::SQLv1Antlr4Parser>(nullptr);
diff --git a/yql/essentials/sql/v1/complete/syntax/grammar.h b/yql/essentials/sql/v1/complete/syntax/grammar.h
index b128129e95f..4a259493dbc 100644
--- a/yql/essentials/sql/v1/complete/syntax/grammar.h
+++ b/yql/essentials/sql/v1/complete/syntax/grammar.h
@@ -6,6 +6,7 @@
 
 #include <unordered_set>
 #include <string>
+#include <string_view>
 
 #ifdef TOKEN_QUERY // Conflict with the winnt.h
     #undef TOKEN_QUERY
@@ -21,6 +22,8 @@ namespace NSQLComplete {
     public:
         virtual const antlr4::dfa::Vocabulary& GetVocabulary() const = 0;
         virtual const std::string& SymbolizedRule(TRuleId rule) const = 0;
+        virtual TRuleId GetRuleId(std::string_view symbolized) const = 0;
+        virtual const std::vector<std::string>& GetAllRules() const = 0;
         virtual const std::unordered_set<TTokenId>& GetAllTokens() const = 0;
         virtual const std::unordered_set<TTokenId>& GetKeywordTokens() const = 0;
         virtual const std::unordered_set<TTokenId>& GetPunctuationTokens() const = 0;
diff --git a/yql/essentials/sql/v1/complete/syntax/local.cpp b/yql/essentials/sql/v1/complete/syntax/local.cpp
index a9803feae85..772f5b78dd5 100644
--- a/yql/essentials/sql/v1/complete/syntax/local.cpp
+++ b/yql/essentials/sql/v1/complete/syntax/local.cpp
@@ -49,10 +49,11 @@ namespace NSQLComplete {
             TDefaultYQLGrammar>;
 
     public:
-        explicit TSpecializedLocalSyntaxAnalysis(TLexerSupplier lexer)
+        explicit TSpecializedLocalSyntaxAnalysis(
+            TLexerSupplier lexer, const THashSet<TString>& IgnoredRules)
             : Grammar_(&GetSqlGrammar())
             , Lexer_(lexer(/* ansi = */ IsAnsiLexer))
-            , C3_(ComputeC3Config())
+            , C3_(ComputeC3Config(IgnoredRules))
         {
         }
 
@@ -105,10 +106,11 @@ namespace NSQLComplete {
         }
 
     private:
-        IC3Engine::TConfig ComputeC3Config() const {
+        IC3Engine::TConfig ComputeC3Config(const THashSet<TString>& IgnoredRules) const {
             return {
                 .IgnoredTokens = ComputeIgnoredTokens(),
                 .PreferredRules = ComputePreferredRules(),
+                .IgnoredRules = ComputeIgnoredRules(IgnoredRules),
             };
         }
 
@@ -127,6 +129,15 @@ namespace NSQLComplete {
             return GetC3PreferredRules();
         }
 
+        std::unordered_set<TRuleId> ComputeIgnoredRules(const THashSet<TString>& IgnoredRules) const {
+            std::unordered_set<TRuleId> ignored;
+            ignored.reserve(IgnoredRules.size());
+            for (const auto& ruleName : IgnoredRules) {
+                ignored.emplace(Grammar_->GetRuleId(ruleName));
+            }
+            return ignored;
+        }
+
         TC3Candidates C3Complete(TCompletionInput statement, const TCursorTokenContext& context) {
             auto enclosing = context.Enclosing();
 
@@ -311,9 +322,10 @@ namespace NSQLComplete {
 
     class TLocalSyntaxAnalysis: public ILocalSyntaxAnalysis {
     public:
-        explicit TLocalSyntaxAnalysis(TLexerSupplier lexer)
-            : DefaultEngine_(lexer)
-            , AnsiEngine_(lexer)
+        explicit TLocalSyntaxAnalysis(
+            TLexerSupplier lexer, const THashSet<TString>& IgnoredRules)
+            : DefaultEngine_(lexer, IgnoredRules)
+            , AnsiEngine_(lexer, IgnoredRules)
         {
         }
 
@@ -335,8 +347,9 @@ namespace NSQLComplete {
         TSpecializedLocalSyntaxAnalysis</* IsAnsiLexer = */ true> AnsiEngine_;
     };
 
-    ILocalSyntaxAnalysis::TPtr MakeLocalSyntaxAnalysis(TLexerSupplier lexer) {
-        return MakeHolder<TLocalSyntaxAnalysis>(lexer);
+    ILocalSyntaxAnalysis::TPtr MakeLocalSyntaxAnalysis(
+        TLexerSupplier lexer, const THashSet<TString>& IgnoredRules) {
+        return MakeHolder<TLocalSyntaxAnalysis>(lexer, IgnoredRules);
     }
 
 } // namespace NSQLComplete
diff --git a/yql/essentials/sql/v1/complete/syntax/local.h b/yql/essentials/sql/v1/complete/syntax/local.h
index 635485d2b7c..0017afa0684 100644
--- a/yql/essentials/sql/v1/complete/syntax/local.h
+++ b/yql/essentials/sql/v1/complete/syntax/local.h
@@ -66,6 +66,7 @@ namespace NSQLComplete {
         virtual ~ILocalSyntaxAnalysis() = default;
     };
 
-    ILocalSyntaxAnalysis::TPtr MakeLocalSyntaxAnalysis(TLexerSupplier lexer);
+    ILocalSyntaxAnalysis::TPtr MakeLocalSyntaxAnalysis(
+        TLexerSupplier lexer, const THashSet<TString>& IgnoredRules);
 
 } // namespace NSQLComplete