YQL-19747: Refactor SQL completion engine

vityaman · vityaman · commit 633c9e434ac3 · 2025-06-26T17:24:30.000+03:00
Some refactorings to compact the code:

- YQL-19747: Pass TParsedInput to visitor
- YQL-19747: Add TParsedInput
- YQL-19747: Cosmetics
- YQL-19747: Refactor configuration
- YQL-19747: Refactor name to candidate mapping
commit_hash:44dfe7dc7bcc627ef9c20696077f2d962a3014f6
diff --git a/yql/essentials/sql/v1/complete/analysis/global/column.cpp b/yql/essentials/sql/v1/complete/analysis/global/column.cpp
@@ -220,8 +220,8 @@ namespace NSQLComplete {
 
         class TVisitor: public TSQLv1NarrowingVisitor {
         public:
-            TVisitor(antlr4::TokenStream* tokens, size_t cursorPosition)
-                : TSQLv1NarrowingVisitor(tokens, cursorPosition)
+            TVisitor(const TParsedInput& input)
+                : TSQLv1NarrowingVisitor(input)
             {
             }
 
@@ -247,12 +247,9 @@ namespace NSQLComplete {
 
     } // namespace
 
-    TMaybe<TColumnContext> InferColumnContext(
-        SQLv1::Sql_queryContext* ctx,
-        antlr4::TokenStream* tokens,
-        size_t cursorPosition) {
+    TMaybe<TColumnContext> InferColumnContext(TParsedInput input) {
         // TODO: add utility `auto ToMaybe<T>(std::any any) -> TMaybe<T>`
-        std::any result = TVisitor(tokens, cursorPosition).visit(ctx);
+        std::any result = TVisitor(input).visit(input.SqlQuery);
         if (!result.has_value()) {
             return Nothing();
         }
diff --git a/yql/essentials/sql/v1/complete/analysis/global/column.h b/yql/essentials/sql/v1/complete/analysis/global/column.h
@@ -1,13 +1,10 @@
 #pragma once
 
 #include "global.h"
-#include "parse_tree.h"
+#include "input.h"
 
 namespace NSQLComplete {
 
-    TMaybe<TColumnContext> InferColumnContext(
-        SQLv1::Sql_queryContext* ctx,
-        antlr4::TokenStream* tokens,
-        size_t cursorPosition);
+    TMaybe<TColumnContext> InferColumnContext(TParsedInput input);
 
 } // namespace NSQLComplete
diff --git a/yql/essentials/sql/v1/complete/analysis/global/function.cpp b/yql/essentials/sql/v1/complete/analysis/global/function.cpp
@@ -8,8 +8,8 @@ namespace NSQLComplete {
 
         class TVisitor: public TSQLv1NarrowingVisitor {
         public:
-            TVisitor(antlr4::TokenStream* tokens, size_t cursorPosition)
-                : TSQLv1NarrowingVisitor(tokens, cursorPosition)
+            TVisitor(const TParsedInput& input)
+                : TSQLv1NarrowingVisitor(input)
             {
             }
 
@@ -37,11 +37,8 @@ namespace NSQLComplete {
 
     } // namespace
 
-    TMaybe<TString> EnclosingFunction(
-        SQLv1::Sql_queryContext* ctx,
-        antlr4::TokenStream* tokens,
-        size_t cursorPosition) {
-        std::any result = TVisitor(tokens, cursorPosition).visit(ctx);
+    TMaybe<TString> EnclosingFunction(TParsedInput input) {
+        std::any result = TVisitor(input).visit(input.SqlQuery);
         if (!result.has_value()) {
             return Nothing();
         }
diff --git a/yql/essentials/sql/v1/complete/analysis/global/function.h b/yql/essentials/sql/v1/complete/analysis/global/function.h
@@ -1,15 +1,12 @@
 #pragma once
 
-#include "parse_tree.h"
+#include "input.h"
 
 #include <util/generic/maybe.h>
 #include <util/generic/string.h>
 
 namespace NSQLComplete {
 
-    TMaybe<TString> EnclosingFunction(
-        SQLv1::Sql_queryContext* ctx,
-        antlr4::TokenStream* tokens,
-        size_t cursorPosition);
+    TMaybe<TString> EnclosingFunction(TParsedInput input);
 
 } // namespace NSQLComplete
diff --git a/yql/essentials/sql/v1/complete/analysis/global/global.cpp b/yql/essentials/sql/v1/complete/analysis/global/global.cpp
@@ -2,6 +2,7 @@
 
 #include "column.h"
 #include "function.h"
+#include "input.h"
 #include "named_node.h"
 #include "parse_tree.h"
 #include "use.h"
@@ -87,18 +88,10 @@ namespace NSQLComplete {
     template <bool IsAnsiLexer>
     class TSpecializedGlobalAnalysis: public IGlobalAnalysis {
     public:
-        using TDefaultYQLGrammar = TAntlrGrammar<
-            NALADefaultAntlr4::SQLv1Antlr4Lexer,
-            NALADefaultAntlr4::SQLv1Antlr4Parser>;
-
-        using TAnsiYQLGrammar = TAntlrGrammar<
-            NALAAnsiAntlr4::SQLv1Antlr4Lexer,
-            NALAAnsiAntlr4::SQLv1Antlr4Parser>;
-
-        using G = std::conditional_t<
+        using TLexer = std::conditional_t<
             IsAnsiLexer,
-            TAnsiYQLGrammar,
-            TDefaultYQLGrammar>;
+            NALAAnsiAntlr4::SQLv1Antlr4Lexer,
+            NALADefaultAntlr4::SQLv1Antlr4Lexer>;
 
         TSpecializedGlobalAnalysis()
             : Chars_()
@@ -128,11 +121,17 @@ namespace NSQLComplete {
 
             TGlobalContext ctx;
 
-            // TODO(YQL-19747): Add ~ParseContext(Tokens, ParseTree, CursorPosition)
-            ctx.Use = FindUseStatement(sqlQuery, &Tokens_, input.CursorPosition, env);
-            ctx.Names = CollectNamedNodes(sqlQuery, &Tokens_, input.CursorPosition);
-            ctx.EnclosingFunction = EnclosingFunction(sqlQuery, &Tokens_, input.CursorPosition);
-            ctx.Column = InferColumnContext(sqlQuery, &Tokens_, input.CursorPosition);
+            TParsedInput parsed = {
+                .Original = input,
+                .Tokens = &Tokens_,
+                .Parser = &Parser_,
+                .SqlQuery = sqlQuery,
+            };
+
+            ctx.Use = FindUseStatement(parsed, env);
+            ctx.Names = CollectNamedNodes(parsed);
+            ctx.EnclosingFunction = EnclosingFunction(parsed);
+            ctx.Column = InferColumnContext(parsed);
 
             if (ctx.Use && ctx.Column) {
                 EnrichTableClusters(*ctx.Column, *ctx.Use);
@@ -176,9 +175,9 @@ namespace NSQLComplete {
         }
 
         antlr4::ANTLRInputStream Chars_;
-        G::TLexer Lexer_;
+        TLexer Lexer_;
         antlr4::CommonTokenStream Tokens_;
-        TDefaultYQLGrammar::TParser Parser_;
+        SQLv1 Parser_;
     };
 
     class TGlobalAnalysis: public IGlobalAnalysis {
diff --git a/yql/essentials/sql/v1/complete/analysis/global/input.cpp b/yql/essentials/sql/v1/complete/analysis/global/input.cpp
@@ -0,0 +1 @@
+#include "input.h"
diff --git a/yql/essentials/sql/v1/complete/analysis/global/input.h b/yql/essentials/sql/v1/complete/analysis/global/input.h
@@ -0,0 +1,16 @@
+#pragma once
+
+#include "parse_tree.h"
+
+#include <yql/essentials/sql/v1/complete/core/input.h>
+
+namespace NSQLComplete {
+
+    struct TParsedInput {
+        TCompletionInput Original;
+        antlr4::CommonTokenStream* Tokens;
+        SQLv1* Parser;
+        SQLv1::Sql_queryContext* SqlQuery;
+    };
+
+} // namespace NSQLComplete
diff --git a/yql/essentials/sql/v1/complete/analysis/global/named_node.cpp b/yql/essentials/sql/v1/complete/analysis/global/named_node.cpp
@@ -14,11 +14,8 @@ namespace NSQLComplete {
 
         class TVisitor: public TSQLv1NarrowingVisitor {
         public:
-            TVisitor(
-                antlr4::TokenStream* tokens,
-                size_t cursorPosition,
-                THashSet<TString>* names)
-                : TSQLv1NarrowingVisitor(tokens, cursorPosition)
+            TVisitor(const TParsedInput& input, THashSet<TString>* names)
+                : TSQLv1NarrowingVisitor(input)
                 , Names_(names)
             {
             }
@@ -110,12 +107,9 @@ namespace NSQLComplete {
 
     } // namespace
 
-    TVector<TString> CollectNamedNodes(
-        SQLv1::Sql_queryContext* ctx,
-        antlr4::TokenStream* tokens,
-        size_t cursorPosition) {
+    TVector<TString> CollectNamedNodes(TParsedInput input) {
         THashSet<TString> names;
-        TVisitor(tokens, cursorPosition, &names).visit(ctx);
+        TVisitor(input, &names).visit(input.SqlQuery);
         return TVector<TString>(begin(names), end(names));
     }
 
diff --git a/yql/essentials/sql/v1/complete/analysis/global/named_node.h b/yql/essentials/sql/v1/complete/analysis/global/named_node.h
@@ -1,15 +1,12 @@
 #pragma once
 
-#include "parse_tree.h"
+#include "input.h"
 
 #include <util/generic/string.h>
 #include <util/generic/vector.h>
 
 namespace NSQLComplete {
 
-    TVector<TString> CollectNamedNodes(
-        SQLv1::Sql_queryContext* ctx,
-        antlr4::TokenStream* tokens,
-        size_t cursorPosition);
+    TVector<TString> CollectNamedNodes(TParsedInput input);
 
 } // namespace NSQLComplete
diff --git a/yql/essentials/sql/v1/complete/analysis/global/narrowing_visitor.cpp b/yql/essentials/sql/v1/complete/analysis/global/narrowing_visitor.cpp
@@ -2,9 +2,9 @@
 
 namespace NSQLComplete {
 
-    TSQLv1NarrowingVisitor::TSQLv1NarrowingVisitor(antlr4::TokenStream* tokens, size_t cursorPosition)
-        : Tokens_(tokens)
-        , CursorPosition_(cursorPosition)
+    TSQLv1NarrowingVisitor::TSQLv1NarrowingVisitor(const TParsedInput& input)
+        : Tokens_(input.Tokens)
+        , CursorPosition_(input.Original.CursorPosition)
     {
     }
 
diff --git a/yql/essentials/sql/v1/complete/analysis/global/narrowing_visitor.h b/yql/essentials/sql/v1/complete/analysis/global/narrowing_visitor.h
@@ -1,13 +1,13 @@
 #pragma once
 
 #include "base_visitor.h"
-#include "parse_tree.h"
+#include "input.h"
 
 namespace NSQLComplete {
 
     class TSQLv1NarrowingVisitor: public TSQLv1BaseVisitor {
     public:
-        TSQLv1NarrowingVisitor(antlr4::TokenStream* tokens, size_t cursorPosition);
+        TSQLv1NarrowingVisitor(const TParsedInput& input);
 
     protected:
         bool shouldVisitNextChild(antlr4::tree::ParseTree* node, const std::any& /*currentResult*/) override;
diff --git a/yql/essentials/sql/v1/complete/analysis/global/use.cpp b/yql/essentials/sql/v1/complete/analysis/global/use.cpp
@@ -9,11 +9,8 @@ namespace NSQLComplete {
 
         class TVisitor: public TSQLv1NarrowingVisitor {
         public:
-            TVisitor(
-                antlr4::TokenStream* tokens,
-                size_t cursorPosition,
-                const TEnvironment* env)
-                : TSQLv1NarrowingVisitor(tokens, cursorPosition)
+            TVisitor(const TParsedInput& input, const TEnvironment* env)
+                : TSQLv1NarrowingVisitor(input)
                 , Env_(env)
             {
             }
@@ -78,12 +75,9 @@ namespace NSQLComplete {
 
     } // namespace
 
-    TMaybe<TUseContext> FindUseStatement(
-        SQLv1::Sql_queryContext* ctx,
-        antlr4::TokenStream* tokens,
-        size_t cursorPosition,
-        const TEnvironment& env) {
-        std::any result = TVisitor(tokens, cursorPosition, &env).visit(ctx);
+    // TODO(YQL-19747): Use any to maybe conversion function
+    TMaybe<TUseContext> FindUseStatement(TParsedInput input, const TEnvironment& env) {
+        std::any result = TVisitor(input, &env).visit(input.SqlQuery);
         if (!result.has_value()) {
             return Nothing();
         }
diff --git a/yql/essentials/sql/v1/complete/analysis/global/use.h b/yql/essentials/sql/v1/complete/analysis/global/use.h
@@ -1,18 +1,14 @@
 #pragma once
 
 #include "global.h"
-#include "parse_tree.h"
+#include "input.h"
 
 #include <util/generic/ptr.h>
 #include <util/generic/maybe.h>
 #include <util/generic/string.h>
 
 namespace NSQLComplete {
 
-    TMaybe<TUseContext> FindUseStatement(
-        SQLv1::Sql_queryContext* ctx,
-        antlr4::TokenStream* tokens,
-        size_t cursorPosition,
-        const TEnvironment& env);
+    TMaybe<TUseContext> FindUseStatement(TParsedInput input, const TEnvironment& env);
 
 } // namespace NSQLComplete
diff --git a/yql/essentials/sql/v1/complete/analysis/global/ya.make b/yql/essentials/sql/v1/complete/analysis/global/ya.make
@@ -6,6 +6,7 @@ SRCS(
     evaluate.cpp
     function.cpp
     global.cpp
+    input.cpp
     named_node.cpp
     narrowing_visitor.cpp
     parse_tree.cpp
diff --git a/yql/essentials/sql/v1/complete/configuration.cpp b/yql/essentials/sql/v1/complete/configuration.cpp
@@ -0,0 +1,62 @@
+#include "configuration.h"
+
+#include <yql/essentials/sql/v1/complete/syntax/grammar.h>
+
+namespace NSQLComplete {
+
+    TConfiguration MakeConfiguration(THashSet<TString> allowedStmts) {
+        allowedStmts.emplace("sql_stmt");
+
+        TConfiguration config;
+        for (const std::string& name : GetSqlGrammar().GetAllRules()) {
+            if (name.ends_with("_stmt") && !allowedStmts.contains(name)) {
+                config.IgnoredRules_.emplace(name);
+            }
+        }
+        return config;
+    }
+
+    TConfiguration MakeYDBConfiguration() {
+        TConfiguration config;
+        config.IgnoredRules_ = {
+            "use_stmt",
+            "import_stmt",
+            "export_stmt",
+        };
+        return config;
+    }
+
+    TConfiguration MakeYQLConfiguration() {
+        auto config = MakeConfiguration(/* allowedStmts = */ {
+            "lambda_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",
+        });
+
+        config.DisabledPreviousByToken_ = {};
+
+        config.ForcedPreviousByToken_ = {
+            {"PARALLEL", {}},
+            {"TABLESTORE", {}},
+            {"FOR", {"EVALUATE"}},
+            {"IF", {"EVALUATE"}},
+            {"EXTERNAL", {"USING"}},
+        };
+
+        return config;
+    }
+
+} // namespace NSQLComplete
diff --git a/yql/essentials/sql/v1/complete/configuration.h b/yql/essentials/sql/v1/complete/configuration.h
@@ -0,0 +1,28 @@
+#pragma once
+
+#include <util/generic/string.h>
+#include <util/generic/hash.h>
+#include <util/generic/hash_set.h>
+
+namespace NSQLComplete {
+
+    struct TConfiguration {
+        friend class TSqlCompletionEngine;
+        friend TConfiguration MakeYDBConfiguration();
+        friend TConfiguration MakeYQLConfiguration();
+        friend TConfiguration MakeConfiguration(THashSet<TString> allowedStmts);
+
+    public:
+        size_t Limit = 256;
+
+    private:
+        THashSet<TString> IgnoredRules_;
+        THashMap<TString, THashSet<TString>> DisabledPreviousByToken_;
+        THashMap<TString, THashSet<TString>> ForcedPreviousByToken_;
+    };
+
+    TConfiguration MakeYDBConfiguration();
+
+    TConfiguration MakeYQLConfiguration();
+
+} // namespace NSQLComplete
diff --git a/yql/essentials/sql/v1/complete/name/service/name_service.h b/yql/essentials/sql/v1/complete/name/service/name_service.h
diff --git a/yql/essentials/sql/v1/complete/name/service/ranking/ranking.cpp b/yql/essentials/sql/v1/complete/name/service/ranking/ranking.cpp
diff --git a/yql/essentials/sql/v1/complete/name/service/schema/name_service.cpp b/yql/essentials/sql/v1/complete/name/service/schema/name_service.cpp
diff --git a/yql/essentials/sql/v1/complete/name_mapping.cpp b/yql/essentials/sql/v1/complete/name_mapping.cpp
diff --git a/yql/essentials/sql/v1/complete/name_mapping.h b/yql/essentials/sql/v1/complete/name_mapping.h
diff --git a/yql/essentials/sql/v1/complete/sql_complete.cpp b/yql/essentials/sql/v1/complete/sql_complete.cpp
diff --git a/yql/essentials/sql/v1/complete/sql_complete.h b/yql/essentials/sql/v1/complete/sql_complete.h
diff --git a/yql/essentials/sql/v1/complete/ya.make b/yql/essentials/sql/v1/complete/ya.make

Original file line number	Diff line number	Diff line change
`@@ -2,9 +2,9 @@`
`2`	`2`
`3`	`3`	`namespace NSQLComplete {`
`4`	`4`
`5`		`- TSQLv1NarrowingVisitor::TSQLv1NarrowingVisitor(antlr4::TokenStream* tokens, size_t cursorPosition)`
`6`		`- : Tokens_(tokens)`
`7`		`- , CursorPosition_(cursorPosition)`
	`5`	`+ TSQLv1NarrowingVisitor::TSQLv1NarrowingVisitor(const TParsedInput& input)`
	`6`	`+ : Tokens_(input.Tokens)`
	`7`	`+ , CursorPosition_(input.Original.CursorPosition)`
`8`	`8`	`{`
`9`	`9`	`}`
`10`	`10`