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

1 files changed, 110 insertions, 2 deletions

diff --git a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
index c65eba0e2d..4fb6dfea58 100644
--- a/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
+++ b/yql/essentials/sql/v1/complete/sql_complete_ut.cpp
@@ -1,5 +1,9 @@
 #include "sql_complete.h"
 
+#include <yql/essentials/sql/v1/complete/name/fallback/name_service.h>
+#include <yql/essentials/sql/v1/complete/name/static/name_service.h>
+
+#include <yql/essentials/sql/v1/lexer/lexer.h>
 #include <yql/essentials/sql/v1/lexer/antlr4_pure/lexer.h>
 #include <yql/essentials/sql/v1/lexer/antlr4_pure_ansi/lexer.h>
 
@@ -7,16 +11,48 @@
 
 using namespace NSQLComplete;
 
+class TDummyException: public std::runtime_error {
+public:
+    TDummyException()
+        : std::runtime_error("T_T") {
+    }
+};
+
+class TFailingNameService: public INameService {
+public:
+    TFuture<TNameResponse> Lookup(TNameRequest) override {
+        auto e = std::make_exception_ptr(TDummyException());
+        return NThreading::MakeErrorFuture<TNameResponse>(e);
+    }
+};
+
+class TSilentNameService: public INameService {
+public:
+    TFuture<TNameResponse> Lookup(TNameRequest) override {
+        auto promise = NThreading::NewPromise<TNameResponse>();
+        return promise.GetFuture();
+    }
+};
+
 Y_UNIT_TEST_SUITE(SqlCompleteTests) {
     using ECandidateKind::Keyword;
+    using ECandidateKind::TypeName;
 
-    ISqlCompletionEngine::TPtr MakeSqlCompletionEngineUT() {
+    TLexerSupplier MakePureLexerSupplier() {
         NSQLTranslationV1::TLexers lexers;
         lexers.Antlr4Pure = NSQLTranslationV1::MakeAntlr4PureLexerFactory();
         lexers.Antlr4PureAnsi = NSQLTranslationV1::MakeAntlr4PureAnsiLexerFactory();
-        return MakeSqlCompletionEngine([lexers = std::move(lexers)](bool ansi) {
+        return [lexers = std::move(lexers)](bool ansi) {
             return NSQLTranslationV1::MakeLexer(lexers, ansi, /* antlr4 = */ true, /* pure = */ true);
+        };
+    }
+
+    ISqlCompletionEngine::TPtr MakeSqlCompletionEngineUT() {
+        TLexerSupplier lexer = MakePureLexerSupplier();
+        INameService::TPtr names = MakeStaticNameService({
+            .Types = {"Uint64"},
         });
+        return MakeSqlCompletionEngine(std::move(lexer), std::move(names));
     }
 
     TVector<TCandidate> Complete(ISqlCompletionEngine::TPtr& engine, TStringBuf prefix) {
@@ -291,6 +327,31 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
         UNIT_ASSERT_VALUES_EQUAL(Complete(engine, {"UPSERT "}), expected);
     }
 
+    Y_UNIT_TEST(TypeName) {
+        TVector<TCandidate> expected = {
+            {Keyword, "CALLABLE"},
+            {Keyword, "DECIMAL"},
+            {Keyword, "DICT"},
+            {Keyword, "ENUM"},
+            {Keyword, "FLOW"},
+            {Keyword, "LIST"},
+            {Keyword, "OPTIONAL"},
+            {Keyword, "RESOURCE"},
+            {Keyword, "SET"},
+            {Keyword, "STREAM"},
+            {Keyword, "STRUCT"},
+            {Keyword, "TAGGED"},
+            {Keyword, "TUPLE"},
+            {Keyword, "VARIANT"},
+            {TypeName, "Uint64"},
+        };
+
+        auto engine = MakeSqlCompletionEngineUT();
+        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, {"CREATE TABLE table (id "}), expected);
+        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, {"SELECT CAST (1 AS "}), expected);
+        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, {"SELECT OPTIONAL<"}), expected);
+    }
+
     Y_UNIT_TEST(UTF8Wide) {
         auto engine = MakeSqlCompletionEngineUT();
         UNIT_ASSERT_VALUES_EQUAL(Complete(engine, {"\xF0\x9F\x98\x8A"}).size(), 0);
@@ -340,4 +401,51 @@ Y_UNIT_TEST_SUITE(SqlCompleteTests) {
         UNIT_ASSERT_VALUES_EQUAL_C(Complete(engine, "!;").size(), 0, "Lexer failing");
     }
 
+    Y_UNIT_TEST(DefaultNameSet) {
+        auto set = MakeDefaultNameSet();
+        auto service = MakeStaticNameService(std::move(set));
+        auto engine = MakeSqlCompletionEngine(MakePureLexerSupplier(), std::move(service));
+        {
+            TVector<TCandidate> expected = {
+                {TypeName, "Uint16"},
+                {TypeName, "Uint32"},
+                {TypeName, "Uint64"},
+                {TypeName, "Uint8"},
+                {TypeName, "Utf8"},
+                {TypeName, "Uuid"},
+            };
+            UNIT_ASSERT_VALUES_EQUAL(Complete(engine, {"SELECT OPTIONAL<U"}), expected);
+        }
+    }
+
+    Y_UNIT_TEST(OnFailingNameService) {
+        auto service = MakeHolder<TFailingNameService>();
+        auto engine = MakeSqlCompletionEngine(MakePureLexerSupplier(), std::move(service));
+        UNIT_ASSERT_NO_EXCEPTION(Complete(engine, {""}));
+        UNIT_ASSERT_EXCEPTION(Complete(engine, {"SELECT OPTIONAL<U"}), TDummyException);
+        UNIT_ASSERT_EXCEPTION(Complete(engine, {"SELECT CAST (1 AS "}).size(), TDummyException);
+    }
+
+    Y_UNIT_TEST(OnSilentNameService) {
+        auto silent = MakeHolder<TSilentNameService>();
+        auto deadlined = MakeDeadlinedNameService(std::move(silent), TDuration::MilliSeconds(1));
+
+        auto engine = MakeSqlCompletionEngine(MakePureLexerSupplier(), std::move(deadlined));
+        UNIT_ASSERT_EXCEPTION(Complete(engine, {"SELECT OPTIONAL<U"}), NThreading::TFutureException);
+        UNIT_ASSERT_EXCEPTION(Complete(engine, {"SELECT OPTIONAL<"}), NThreading::TFutureException);
+    }
+
+    Y_UNIT_TEST(OnFallbackNameService) {
+        auto silent = MakeHolder<TSilentNameService>();
+        auto primary = MakeDeadlinedNameService(std::move(silent), TDuration::MilliSeconds(1));
+
+        auto standby = MakeStaticNameService(MakeDefaultNameSet());
+
+        auto fallback = MakeFallbackNameService(std::move(primary), std::move(standby));
+
+        auto engine = MakeSqlCompletionEngine(MakePureLexerSupplier(), std::move(fallback));
+        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, {"SELECT CAST (1 AS U"}).size(), 6);
+        UNIT_ASSERT_VALUES_EQUAL(Complete(engine, {"SELECT CAST (1 AS "}).size(), 47);
+    }
+
 } // Y_UNIT_TEST_SUITE(SqlCompleteTests)