Merge pull request #16402 from ydb-platform/merge-libs-250328-1230

1 files changed, 22 insertions, 7 deletions

diff --git a/yql/essentials/sql/v1/complete/sql_context.cpp b/yql/essentials/sql/v1/complete/sql_context.cpp
index 2bd1a2af98..4cc809479b 100644
--- a/yql/essentials/sql/v1/complete/sql_context.cpp
+++ b/yql/essentials/sql/v1/complete/sql_context.cpp
@@ -4,14 +4,18 @@
 #include "sql_syntax.h"
 
 #include <yql/essentials/core/issue/yql_issue.h>
+
+#include <util/generic/algorithm.h>
+#include <util/stream/output.h>
+
+#ifdef TOKEN_QUERY // Conflict with the winnt.h
+#undef TOKEN_QUERY
+#endif
 #include <yql/essentials/parser/antlr_ast/gen/v1_antlr4/SQLv1Antlr4Lexer.h>
 #include <yql/essentials/parser/antlr_ast/gen/v1_antlr4/SQLv1Antlr4Parser.h>
 #include <yql/essentials/parser/antlr_ast/gen/v1_ansi_antlr4/SQLv1Antlr4Lexer.h>
 #include <yql/essentials/parser/antlr_ast/gen/v1_ansi_antlr4/SQLv1Antlr4Parser.h>
 
-#include <util/generic/algorithm.h>
-#include <util/stream/output.h>
-
 namespace NSQLComplete {
 
     template <bool IsAnsiLexer>
@@ -44,9 +48,10 @@ namespace NSQLComplete {
                 return {};
             }
 
-            auto tokens = C3.Complete(prefix);
+            auto candidates = C3.Complete(prefix);
             return {
-                .Keywords = SiftedKeywords(tokens),
+                .Keywords = SiftedKeywords(candidates),
+                .IsTypeName = IsTypeNameMatched(candidates),
             };
         }
 
@@ -68,12 +73,15 @@ namespace NSQLComplete {
 
         std::unordered_set<TRuleId> ComputePreferredRules() {
             const auto& keywordRules = Grammar->GetKeywordRules();
+            const auto& typeNameRules = Grammar->GetTypeNameRules();
 
             std::unordered_set<TRuleId> preferredRules;
 
             // Excludes tokens obtained from keyword rules
             preferredRules.insert(std::begin(keywordRules), std::end(keywordRules));
 
+            preferredRules.insert(std::begin(typeNameRules), std::end(typeNameRules));
+
             return preferredRules;
         }
 
@@ -97,12 +105,12 @@ namespace NSQLComplete {
             return true;
         }
 
-        TVector<TString> SiftedKeywords(const TVector<TSuggestedToken>& tokens) {
+        TVector<TString> SiftedKeywords(const TC3Candidates& candidates) {
             const auto& vocabulary = Grammar->GetVocabulary();
             const auto& keywordTokens = Grammar->GetKeywordTokens();
 
             TVector<TString> keywords;
-            for (const auto& token : tokens) {
+            for (const auto& token : candidates.Tokens) {
                 if (keywordTokens.contains(token.Number)) {
                     keywords.emplace_back(vocabulary.getDisplayName(token.Number));
                 }
@@ -110,6 +118,13 @@ namespace NSQLComplete {
             return keywords;
         }
 
+        bool IsTypeNameMatched(const TC3Candidates& candidates) {
+            const auto& typeNameRules = Grammar->GetTypeNameRules();
+            return FindIf(candidates.Rules, [&](const TMatchedRule& rule) {
+                       return Find(typeNameRules, rule.Index) != std::end(typeNameRules);
+                   }) != std::end(candidates.Rules);
+        }
+
         const ISqlGrammar* Grammar;
         NSQLTranslation::ILexer::TPtr Lexer_;
         TC3Engine<G> C3;