Restore async replication (#14465)

Author: CyberROFL

ydb/core/tx/schemeshard/schemeshard_import__create.cpp

@@ -848,7 +848,8 @@ struct TSchemeShard::TImport::TTxProgress: public TSchemeShard::TXxport::TTxBase
                 << importInfo->Settings.items(msg.ItemIdx).source_prefix() << NYdb::NDump::NFiles::CreateView().FileName;
 
             NYql::TIssues issues;
-            if (!NYdb::NDump::RewriteCreateViewQuery(item.CreationQuery, database, true, item.DstPathName, source, issues)) {
+            if (!NYdb::NDump::RewriteCreateViewQuery(item.CreationQuery, database, true, item.DstPathName, issues)) {
+                issues.AddIssue(TStringBuilder() << "path: " << source);
                 return CancelAndPersist(db, importInfo, msg.ItemIdx, issues.ToString(), "invalid view creation query");
             }
             item.SchemeQueryExecutor = ctx.Register(CreateSchemeQueryExecutor(

ydb/public/lib/ydb_cli/dump/restore_impl.cpp

@@ -9,6 +9,7 @@
 #include <ydb/public/lib/ydb_cli/common/retry_func.h>
 #include <ydb/public/lib/ydb_cli/dump/files/files.h>
 #include <ydb/public/lib/ydb_cli/dump/util/log.h>
+#include <ydb/public/lib/ydb_cli/dump/util/rewrite_query.h>
 #include <ydb/public/lib/ydb_cli/dump/util/util.h>
 #include <ydb/public/lib/ydb_cli/dump/util/view_utils.h>
 #include <ydb-cpp-sdk/client/proto/accessor.h>
@@ -48,12 +49,24 @@ bool IsFileExists(const TFsPath& path) {
     return path.Exists() && path.IsFile();
 }
 
-template <typename TProtoType>
-TProtoType ReadProtoFromFile(const TFsPath& fsDirPath, const TLog* log, const NFiles::TFileInfo& fileInfo) {
+TString ReadFromFile(const TFsPath& fsDirPath, const TLog* log, const NFiles::TFileInfo& fileInfo) {
     const auto fsPath = fsDirPath.Child(fileInfo.FileName);
     LOG_IMPL(log, ELogPriority::TLOG_DEBUG, "Read " << fileInfo.LogObjectType << " from " << fsPath.GetPath().Quote());
+    return TFileInput(fsPath).ReadAll();
+}
+
+TString ReadViewQuery(const TFsPath& fsDirPath, const TLog* log) {
+    return ReadFromFile(fsDirPath, log, NFiles::CreateView());
+}
+
+TString ReadAsyncReplicationQuery(const TFsPath& fsDirPath, const TLog* log) {
+    return ReadFromFile(fsDirPath, log, NFiles::CreateAsyncReplication());
+}
+
+template <typename TProtoType>
+TProtoType ReadProtoFromFile(const TFsPath& fsDirPath, const TLog* log, const NFiles::TFileInfo& fileInfo) {
     TProtoType proto;
-    Y_ENSURE(google::protobuf::TextFormat::ParseFromString(TFileInput(fsPath).ReadAll(), &proto));
+    Y_ENSURE(google::protobuf::TextFormat::ParseFromString(ReadFromFile(fsDirPath, log, fileInfo), &proto));
     return proto;
 }
 
@@ -465,6 +478,10 @@ TRestoreResult TRestoreClient::RestoreFolder(
         return RestoreCoordinationNode(fsPath, objectDbPath, settings, oldEntries.contains(objectDbPath));
     }
 
+    if (IsFileExists(fsPath.Child(NFiles::CreateAsyncReplication().FileName))) {
+        return RestoreReplication(fsPath, objectDbPath, settings, oldEntries.contains(objectDbPath));
+    }
+
     if (IsFileExists(fsPath.Child(NFiles::Empty().FileName))) {
         return RestoreEmptyDir(fsPath, objectDbPath, settings, oldEntries.contains(objectDbPath));
     }
@@ -486,6 +503,8 @@ TRestoreResult TRestoreClient::RestoreFolder(
             result = RestoreTopic(child, childDbPath, settings, oldEntries.contains(childDbPath));
         } else if (IsFileExists(child.Child(NFiles::CreateCoordinationNode().FileName))) {
             result = RestoreCoordinationNode(child, childDbPath, settings, oldEntries.contains(childDbPath));
+        } else if (IsFileExists(child.Child(NFiles::CreateAsyncReplication().FileName))) {
+            result = RestoreReplication(child, childDbPath, settings, oldEntries.contains(childDbPath));
         } else if (child.IsDirectory()) {
             result = RestoreFolder(child, dbRestoreRoot, Join('/', dbPathRelativeToRestoreRoot, child.GetName()), settings, oldEntries);
         }
@@ -520,13 +539,12 @@ TRestoreResult TRestoreClient::RestoreView(
     const TString dbPath = dbRestoreRoot + dbPathRelativeToRestoreRoot;
     LOG_I("Restore view " << fsPath.GetPath().Quote() << " to " << dbPath.Quote());
 
-    const auto createViewFile = fsPath.Child(NFiles::CreateView().FileName);
-    TString query = TFileInput(createViewFile).ReadAll();
+    TString query = ReadViewQuery(fsPath, Log.get());
 
     NYql::TIssues issues;
     const bool isDb = IsDatabase(SchemeClient, dbRestoreRoot);
-    if (!RewriteCreateViewQuery(query, dbRestoreRoot, isDb, dbPath, createViewFile.GetPath().Quote(), issues)) {
-        return Result<TRestoreResult>(dbPath, EStatus::BAD_REQUEST, issues.ToString());
+    if (!RewriteCreateViewQuery(query, dbRestoreRoot, isDb, dbPath, issues)) {
+        return Result<TRestoreResult>(fsPath.GetPath(), EStatus::BAD_REQUEST, issues.ToString());
     }
 
     if (settings.DryRun_) {
@@ -584,6 +602,44 @@ TRestoreResult TRestoreClient::RestoreTopic(
     return Result<TRestoreResult>(dbPath, std::move(result));
 }
 
+TRestoreResult TRestoreClient::RestoreReplication(
+    const TFsPath& fsPath,
+    const TString& dbPath,
+    const TRestoreSettings& settings,
+    bool isAlreadyExisting)
+{
+    LOG_D("Process " << fsPath.GetPath().Quote());
+
+    if (auto error = ErrorOnIncomplete(fsPath)) {
+        return *error;
+    }
+
+    LOG_I("Restore async replication " << fsPath.GetPath().Quote() << " to " << dbPath.Quote());
+
+    if (settings.DryRun_) {
+        return CheckExistenceAndType(SchemeClient, dbPath, ESchemeEntryType::Replication);
+    }
+
+    auto query = ReadAsyncReplicationQuery(fsPath, Log.get());
+
+    NYql::TIssues issues;
+    if (!RewriteCreateQuery(query, "CREATE ASYNC REPLICATION `{}`", dbPath, issues)) {
+        return Result<TRestoreResult>(fsPath.GetPath(), EStatus::BAD_REQUEST, issues.ToString());
+    }
+
+    auto result = QueryClient.RetryQuerySync([&](NQuery::TSession session) {
+        return session.ExecuteQuery(query, NQuery::TTxControl::NoTx()).ExtractValueSync();
+    });
+
+    if (result.IsSuccess()) {
+        LOG_D("Created " << dbPath.Quote());
+        return RestorePermissions(fsPath, dbPath, settings, isAlreadyExisting);
+    }
+
+    LOG_E("Failed to create " << dbPath.Quote());
+    return Result<TRestoreResult>(dbPath, std::move(result));
+}
+
 TRestoreResult TRestoreClient::RestoreRateLimiter(
     const TFsPath& fsPath,
     const TString& coordinationNodePath,

ydb/public/lib/ydb_cli/dump/restore_impl.h

@@ -131,6 +131,7 @@ class TRestoreClient {
     TRestoreResult RestoreTable(const TFsPath& fsPath, const TString& dbPath, const TRestoreSettings& settings, bool isAlreadyExisting);
     TRestoreResult RestoreView(const TFsPath& fsPath, const TString& dbRestoreRoot, const TString& dbPathRelativeToRestoreRoot, const TRestoreSettings& settings, bool isAlreadyExisting);
     TRestoreResult RestoreTopic(const TFsPath& fsPath, const TString& dbPath, const TRestoreSettings& settings, bool isAlreadyExisting);
+    TRestoreResult RestoreReplication(const TFsPath& fsPath, const TString& dbPath, const TRestoreSettings& settings, bool isAlreadyExisting);
     TRestoreResult RestoreCoordinationNode(const TFsPath& fsPath, const TString& dbPath, const TRestoreSettings& settings, bool isAlreadyExisting);
     TRestoreResult RestoreDependentResources(const TFsPath& fsPath, const TString& dbPath);
     TRestoreResult RestoreRateLimiter(const TFsPath& fsPath, const TString& coordinationNodePath, const TString& resourcePath);

ydb/public/lib/ydb_cli/dump/util/rewrite_query.cpp

@@ -0,0 +1,21 @@
+#include "rewrite_query.h"
+
+#include <util/string/builder.h>
+
+#include <re2/re2.h>
+
+#include <format>
+
+namespace NYdb::NDump {
+
+bool RewriteCreateQuery(TString& query, std::string_view pattern, const std::string& dbPath, NYql::TIssues& issues) {
+    const auto searchPattern = std::vformat(pattern, std::make_format_args("\\S+"));
+    if (re2::RE2::Replace(&query, searchPattern, std::vformat(pattern, std::make_format_args(dbPath)))) {
+        return true;
+    }
+
+    issues.AddIssue(TStringBuilder() << "Pattern: \"" << pattern << "\" was not found: " << query.Quote());
+    return false;
+}
+
+}

ydb/public/lib/ydb_cli/dump/util/rewrite_query.h

@@ -0,0 +1,7 @@
+#include <yql/essentials/public/issue/yql_issue.h>
+
+namespace NYdb::NDump {
+
+bool RewriteCreateQuery(TString& query, std::string_view pattern, const std::string& dbPath, NYql::TIssues& issues);
+
+}

ydb/public/lib/ydb_cli/dump/util/view_utils.cpp

@@ -1,3 +1,4 @@
+#include "rewrite_query.h"
 #include "view_utils.h"
 
 #include <yql/essentials/parser/proto_ast/gen/v1/SQLv1Lexer.h>
@@ -361,7 +362,7 @@ TString BuildCreateViewQuery(
 }
 
 bool RewriteCreateViewQuery(TString& query, const TString& restoreRoot, bool restoreRootIsDatabase,
-    const TString& dbPath, const TString& source, NYql::TIssues& issues
+    const TString& dbPath, NYql::TIssues& issues
 ) {
     const auto backupRoot = GetBackupRoot(query);
 
@@ -373,16 +374,7 @@ bool RewriteCreateViewQuery(TString& query, const TString& restoreRoot, bool res
         return false;
     }
 
-    constexpr TStringBuf pattern = R"(CREATE VIEW IF NOT EXISTS `\S+` )";
-    if (!re2::RE2::Replace(&query, pattern, std::format(R"(CREATE VIEW IF NOT EXISTS `{}` )", dbPath.c_str()))) {
-        issues.AddIssue(TStringBuilder()
-            << "Cannot restore a view from the source: " << source
-            << ". Pattern: \"" << pattern << "\", was not found in the create view statement: " << query.Quote()
-        );
-        return false;
-    }
-
-    return true;
+    return RewriteCreateQuery(query, "CREATE VIEW IF NOT EXISTS `{}`", dbPath, issues);
 }
 
 }

ydb/public/lib/ydb_cli/dump/util/view_utils.h

@@ -8,7 +8,7 @@ TString BuildCreateViewQuery(
 );
 
 bool RewriteCreateViewQuery(TString& query, const TString& restoreRoot, bool restoreRootIsDatabase,
-    const TString& dbPath, const TString& source, NYql::TIssues& issues
+    const TString& dbPath, NYql::TIssues& issues
 );
 
 }

ydb/public/lib/ydb_cli/dump/util/ya.make

@@ -1,6 +1,7 @@
 LIBRARY()
 
 SRCS(
+    rewrite_query.cpp
     util.cpp
     view_utils.cpp
 )