YQL-19747 Complete after PRAGMA and multi-token names

- [x] Complete after PRAGMA
- [x] Complete multi-token names correctly, for example, `yt.` returns only `DisableStrict`, not `yt.DisableStrict` and `DateTime::` returns `Split`, not `DateTime::Split`.

I tried to implement it using `CompletedToken` edition, but not all completion environments support candidates with various `contextLen` (`Replxx` does not). So I decided that completions should rewrite only the current token, not sequences. For example, on `DateTime::Spl` rewrite only `Spl`. It makes sense as multi-token names have some namespace separated by a punctuation, so used types only namespace and gets names inside of it.

---

Pull Request resolved: https://github.com/ytsaurus/ytsaurus/pull/1181
commit_hash:9d8967ac43b9348f6dbb53837d92a9dcc9b51f48

4 files changed, 89 insertions, 5 deletions

diff --git a/yql/essentials/sql/v1/complete/syntax/local.cpp b/yql/essentials/sql/v1/complete/syntax/local.cpp
index cac43e5a320..28ace474f8e 100644
--- a/yql/essentials/sql/v1/complete/syntax/local.cpp
+++ b/yql/essentials/sql/v1/complete/syntax/local.cpp
@@ -53,10 +53,14 @@ namespace NSQLComplete {
             }
 
             auto candidates = C3.Complete(prefix);
+
+            NSQLTranslation::TParsedTokenList tokens = Tokenized(prefix);
+
             return {
                 .Keywords = SiftedKeywords(candidates),
+                .Pragma = PragmaMatch(tokens, candidates),
                 .IsTypeName = IsTypeNameMatched(candidates),
-                .IsFunctionName = IsFunctionNameMatched(candidates),
+                .Function = FunctionMatch(tokens, candidates),
             };
         }
 
@@ -131,16 +135,73 @@ namespace NSQLComplete {
             return name;
         }
 
+        std::optional<TLocalSyntaxContext::TPragma> PragmaMatch(
+            const NSQLTranslation::TParsedTokenList& tokens, const TC3Candidates& candidates) {
+            bool isMatched = AnyOf(candidates.Rules, [&](const TMatchedRule& rule) {
+                return IsLikelyPragmaStack(rule.ParserCallStack);
+            });
+            if (!isMatched) {
+                return std::nullopt;
+            }
+
+            TLocalSyntaxContext::TPragma pragma;
+            if (EndsWith(tokens, {"ID_PLAIN", "DOT"})) {
+                pragma.Namespace = tokens[tokens.size() - 2].Content;
+            } else if (EndsWith(tokens, {"ID_PLAIN", "DOT", ""})) {
+                pragma.Namespace = tokens[tokens.size() - 3].Content;
+            }
+            return pragma;
+        }
+
         bool IsTypeNameMatched(const TC3Candidates& candidates) {
             return AnyOf(candidates.Rules, [&](const TMatchedRule& rule) {
                 return IsLikelyTypeStack(rule.ParserCallStack);
             });
         }
 
-        bool IsFunctionNameMatched(const TC3Candidates& candidates) {
-            return AnyOf(candidates.Rules, [&](const TMatchedRule& rule) {
+        std::optional<TLocalSyntaxContext::TFunction> FunctionMatch(
+            const NSQLTranslation::TParsedTokenList& tokens, const TC3Candidates& candidates) {
+            bool isMatched = AnyOf(candidates.Rules, [&](const TMatchedRule& rule) {
                 return IsLikelyFunctionStack(rule.ParserCallStack);
             });
+            if (!isMatched) {
+                return std::nullopt;
+            }
+
+            TLocalSyntaxContext::TFunction function;
+            if (EndsWith(tokens, {"ID_PLAIN", "NAMESPACE"})) {
+                function.Namespace = tokens[tokens.size() - 2].Content;
+            } else if (EndsWith(tokens, {"ID_PLAIN", "NAMESPACE", ""})) {
+                function.Namespace = tokens[tokens.size() - 3].Content;
+            }
+            return function;
+        }
+
+        NSQLTranslation::TParsedTokenList Tokenized(const TStringBuf text) {
+            NSQLTranslation::TParsedTokenList tokens;
+            NYql::TIssues issues;
+            if (!NSQLTranslation::Tokenize(
+                    *Lexer_, TString(text), /* queryName = */ "",
+                    tokens, issues, /* maxErrors = */ 0)) {
+                return {};
+            }
+            Y_ENSURE(!tokens.empty() && tokens.back().Name == "EOF");
+            tokens.pop_back();
+            return tokens;
+        }
+
+        bool EndsWith(
+            const NSQLTranslation::TParsedTokenList& tokens,
+            const TVector<TStringBuf>& pattern) {
+            if (tokens.size() < pattern.size()) {
+                return false;
+            }
+            for (yssize_t i = tokens.ysize() - 1, j = pattern.ysize() - 1; 0 <= j; --i, --j) {
+                if (!pattern[j].empty() && tokens[i].Name != pattern[j]) {
+                    return false;
+                }
+            }
+            return true;
         }
 
         const ISqlGrammar* Grammar;
diff --git a/yql/essentials/sql/v1/complete/syntax/local.h b/yql/essentials/sql/v1/complete/syntax/local.h
index 79984d00e2b..24e78108e10 100644
--- a/yql/essentials/sql/v1/complete/syntax/local.h
+++ b/yql/essentials/sql/v1/complete/syntax/local.h
@@ -10,9 +10,18 @@
 namespace NSQLComplete {
 
     struct TLocalSyntaxContext {
+        struct TPragma {
+            TString Namespace;
+        };
+
+        struct TFunction {
+            TString Namespace;
+        };
+
         TVector<TString> Keywords;
+        std::optional<TPragma> Pragma;
         bool IsTypeName;
-        bool IsFunctionName;
+        std::optional<TFunction> Function;
     };
 
     class ILocalSyntaxAnalysis {
diff --git a/yql/essentials/sql/v1/complete/syntax/parser_call_stack.cpp b/yql/essentials/sql/v1/complete/syntax/parser_call_stack.cpp
index 57e058fa900..1bfcac47266 100644
--- a/yql/essentials/sql/v1/complete/syntax/parser_call_stack.cpp
+++ b/yql/essentials/sql/v1/complete/syntax/parser_call_stack.cpp
@@ -24,6 +24,11 @@ namespace NSQLComplete {
         RULE(Keyword_compat),
     };
 
+    const TVector<TRuleId> PragmaNameRules = {
+        RULE(Opt_id_prefix_or_type),
+        RULE(An_id),
+    };
+
     const TVector<TRuleId> TypeNameRules = {
         RULE(Type_name_simple),
         RULE(An_id_or_type),
@@ -62,6 +67,11 @@ namespace NSQLComplete {
         return Find(stack, rule) != std::end(stack);
     }
 
+    bool IsLikelyPragmaStack(const TParserCallStack& stack) {
+        return EndsWith({RULE(Pragma_stmt), RULE(Opt_id_prefix_or_type)}, stack) ||
+               EndsWith({RULE(Pragma_stmt), RULE(An_id)}, stack);
+    }
+
     bool IsLikelyTypeStack(const TParserCallStack& stack) {
         return EndsWith({RULE(Type_name_simple)}, stack) ||
                (Contains({RULE(Invoke_expr),
@@ -75,12 +85,14 @@ namespace NSQLComplete {
         return EndsWith({RULE(Unary_casual_subexpr), RULE(Id_expr)}, stack) ||
                EndsWith({RULE(Unary_casual_subexpr),
                          RULE(Atom_expr),
-                         RULE(An_id_or_type)}, stack);
+                         RULE(An_id_or_type)}, stack) ||
+               EndsWith({RULE(Atom_expr), RULE(Id_or_type)}, stack);
     }
 
     std::unordered_set<TRuleId> GetC3PreferredRules() {
         std::unordered_set<TRuleId> preferredRules;
         preferredRules.insert(std::begin(KeywordRules), std::end(KeywordRules));
+        preferredRules.insert(std::begin(PragmaNameRules), std::end(PragmaNameRules));
         preferredRules.insert(std::begin(TypeNameRules), std::end(TypeNameRules));
         preferredRules.insert(std::begin(FunctionNameRules), std::end(FunctionNameRules));
         return preferredRules;
diff --git a/yql/essentials/sql/v1/complete/syntax/parser_call_stack.h b/yql/essentials/sql/v1/complete/syntax/parser_call_stack.h
index 756586988db..94533bddaae 100644
--- a/yql/essentials/sql/v1/complete/syntax/parser_call_stack.h
+++ b/yql/essentials/sql/v1/complete/syntax/parser_call_stack.h
@@ -4,6 +4,8 @@
 
 namespace NSQLComplete {
 
+    bool IsLikelyPragmaStack(const TParserCallStack& stack);
+
     bool IsLikelyTypeStack(const TParserCallStack& stack);
 
     bool IsLikelyFunctionStack(const TParserCallStack& stack);