YQL-19747: Respect quotable columns

Table can have not `ID_PLAIN` columns and we
should return them as `ID_QUOTED`.
commit_hash:f5850228a3191ab84e3740fc1c36cf8882b08c40

7 files changed, 78 insertions, 3 deletions

diff --git a/yql/essentials/sql/v1/complete/name_mapping.cpp b/yql/essentials/sql/v1/complete/name_mapping.cpp
index 541188ea8c8..e8e317e0b57 100644
--- a/yql/essentials/sql/v1/complete/name_mapping.cpp
+++ b/yql/essentials/sql/v1/complete/name_mapping.cpp
@@ -4,6 +4,25 @@
 
 namespace NSQLComplete {
 
+    namespace {
+
+        TString ToIdentifier(TString content, const TLocalSyntaxContext& context) {
+            if (IsPlain(content) && !context.IsQuoted.AtLhs && !context.IsQuoted.AtRhs) {
+                return content;
+            }
+
+            SubstGlobal(content, "`", "``");
+            if (!context.IsQuoted.AtLhs) {
+                content.prepend('`');
+            }
+            if (!context.IsQuoted.AtRhs) {
+                content.append('`');
+            }
+            return content;
+        }
+
+    } // namespace
+
     TCandidate ToCandidate(TKeyword name, TLocalSyntaxContext& context) {
         TVector<TString>& seq = context.Keywords[name.Content];
         seq.insert(std::begin(seq), name.Content);
@@ -98,6 +117,8 @@ namespace NSQLComplete {
     }
 
     TCandidate ToCandidate(TColumnName name, TLocalSyntaxContext& context) {
+        name.Indentifier = ToIdentifier(std::move(name.Indentifier), context);
+
         if (context.Column->Table.empty() && !name.TableAlias.empty()) {
             name.Indentifier.prepend('.');
             name.Indentifier.prepend(name.TableAlias);
@@ -113,7 +134,9 @@ namespace NSQLComplete {
         return {ECandidateKind::BindingName, std::move(name.Indentifier)};
     }
 
-    TCandidate ToCandidate(TUnknownName name) {
+    TCandidate ToCandidate(TUnknownName name, TLocalSyntaxContext& context) {
+        name.Content = ToIdentifier(std::move(name.Content), context);
+
         return {ECandidateKind::UnknownName, std::move(name.Content)};
     }
 
@@ -125,7 +148,8 @@ namespace NSQLComplete {
                 std::is_same_v<T, TFolderName> ||
                 std::is_same_v<T, TTableName> ||
                 std::is_same_v<T, TColumnName> ||
-                std::is_same_v<T, TBindingName>;
+                std::is_same_v<T, TBindingName> ||
+                std::is_same_v<T, TUnknownName>;
 
             if constexpr (IsContextSensitive) {
                 return ToCandidate(std::move(name), context);
diff --git a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
index ba422168652..2eab708acd3 100644
--- a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
+++ b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
@@ -121,7 +121,10 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
                 }}
             }},
             "saurus": { "type": "Folder", "entries": {
-                "maxim": { "type": "Table", "columns": {} }
+                "maxim": { "type": "Table", "columns": {
+                   "Y Q L": {},
+                   "o``o": {}
+                }}
             }}
         })";
 
@@ -1280,6 +1283,21 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
         }
     }
 
+    Y_UNIT_TEST(ColumnQuoted) {
+        auto engine = MakeSqlCompletionEngineUT();
+
+        TVector<TCandidate> expected = {
+            {ColumnName, "`Y Q L`"},
+            {ColumnName, "`o````o`"},
+        };
+
+        UNIT_ASSERT_VALUES_EQUAL(CompleteTop(2, engine, "SELECT # FROM saurus.maxim"), expected);
+        UNIT_ASSERT_VALUES_EQUAL(CompleteTop(1, engine, "SELECT Y# FROM saurus.maxim").at(0), expected.at(0));
+        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, "SELECT `# FROM saurus.maxim").size(), 0);
+        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, "SELECT `#` FROM saurus.maxim").size(), 0);
+        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, "SELECT `Y #` FROM saurus.maxim").size(), 0);
+    }
+
     Y_UNIT_TEST(Typing) {
         const auto queryUtf16 = TUtf16String::FromUtf8(
             "SELECT \n"
diff --git a/yql/essentials/sql/v1/complete/syntax/format.cpp b/yql/essentials/sql/v1/complete/syntax/format.cpp
index 3037faced5c..f85c93dd085 100644
--- a/yql/essentials/sql/v1/complete/syntax/format.cpp
+++ b/yql/essentials/sql/v1/complete/syntax/format.cpp
@@ -40,6 +40,10 @@ namespace NSQLComplete {
         return text;
     }
 
+    bool IsPlain(TStringBuf content) {
+        return GetSqlGrammar().IsPlainIdentifier(content);
+    }
+
     TString Quoted(TString content) {
         content.prepend('`');
         content.append('`');
diff --git a/yql/essentials/sql/v1/complete/syntax/format.h b/yql/essentials/sql/v1/complete/syntax/format.h
index 58e5d1f1e4a..ea3274ccf5d 100644
--- a/yql/essentials/sql/v1/complete/syntax/format.h
+++ b/yql/essentials/sql/v1/complete/syntax/format.h
@@ -6,6 +6,7 @@
 namespace NSQLComplete {
 
     TString FormatKeywords(const TVector<TString>& seq);
+    bool IsPlain(TStringBuf content);
     TString Quoted(TString content);
     TString Unquoted(TString content);
 
diff --git a/yql/essentials/sql/v1/complete/syntax/grammar.cpp b/yql/essentials/sql/v1/complete/syntax/grammar.cpp
index 3c0dd579e4f..9e15a8e2347 100644
--- a/yql/essentials/sql/v1/complete/syntax/grammar.cpp
+++ b/yql/essentials/sql/v1/complete/syntax/grammar.cpp
@@ -1,7 +1,12 @@
 #include "grammar.h"
 
+#include <yql/essentials/sql/v1/lexer/regex/regex.h>
 #include <yql/essentials/sql/v1/reflect/sql_reflect.h>
 
+#include <contrib/libs/re2/re2/re2.h>
+
+#include <util/generic/algorithm.h>
+
 namespace NSQLComplete {
 
     class TSqlGrammar: public ISqlGrammar {
@@ -11,6 +16,7 @@ namespace NSQLComplete {
             , AllTokens_(ComputeAllTokens())
             , KeywordTokens_(ComputeKeywordTokens(grammar))
             , PunctuationTokens_(ComputePunctuationTokens(grammar))
+            , IdPlainRegex_(ComputeIdPlainRegex(grammar))
         {
         }
 
@@ -52,6 +58,10 @@ namespace NSQLComplete {
             return Parser_->getRuleNames();
         }
 
+        bool IsPlainIdentifier(TStringBuf content) const override {
+            return RE2::FullMatch(content, IdPlainRegex_);
+        }
+
     private:
         static THolder<antlr4::Parser> MakeDummyParser() {
             return MakeHolder<NALADefaultAntlr4::SQLv1Antlr4Parser>(nullptr);
@@ -94,10 +104,23 @@ namespace NSQLComplete {
             return punctuationTokens;
         }
 
+        static TString ComputeIdPlainRegex(const NSQLReflect::TLexerGrammar& grammar) {
+            TVector<std::tuple<TString, TString>> regexes =
+                NSQLTranslationV1::MakeRegexByOtherName(grammar, /* ansi = */ false);
+
+            std::tuple<TString, TString>* regex = FindIfPtr(regexes, [&](const auto& x) {
+                return std::get<0>(x) == "ID_PLAIN";
+            });
+
+            Y_ENSURE(regex, "ID_PLAIN regex not found");
+            return std::get<1>(*regex);
+        }
+
         const THolder<antlr4::Parser> Parser_;
         const std::unordered_set<TTokenId> AllTokens_;
         const std::unordered_set<TTokenId> KeywordTokens_;
         const std::unordered_set<TTokenId> PunctuationTokens_;
+        const RE2 IdPlainRegex_;
     };
 
     const ISqlGrammar& GetSqlGrammar() {
diff --git a/yql/essentials/sql/v1/complete/syntax/grammar.h b/yql/essentials/sql/v1/complete/syntax/grammar.h
index 7e6bb398439..9db650b58c7 100644
--- a/yql/essentials/sql/v1/complete/syntax/grammar.h
+++ b/yql/essentials/sql/v1/complete/syntax/grammar.h
@@ -4,6 +4,8 @@
 
 #include <contrib/libs/antlr4_cpp_runtime/src/Vocabulary.h>
 
+#include <util/generic/string.h>
+
 #include <unordered_set>
 #include <string>
 #include <string_view>
@@ -28,6 +30,7 @@ namespace NSQLComplete {
         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;
+        virtual bool IsPlainIdentifier(TStringBuf content) const = 0;
     };
 
     const ISqlGrammar& GetSqlGrammar();
diff --git a/yql/essentials/sql/v1/complete/syntax/ya.make b/yql/essentials/sql/v1/complete/syntax/ya.make
index f60a05f7b59..7c6fa046f30 100644
--- a/yql/essentials/sql/v1/complete/syntax/ya.make
+++ b/yql/essentials/sql/v1/complete/syntax/ya.make
@@ -13,10 +13,12 @@ PEERDIR(
     yql/essentials/parser/lexer_common
     yql/essentials/sql/settings
     yql/essentials/sql/v1/lexer
+    yql/essentials/sql/v1/lexer/regex
     yql/essentials/sql/v1/reflect
     yql/essentials/sql/v1/complete/antlr4
     yql/essentials/sql/v1/complete/core
     yql/essentials/sql/v1/complete/text
+    contrib/libs/re2
 )
 
 END()