Add YDB CLI commands for update and get state (#20114)

4 files changed, 245 insertions, 0 deletions

diff --git a/ydb/public/lib/ydb_cli/commands/ya.make b/ydb/public/lib/ydb_cli/commands/ya.make
index 4b671334caf..86e60e67e7f 100644
--- a/ydb/public/lib/ydb_cli/commands/ya.make
+++ b/ydb/public/lib/ydb_cli/commands/ya.make
@@ -12,6 +12,7 @@ SRCS(
     topic_readwrite_scenario.cpp
     ydb_admin.cpp
     ydb_benchmark.cpp
+    ydb_bridge.cpp
     ydb_cluster.cpp
     ydb_debug.cpp
     ydb_dynamic_config.cpp
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_bridge.cpp b/ydb/public/lib/ydb_cli/commands/ydb_bridge.cpp
new file mode 100644
index 00000000000..c94c1a4667a
--- /dev/null
+++ b/ydb/public/lib/ydb_cli/commands/ydb_bridge.cpp
@@ -0,0 +1,198 @@
+#include "ydb_bridge.h"
+
+#include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/draft/ydb_bridge.h>
+#include <ydb/public/lib/ydb_cli/common/command_utils.h>
+#include <library/cpp/json/json_reader.h>
+#include <library/cpp/json/json_value.h>
+#include <library/cpp/json/json_writer.h>
+
+#include <ydb/public/lib/ydb_cli/common/pretty_table.h>
+
+#include <util/stream/file.h>
+#include <util/string/cast.h>
+#include <util/string/split.h>
+
+namespace NYdb::NConsoleClient {
+
+namespace {
+    Ydb::Bridge::PileState ParsePileState(TString stateStr) {
+        TString stateStrUpper = stateStr;
+        stateStrUpper.to_upper();
+
+        if (stateStrUpper == "DISCONNECTED") {
+            return Ydb::Bridge::PileState::DISCONNECTED;
+        }
+        if (stateStrUpper == "NOT_SYNCHRONIZED") {
+            return Ydb::Bridge::PileState::NOT_SYNCHRONIZED;
+        }
+        if (stateStrUpper == "SYNCHRONIZED") {
+            return Ydb::Bridge::PileState::SYNCHRONIZED;
+        }
+        if (stateStrUpper == "PROMOTE") {
+            return Ydb::Bridge::PileState::PROMOTE;
+        }
+        if (stateStrUpper == "PRIMARY") {
+            return Ydb::Bridge::PileState::PRIMARY;
+        }
+        ythrow yexception() << "Invalid pile state: \"" << stateStr
+            << "\". Please use one of: DISCONNECTED, NOT_SYNCHRONIZED, SYNCHRONIZED, PROMOTE, PRIMARY.";
+    }
+
+    TString PileStateToString(NYdb::NBridge::EPileState state) {
+        switch (state) {
+            case NYdb::NBridge::EPileState::DISCONNECTED: return "DISCONNECTED";
+            case NYdb::NBridge::EPileState::NOT_SYNCHRONIZED: return "NOT_SYNCHRONIZED";
+            case NYdb::NBridge::EPileState::SYNCHRONIZED: return "SYNCHRONIZED";
+            case NYdb::NBridge::EPileState::PROMOTE: return "PROMOTE";
+            case NYdb::NBridge::EPileState::PRIMARY: return "PRIMARY";
+        }
+        return "UNKNOWN";
+    }
+}
+
+TCommandBridge::TCommandBridge(bool allowEmptyDatabase)
+    : TClientCommandTree("bridge", {}, "Manage cluster in bridge mode")
+{
+    AddCommand(std::make_unique<TCommandBridgeUpdate>(allowEmptyDatabase));
+    AddCommand(std::make_unique<TCommandBridgeGet>(allowEmptyDatabase));
+}
+
+TCommandBridgeUpdate::TCommandBridgeUpdate(bool allowEmptyDatabase)
+    : TYdbCommand("update", {}, "Update cluster state in bridge mode")
+    , AllowEmptyDatabase(allowEmptyDatabase)
+{
+}
+
+void TCommandBridgeUpdate::Config(TConfig& config) {
+    TYdbCommand::Config(config);
+    config.Opts->AddLongOption("set", "Set new state for a pile. Format: <pile-id>:<state>. Can be used multiple times.")
+        .RequiredArgument("ID:STATE").Handler([this](const TString& value) {
+            PileStateUpdates.push_back(value);
+        });
+    config.Opts->AddLongOption('f', "file", "Path to a JSON file with state updates.")
+        .RequiredArgument("PATH").StoreResult(&FilePath);
+    config.Opts->MutuallyExclusive("set", "file");
+
+    config.AllowEmptyDatabase = AllowEmptyDatabase;
+    config.SetFreeArgsNum(0);
+}
+
+void TCommandBridgeUpdate::Parse(TConfig& config) {
+    TYdbCommand::Parse(config);
+
+    if (PileStateUpdates.empty() && FilePath.empty()) {
+        ythrow yexception() << "Either --set or --file must be specified.";
+    }
+
+    if (!PileStateUpdates.empty()) {
+        for (const auto& updateStr : PileStateUpdates) {
+            TStringBuf pileIdStr, stateStr;
+            if (!TStringBuf(updateStr).TrySplit(':', pileIdStr, stateStr) || pileIdStr.empty() || stateStr.empty()) {
+                ythrow yexception() << "Invalid format for --set option. Expected '<pile-id>:<state>'.";
+            }
+            NYdb::NBridge::TPileStateUpdate update;
+            update.PileId = FromString<ui32>(pileIdStr);
+            update.State = static_cast<NYdb::NBridge::EPileState>(ParsePileState(TString(stateStr)));
+            Updates.push_back(update);
+        }
+    }
+
+    if (!FilePath.empty()) {
+        TString jsonStr = TFileInput(FilePath).ReadAll();
+        NJson::TJsonValue jsonValue;
+        if (!NJson::ReadJsonTree(jsonStr, &jsonValue)) {
+            ythrow yexception() << "Failed to parse JSON from file \"" << FilePath << "\"";
+        }
+        if (!jsonValue.IsArray()) {
+            ythrow yexception() << "Root of the JSON document must be an array. File should start with '['. "
+                << "Location: " << FilePath;
+        }
+
+        for (const auto& item : jsonValue.GetArray()) {
+            if (!item.IsMap() || !item.Has("pile_id") || !item.Has("state")) {
+                ythrow yexception() << "Invalid object in JSON array: each item must be an object with string \"state\" and integer \"pile_id\" keys.";
+            }
+            NYdb::NBridge::TPileStateUpdate update;
+            update.PileId = item["pile_id"].GetInteger();
+            update.State = static_cast<NYdb::NBridge::EPileState>(ParsePileState(item["state"].GetString()));
+            Updates.push_back(update);
+        }
+    }
+}
+
+int TCommandBridgeUpdate::Run(TConfig& config) {
+    auto driver = std::make_unique<TDriver>(CreateDriver(config));
+    auto client = NYdb::NBridge::TBridgeClient(*driver);
+
+    auto result = client.UpdateClusterState(Updates).GetValueSync();
+    NStatusHelpers::ThrowOnErrorOrPrintIssues(result);
+
+    Cout << "Cluster state updated successfully." << Endl;
+
+    return EXIT_SUCCESS;
+}
+
+TCommandBridgeGet::TCommandBridgeGet(bool allowEmptyDatabase)
+    : TYdbReadOnlyCommand("get", {"state"}, "Get current bridge cluster state")
+    , AllowEmptyDatabase(allowEmptyDatabase)
+{}
+
+void TCommandBridgeGet::Config(TConfig& config) {
+    TYdbReadOnlyCommand::Config(config);
+    config.AllowEmptyDatabase = AllowEmptyDatabase;
+    config.SetFreeArgsNum(0);
+    AddOutputFormats(config, {
+        EDataFormat::Pretty,
+        EDataFormat::Json,
+        EDataFormat::Csv
+    });
+}
+
+void TCommandBridgeGet::Parse(TConfig& config) {
+    TClientCommand::Parse(config);
+    ParseOutputFormats();
+}
+
+int TCommandBridgeGet::Run(TConfig& config) {
+    auto driver = std::make_unique<TDriver>(CreateDriver(config));
+    auto client = NYdb::NBridge::TBridgeClient(*driver);
+
+    auto result = client.GetClusterState().GetValueSync();
+    NStatusHelpers::ThrowOnErrorOrPrintIssues(result);
+
+    const auto& state = result.GetState();
+
+    switch (OutputFormat) {
+        case EDataFormat::Json: {
+            NJson::TJsonValue json(NJson::JSON_ARRAY);
+            for (const auto& s : state) {
+                NJson::TJsonValue item(NJson::JSON_MAP);
+                item.InsertValue("pile_id", s.PileId);
+                item.InsertValue("state", PileStateToString(s.State));
+                json.AppendValue(item);
+            }
+            NJson::WriteJson(&Cout, &json, true);
+            Cout << Endl;
+            break;
+        }
+        case EDataFormat::Csv: {
+            Cout << "pile_id,state" << Endl;
+            for (const auto& s : state) {
+                Cout << s.PileId << "," << PileStateToString(s.State) << Endl;
+            }
+            break;
+        }
+        default: {
+            TStringStream ss;
+            for (const auto& s : state) {
+                ss << "Pile " << s.PileId << ": " << PileStateToString(s.State) << Endl;
+            }
+            Cout << ss.Str();
+            break;
+        }
+    }
+
+    return EXIT_SUCCESS;
+}
+
+}
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_bridge.h b/ydb/public/lib/ydb_cli/commands/ydb_bridge.h
new file mode 100644
index 00000000000..53784e089a9
--- /dev/null
+++ b/ydb/public/lib/ydb_cli/commands/ydb_bridge.h
@@ -0,0 +1,44 @@
+#pragma once
+
+#include "ydb_command.h"
+#include "ydb_common.h"
+#include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/draft/ydb_bridge.h>
+#include <ydb/public/lib/ydb_cli/common/format.h>
+
+#include <util/generic/string.h>
+#include <vector>
+
+namespace NYdb::NConsoleClient {
+
+class TCommandBridge : public TClientCommandTree {
+public:
+    TCommandBridge(bool allowEmptyDatabase);
+};
+
+class TCommandBridgeUpdate : public TYdbCommand {
+public:
+    TCommandBridgeUpdate(bool allowEmptyDatabase);
+    void Config(TConfig& config) override;
+    void Parse(TConfig& config) override;
+    int Run(TConfig& config) override;
+
+private:
+    bool AllowEmptyDatabase = false;
+
+    TVector<TString> PileStateUpdates;
+    TString FilePath;
+
+    std::vector<NYdb::NBridge::TPileStateUpdate> Updates;
+};
+
+class TCommandBridgeGet : public TYdbReadOnlyCommand, public TCommandWithOutput {
+public:
+    TCommandBridgeGet(bool allowEmptyDatabase);
+    void Config(TConfig& config) override;
+    void Parse(TConfig& config) override;
+    int Run(TConfig& config) override;
+private:
+    bool AllowEmptyDatabase;
+};
+
+}
diff --git a/ydb/public/lib/ydb_cli/commands/ydb_cluster.cpp b/ydb/public/lib/ydb_cli/commands/ydb_cluster.cpp
index 4d203d37170..651101851fd 100644
--- a/ydb/public/lib/ydb_cli/commands/ydb_cluster.cpp
+++ b/ydb/public/lib/ydb_cli/commands/ydb_cluster.cpp
@@ -1,5 +1,6 @@
 #include "ydb_cluster.h"
 
+#include "ydb_bridge.h"
 #include "ydb_dynamic_config.h"
 
 #include <ydb/public/sdk/cpp/include/ydb-cpp-sdk/client/config/config.h>
@@ -20,6 +21,7 @@ TCommandCluster::TCommandCluster()
     AddCommand(std::make_unique<NDynamicConfig::TCommandConfig>(false, true));
     AddCommand(std::make_unique<TCommandClusterDump>());
     AddCommand(std::make_unique<TCommandClusterRestore>());
+    AddHiddenCommand(std::make_unique<TCommandBridge>(true));
 }
 
 TCommandClusterBootstrap::TCommandClusterBootstrap()