Interactive CLI improvements

Interactive CLI improvements

Pull Request resolved: https://github.com/ydb-platform/ydb/pull/371

8 files changed, 145 insertions, 31 deletions

diff --git a/ydb/public/lib/ydb_cli/commands/interactive/interactive_cli.cpp b/ydb/public/lib/ydb_cli/commands/interactive/interactive_cli.cpp
index c6c61661dc8..f47f207b968 100644
--- a/ydb/public/lib/ydb_cli/commands/interactive/interactive_cli.cpp
+++ b/ydb/public/lib/ydb_cli/commands/interactive/interactive_cli.cpp
@@ -5,6 +5,7 @@
 #include <util/folder/path.h>
 #include <util/folder/dirut.h>
 
+#include <ydb/public/lib/ydb_cli/common/query_stats.h>
 #include <ydb/public/lib/ydb_cli/commands/interactive/line_reader.h>
 #include <ydb/public/lib/ydb_cli/commands/ydb_service_scheme.h>
 #include <ydb/public/lib/ydb_cli/commands/ydb_service_table.h>
@@ -16,6 +17,48 @@ namespace NConsoleClient {
 
 namespace {
 
+std::string ToLower(std::string_view value) {
+    size_t value_size = value.size();
+    std::string result;
+    result.resize(value_size);
+
+    for (size_t i = 0; i < value_size; ++i) {
+        result[i] = std::tolower(value[i]);
+    }
+
+    return result;
+}
+
+template <typename Predicate>
+void TrimRight(std::string & value, Predicate predicate) {
+    while (!value.empty() && predicate(value.back()))
+        value.pop_back();
+}
+
+template <typename Predicate>
+void TrimLeft(std::string & value, Predicate predicate) {
+    size_t value_size = value.size();
+    size_t i = 0;
+
+    for (; i < value_size; ++i) {
+        if (!predicate(value[i])) {
+            break;
+        }
+    }
+
+    if (i != 0) {
+        value = value.substr(i);
+    }
+}
+
+void TrimSpacesRight(std::string & value) {
+    TrimRight(value, [](char character) { return std::isspace(character); });
+}
+
+void TrimSpacesLeft(std::string & value) {
+    TrimLeft(value, [](char character) { return std::isspace(character); });
+}
+
 struct Token {
     std::string_view data;
 };
@@ -68,9 +111,46 @@ std::vector<Token> Tokenize(std::string_view input) {
 }
 
 struct InteractiveCLIState {
-    std::string CollectStatsMode = "none";
+    NTable::ECollectQueryStatsMode CollectStatsMode = NTable::ECollectQueryStatsMode::None;
 };
 
+std::optional<NTable::ECollectQueryStatsMode> TryParseCollectStatsMode(const std::vector<Token> & tokens) {
+    size_t tokensSize = tokens.size();
+    if (tokensSize < 2) {
+        return {};
+    }
+
+    if (ToLower(tokens[0].data) != "set") {
+        return {};
+    }
+
+    std::string setQuery;
+    for (size_t i = 1; i < tokensSize; ++i) {
+        setQuery += tokens[i].data;
+    }
+
+    auto position = setQuery.find_first_of('=');
+    if (position == std::string::npos) {
+        return {};
+    }
+
+    std::string name = setQuery.substr(0, position);
+    std::string value = setQuery.substr(position + 1);
+
+    TrimSpacesLeft(name);
+    TrimSpacesRight(name);
+    if (name != "stats") {
+        return {};
+    }
+
+    TrimSpacesLeft(value);
+    TrimRight(value, [](char character) {
+        return std::isspace(character) || character == ';';
+    });
+
+    return NTable::ParseQueryStatsMode(value);
+}
+
 }
 
 TInteractiveCLI::TInteractiveCLI(TClientCommand::TConfig & config, std::string prompt)
@@ -79,15 +159,11 @@ TInteractiveCLI::TInteractiveCLI(TClientCommand::TConfig & config, std::string p
 {}
 
 void TInteractiveCLI::Run() {
-    std::vector<std::string> SQLWords = {"SELECT", "FROM", "WHERE", "GROUP", "ORDER" , "BY", "LIMIT", "OFFSET"};
+    std::vector<std::string> SQLWords = {"SELECT", "FROM", "WHERE", "GROUP", "ORDER" , "BY", "LIMIT", "OFFSET", "EXPLAIN", "AST"};
     std::vector<std::string> Words;
     for (auto & word : SQLWords) {
         Words.push_back(word);
-        for (auto & character : word) {
-            character = std::tolower(character);
-        }
-
-        Words.push_back(word);
+        Words.push_back(ToLower(word));
     }
 
     TFsPath homeDirPath(HomeDir);
@@ -107,15 +183,13 @@ void TInteractiveCLI::Run() {
             auto tokens = Tokenize(line);
             size_t tokensSize = tokens.size();
 
-            if (tokensSize == 4 && tokens[0].data == "SET" && tokens[1].data == "stats" && tokens[2].data == "=") {
-                interactiveCLIState.CollectStatsMode = tokens[3].data;
-                while (!interactiveCLIState.CollectStatsMode.empty() && interactiveCLIState.CollectStatsMode.back() == ';')
-                    interactiveCLIState.CollectStatsMode.pop_back();
+            if (auto collectStatsMode = TryParseCollectStatsMode(tokens)) {
+                interactiveCLIState.CollectStatsMode = *collectStatsMode;
                 continue;
             }
 
-            if (!tokens.empty() && tokens.front().data == "EXPLAIN") {
-                bool printAst = tokensSize >= 2 && tokens[1].data == "AST";
+            if (!tokens.empty() && ToLower(tokens.front().data) == "explain") {
+                bool printAst = tokensSize >= 2 && ToLower(tokens[1].data) == "ast";
                 size_t skipTokens = 1 + printAst;
                 TString explainQuery;
 
@@ -129,10 +203,15 @@ void TInteractiveCLI::Run() {
                 continue;
             }
 
-            TCommandYql yqlCommand(TString(line), TString(interactiveCLIState.CollectStatsMode));
+            TString queryStatsMode(NTable::QueryStatsModeToString(interactiveCLIState.CollectStatsMode));
+            TCommandYql yqlCommand(TString(line), queryStatsMode);
             yqlCommand.Run(Config);
         } catch (TYdbErrorException &error) {
             Cerr << error;
+        } catch (yexception & error) {
+            Cerr << error;
+        } catch (std::exception & error) {
+            Cerr << error.what();
         }
     }
 
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_service_scripting.cpp b/ydb/public/lib/ydb_cli/commands/ydb_service_scripting.cpp
index 5b3fec4d3e4..43e6f1ce538 100644
--- a/ydb/public/lib/ydb_cli/commands/ydb_service_scripting.cpp
+++ b/ydb/public/lib/ydb_cli/commands/ydb_service_scripting.cpp
@@ -113,7 +113,7 @@ int TCommandExecuteYqlScript::Run(TConfig& config) {
         PrintExplainResult(result);
     } else {
         NScripting::TExecuteYqlRequestSettings settings;
-        settings.CollectQueryStats(ParseQueryStatsMode(CollectStatsMode, NTable::ECollectQueryStatsMode::None));
+        settings.CollectQueryStats(ParseQueryStatsModeOrThrow(CollectStatsMode, NTable::ECollectQueryStatsMode::None));
 
         if (FlameGraphPath && (settings.CollectQueryStats_ != NTable::ECollectQueryStatsMode::Full
                                && settings.CollectQueryStats_ != NTable::ECollectQueryStatsMode::Profile)) {
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_service_table.cpp b/ydb/public/lib/ydb_cli/commands/ydb_service_table.cpp
index 7ae96a54776..37d8878cc79 100644
--- a/ydb/public/lib/ydb_cli/commands/ydb_service_table.cpp
+++ b/ydb/public/lib/ydb_cli/commands/ydb_service_table.cpp
@@ -429,7 +429,7 @@ int TCommandExecuteQuery::ExecuteDataQuery(TConfig& config) {
     auto defaultStatsMode = BasicStats ? NTable::ECollectQueryStatsMode::Basic : NTable::ECollectQueryStatsMode::None;
     NTable::TExecDataQuerySettings settings;
     settings.KeepInQueryCache(true);
-    settings.CollectQueryStats(ParseQueryStatsMode(CollectStatsMode, defaultStatsMode));
+    settings.CollectQueryStats(ParseQueryStatsModeOrThrow(CollectStatsMode, defaultStatsMode));
 
     NTable::TTxSettings txSettings;
     if (TxMode) {
@@ -535,7 +535,7 @@ int TCommandExecuteQuery::ExecuteScanQuery(TConfig& config) {
 
     auto defaultStatsMode = BasicStats ? NTable::ECollectQueryStatsMode::Basic : NTable::ECollectQueryStatsMode::None;
     NTable::TStreamExecScanQuerySettings settings;
-    settings.CollectQueryStats(ParseQueryStatsMode(CollectStatsMode, defaultStatsMode));
+    settings.CollectQueryStats(ParseQueryStatsModeOrThrow(CollectStatsMode, defaultStatsMode));
 
     NTable::TAsyncScanQueryPartIterator asyncResult;
     SetInterruptHandlers();
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_yql.cpp b/ydb/public/lib/ydb_cli/commands/ydb_yql.cpp
index f4ef469fc33..fabd6551018 100644
--- a/ydb/public/lib/ydb_cli/commands/ydb_yql.cpp
+++ b/ydb/public/lib/ydb_cli/commands/ydb_yql.cpp
@@ -93,7 +93,7 @@ int TCommandYql::RunCommand(TConfig& config, const TString& script) {
     NScripting::TScriptingClient client(driver);
 
     NScripting::TExecuteYqlRequestSettings settings;
-    settings.CollectQueryStats(ParseQueryStatsMode(CollectStatsMode, NTable::ECollectQueryStatsMode::None));
+    settings.CollectQueryStats(ParseQueryStatsModeOrThrow(CollectStatsMode, NTable::ECollectQueryStatsMode::None));
 
     if (FlameGraphPath && (settings.CollectQueryStats_ != NTable::ECollectQueryStatsMode::Full
                            && settings.CollectQueryStats_ != NTable::ECollectQueryStatsMode::Profile)) {
diff --git a/ydb/public/lib/ydb_cli/common/query_stats.cpp b/ydb/public/lib/ydb_cli/common/query_stats.cpp
index aade6b323b0..c32d0cd2348 100644
--- a/ydb/public/lib/ydb_cli/common/query_stats.cpp
+++ b/ydb/public/lib/ydb_cli/common/query_stats.cpp
@@ -1,25 +1,21 @@
 #include "query_stats.h"
+#include <iostream>
 
 #include "common.h"
 
 namespace NYdb {
 namespace NConsoleClient {
 
-NTable::ECollectQueryStatsMode ParseQueryStatsMode(const TString& statsMode,
+NTable::ECollectQueryStatsMode ParseQueryStatsModeOrThrow(const TString& statsMode,
     NTable::ECollectQueryStatsMode defaultMode)
 {
-    if (statsMode) {
-        if (statsMode == "none") {
-            return NTable::ECollectQueryStatsMode::None;
-        } else if (statsMode == "basic") {
-            return NTable::ECollectQueryStatsMode::Basic;
-        } else if (statsMode == "full") {
-            return NTable::ECollectQueryStatsMode::Full;
-        } else if (statsMode == "profile") {
-            return NTable::ECollectQueryStatsMode::Profile;
-        } else {
-            throw TMisuseException() << "Unknown stats collection mode " + statsMode + '.';
+    if (!statsMode.empty()) {
+        auto stats = NTable::ParseQueryStatsMode({statsMode.data(), statsMode.size()});
+        if (stats) {
+            return *stats;
         }
+
+        throw TMisuseException() << "Unknown stats collection mode " + statsMode + '.';
     }
 
     return defaultMode;
diff --git a/ydb/public/lib/ydb_cli/common/query_stats.h b/ydb/public/lib/ydb_cli/common/query_stats.h
index f2aa214e18a..a8a556f0c8c 100644
--- a/ydb/public/lib/ydb_cli/common/query_stats.h
+++ b/ydb/public/lib/ydb_cli/common/query_stats.h
@@ -5,7 +5,7 @@
 namespace NYdb {
 namespace NConsoleClient {
 
-NTable::ECollectQueryStatsMode ParseQueryStatsMode(const TString& statsMode,
+NTable::ECollectQueryStatsMode ParseQueryStatsModeOrThrow(const TString& statsMode,
     NTable::ECollectQueryStatsMode defaultMode);
 
 } // namespace NConsoleClient
diff --git a/ydb/public/sdk/cpp/client/ydb_table/query_stats/stats.cpp b/ydb/public/sdk/cpp/client/ydb_table/query_stats/stats.cpp
index 7dd7622062c..f2d37a91ace 100644
--- a/ydb/public/sdk/cpp/client/ydb_table/query_stats/stats.cpp
+++ b/ydb/public/sdk/cpp/client/ydb_table/query_stats/stats.cpp
@@ -1 +1,36 @@
 #include "stats.h"
+
+namespace NYdb {
+namespace NTable {
+
+std::optional<ECollectQueryStatsMode> ParseQueryStatsMode(std::string_view statsMode)
+{
+    if (statsMode == "none") {
+        return ECollectQueryStatsMode::None;
+    } else if (statsMode == "basic") {
+        return ECollectQueryStatsMode::Basic;
+    } else if (statsMode == "full") {
+        return ECollectQueryStatsMode::Full;
+    } else if (statsMode == "profile") {
+        return ECollectQueryStatsMode::Profile;
+    }
+
+    return {};
+}
+
+std::string_view QueryStatsModeToString(ECollectQueryStatsMode statsMode)
+{
+    switch (statsMode) {
+    case ECollectQueryStatsMode::None:
+        return "none";
+    case ECollectQueryStatsMode::Basic:
+        return "basic";
+    case ECollectQueryStatsMode::Full:
+        return "full";
+    case ECollectQueryStatsMode::Profile:
+        return "profile";
+    }
+}
+
+} // namespace NTable
+} // namespace NYdb
diff --git a/ydb/public/sdk/cpp/client/ydb_table/query_stats/stats.h b/ydb/public/sdk/cpp/client/ydb_table/query_stats/stats.h
index 56dc1199b5e..29a7c87fa1f 100644
--- a/ydb/public/sdk/cpp/client/ydb_table/query_stats/stats.h
+++ b/ydb/public/sdk/cpp/client/ydb_table/query_stats/stats.h
@@ -41,5 +41,9 @@ enum class ECollectQueryStatsMode {
 
 using TQueryStats = NQuery::TExecStats;
 
+std::optional<ECollectQueryStatsMode> ParseQueryStatsMode(std::string_view statsMode);
+
+std::string_view QueryStatsModeToString(ECollectQueryStatsMode statsMode);
+
 } // namespace NTable
 } // namespace NYdb