Add 'export' feature. (#18525)

Author: Tony-Romanov

ydb/apps/etcd_proxy/proxy.cpp

@@ -108,17 +108,23 @@ int TProxy::Run() {
     if (const auto res = Init()) {
         return res;
     }
-    if (Initialize_) {
-        if (const auto res = InitDatabase()) {
+    if (!ExportTo_.empty()) {
+        if (const auto res = ExportDatabase()) {
             return res;
         }
-    }
-    if (!ImportPrefix_.empty()) {
-        if (const auto res = ImportDatabase()) {
-            return res;
+    } else {
+        if (Initialize_) {
+            if (const auto res = InitDatabase()) {
+                return res;
+            }
+        }
+        if (!ImportPrefix_.empty()) {
+            if (const auto res = ImportDatabase()) {
+                return res;
+            }
         }
     }
-    if (!Initialize_ && ImportPrefix_.empty()) {
+    if (!Initialize_ && ImportPrefix_.empty() && ExportTo_.empty()) {
         if (const auto res = StartServer()) {
             return res;
         }
@@ -225,6 +231,75 @@ int TProxy::ImportDatabase() {
     return 0;
 }
 
+int TProxy::ExportDatabase() {
+    auto credentials = grpc::InsecureChannelCredentials();
+    if (!Root.empty() || !Cert.empty() || !Key.empty()) {
+        const grpc::SslCredentialsOptions opts {
+            .pem_root_certs = TFileInput(Root).ReadAll(),
+            .pem_private_key = TFileInput(Key).ReadAll(),
+            .pem_cert_chain = TFileInput(Cert).ReadAll()
+        };
+        credentials = grpc::SslCredentials(opts);
+    }
+
+    const auto channel = grpc::CreateChannel(TString(ExportTo_), credentials);
+    const std::unique_ptr<etcdserverpb::KV::Stub> kv = etcdserverpb::KV::NewStub(channel);
+
+    NYdb::TDriverConfig config;
+    config.SetEndpoint(Endpoint);
+    config.SetDatabase(Database);
+    if (!Token.empty())
+        config.SetAuthToken(Token);
+    if (!CA.empty())
+        config.UseSecureConnection(TFileInput(CA).ReadAll());
+
+    const auto driver = NYdb::TDriver(config);
+    auto client = NYdb::NTable::TTableClient(driver);
+    auto count = 0ULL;
+    if (const auto res = client.CreateSession().ExtractValueSync(); res.IsSuccess()) {
+        NYdb::NTable::TReadTableSettings settings;
+        settings.BatchLimitRows(97ULL).AppendColumns("key").AppendColumns("value").AppendColumns("lease");
+        if (auto it = res.GetSession().ReadTable(Database + Folder + "/current", settings).ExtractValueSync(); it.IsSuccess()) {
+            for (;;) {
+                auto streamPart = it.ReadNext().ExtractValueSync();
+                if (streamPart.EOS())
+                    break;
+
+                if (!streamPart.IsSuccess()) {
+                    std::cout << streamPart.GetIssues().ToString() << std::endl;
+                    return 1;
+                }
+
+                etcdserverpb::TxnRequest txnRequest;
+                for (auto parser = NYdb::TResultSetParser(streamPart.ExtractPart()); parser.TryNextRow();) {
+                    if (const auto lease = NYdb::TValueParser(parser.GetValue("lease")).GetOptionalInt64(); lease && !*lease) {
+                        ++count;
+                        const auto put = txnRequest.add_success()->mutable_request_put();
+                        put->set_key(*NYdb::TValueParser(parser.GetValue("key")).GetOptionalString());
+                        put->set_value(*NYdb::TValueParser(parser.GetValue("value")).GetOptionalString());
+                    }
+                }
+
+                grpc::ClientContext txnCtx;
+                etcdserverpb::TxnResponse txnResponse;
+                if (const auto& status = kv->Txn(&txnCtx, txnRequest, &txnResponse); !status.ok()) {
+                    std::cout << status.error_message() << std::endl;
+                    return 1;
+                }
+            }
+        } else {
+            std::cout << it.GetIssues().ToString() << std::endl;
+            return 1;
+        }
+    } else {
+        std::cout << res.GetIssues().ToString() << std::endl;
+        return 1;
+    }
+
+    std::cout << count << " keys exported successfully." << std::endl;
+    return 0;
+}
+
 int TProxy::Shutdown() {
     if (GRpcServer)
         GRpcServer->Stop();
@@ -252,6 +327,7 @@ TProxy::TProxy(int argc, char** argv)
 
     opts.AddLongOption("import-from", "Import existing data from etcd base").RequiredArgument("ENDPOINT").DefaultValue("localhost:2379").StoreResult(&ImportFrom_);
     opts.AddLongOption("import-prefix", "Prefix of data to import").RequiredArgument("PREFIX").StoreResult(&ImportPrefix_);
+    opts.AddLongOption("export-to", "Export existing data to etcd from ydb").RequiredArgument("ENDPOINT").StoreResult(&ExportTo_);
 
     opts.AddLongOption("ca", "SSL CA certificate file").Optional().RequiredArgument("CA").StoreResult(&Root);
     opts.AddLongOption("cert", "SSL certificate file").Optional().RequiredArgument("CERT").StoreResult(&Cert);

ydb/apps/etcd_proxy/proxy.h

@@ -33,6 +33,7 @@ class TProxy {
     int Discovery();
     int InitDatabase();
     int ImportDatabase();
+    int ExportDatabase();
     int StartServer();
 
     static THolder<NActors::TActorSystemSetup> BuildActorSystemSetup();
@@ -48,7 +49,7 @@ class TProxy {
     std::string Database, Endpoint, Token, CA, Folder;
     ui16 ListeningPort = 2379U;
     std::string Root, Cert, Key;
-    std::string ImportFrom_, ImportPrefix_;
+    std::string ImportFrom_, ImportPrefix_, ExportTo_;
 };
 
 }