YQL-19747: Load schema from JSON for CLI

Now it is possible to pass clusters schema
via a JSON file using `-s` flag.
commit_hash:fd579430968881627c77d4586464952176134dc5

3 files changed, 51 insertions, 4 deletions

diff --git a/yql/essentials/sql/v1/complete/name/service/static/name_service.h b/yql/essentials/sql/v1/complete/name/service/static/name_service.h
index 531d3882ccc..b009ec5b27e 100644
--- a/yql/essentials/sql/v1/complete/name/service/static/name_service.h
+++ b/yql/essentials/sql/v1/complete/name/service/static/name_service.h
@@ -10,7 +10,6 @@ namespace NSQLComplete {
 
     INameService::TPtr MakeStaticNameService(TNameSet names, TFrequencyData frequency);
 
-    // TODO(YQL-19747): Migrate YDB CLI to method from pure TNameSet and TFrequencyData
     INameService::TPtr MakeStaticNameService(TNameSet names, IRanking::TPtr ranking);
 
 } // namespace NSQLComplete
diff --git a/yql/essentials/tools/yql_complete/ya.make b/yql/essentials/tools/yql_complete/ya.make
index 36d845dadab..07c7b7cffdc 100644
--- a/yql/essentials/tools/yql_complete/ya.make
+++ b/yql/essentials/tools/yql_complete/ya.make
@@ -4,9 +4,15 @@ PROGRAM()
 
 PEERDIR(
     library/cpp/getopt
+    library/cpp/iterator
     yql/essentials/sql/v1/complete
+    yql/essentials/sql/v1/complete/name/cluster/static
+    yql/essentials/sql/v1/complete/name/object/simple/static
     yql/essentials/sql/v1/complete/name/service/ranking
     yql/essentials/sql/v1/complete/name/service/static
+    yql/essentials/sql/v1/complete/name/service/cluster
+    yql/essentials/sql/v1/complete/name/service/schema
+    yql/essentials/sql/v1/complete/name/service/union
     yql/essentials/sql/v1/lexer/antlr4_pure
     yql/essentials/sql/v1/lexer/antlr4_pure_ansi
     yql/essentials/utils
diff --git a/yql/essentials/tools/yql_complete/yql_complete.cpp b/yql/essentials/tools/yql_complete/yql_complete.cpp
index 2a940c4247e..bfbf004a572 100644
--- a/yql/essentials/tools/yql_complete/yql_complete.cpp
+++ b/yql/essentials/tools/yql_complete/yql_complete.cpp
@@ -1,7 +1,12 @@
 #include <yql/essentials/sql/v1/complete/sql_complete.h>
+#include <yql/essentials/sql/v1/complete/name/cluster/static/discovery.h>
+#include <yql/essentials/sql/v1/complete/name/object/simple/static/schema_json.h>
 #include <yql/essentials/sql/v1/complete/name/service/ranking/frequency.h>
 #include <yql/essentials/sql/v1/complete/name/service/ranking/ranking.h>
 #include <yql/essentials/sql/v1/complete/name/service/static/name_service.h>
+#include <yql/essentials/sql/v1/complete/name/service/cluster/name_service.h>
+#include <yql/essentials/sql/v1/complete/name/service/schema/name_service.h>
+#include <yql/essentials/sql/v1/complete/name/service/union/name_service.h>
 
 #include <yql/essentials/sql/v1/lexer/antlr4_pure/lexer.h>
 #include <yql/essentials/sql/v1/lexer/antlr4_pure_ansi/lexer.h>
@@ -9,6 +14,9 @@
 #include <yql/essentials/utils/utf8.h>
 
 #include <library/cpp/getopt/last_getopt.h>
+#include <library/cpp/json/json_reader.h>
+#include <library/cpp/iterator/iterate_keys.h>
+#include <library/cpp/iterator/functools.h>
 
 #include <util/charset/utf8.h>
 #include <util/stream/file.h>
@@ -18,6 +26,15 @@ NSQLComplete::TFrequencyData LoadFrequencyDataFromFile(TString filepath) {
     return NSQLComplete::Pruned(NSQLComplete::ParseJsonFrequencyData(text));
 }
 
+NJson::TJsonMap LoadSchemaJsonFromFile(TString filepath) {
+    TString text = TUnbufferedFileInput(filepath).ReadAll();
+    NJson::TJsonMap map;
+    if (!NJson::ReadJsonTree(text, &map)) {
+        ythrow yexception() << "Failed to parse JSON: '" << text << "'";
+    }
+    return map;
+}
+
 NSQLComplete::TLexerSupplier MakePureLexerSupplier() {
     NSQLTranslationV1::TLexers lexers;
     lexers.Antlr4Pure = NSQLTranslationV1::MakeAntlr4PureLexerFactory();
@@ -40,10 +57,12 @@ int Run(int argc, char* argv[]) {
     TString inFileName;
     TString inQueryText;
     TString freqFileName;
+    TString schemaFileName;
     TMaybe<ui64> pos;
     opts.AddLongOption('i', "input", "input file").RequiredArgument("input").StoreResult(&inFileName);
     opts.AddLongOption('q', "query", "input query text").RequiredArgument("query").StoreResult(&inQueryText);
     opts.AddLongOption('f', "freq", "frequences file").StoreResult(&freqFileName);
+    opts.AddLongOption('s', "schema", "schema file").StoreResult(&schemaFileName);
     opts.AddLongOption('p', "pos", "position").StoreResult(&pos);
     opts.SetFreeArgsNum(0);
     opts.AddHelpOption();
@@ -72,12 +91,35 @@ int Run(int argc, char* argv[]) {
     } else {
         frequency = LoadFrequencyDataFromFile(freqFileName);
     }
+    auto ranking = NSQLComplete::MakeDefaultRanking(frequency);
+
+    TVector<NSQLComplete::INameService::TPtr> services;
+
+    services.emplace_back(
+        NSQLComplete::MakeStaticNameService(
+            NSQLComplete::LoadDefaultNameSet(), ranking));
+
+    if (!schemaFileName.empty()) {
+        NJson::TJsonMap schema = LoadSchemaJsonFromFile(schemaFileName);
+
+        services.emplace_back(
+            NSQLComplete::MakeSchemaNameService(
+                NSQLComplete::MakeSimpleSchema(
+                    NSQLComplete::MakeStaticSimpleSchema(schema))));
+
+        auto clustersIt = NFuncTools::Filter(
+            [](const auto& x) { return !x.empty(); },
+            IterateKeys(schema.GetMapSafe()));
+        TVector<TString> clusters(clustersIt.begin(), clustersIt.end());
+
+        services.emplace_back(
+            NSQLComplete::MakeClusterNameService(
+                NSQLComplete::MakeStaticClusterDiscovery(std::move(clusters))));
+    }
 
     auto engine = NSQLComplete::MakeSqlCompletionEngine(
         MakePureLexerSupplier(),
-        NSQLComplete::MakeStaticNameService(
-            NSQLComplete::LoadDefaultNameSet(),
-            std::move(frequency)));
+        NSQLComplete::MakeUnionNameService(std::move(services), ranking));
 
     if (!NYql::IsUtf8(queryString)) {
         ythrow yexception() << "provided input is not UTF encoded";