Check secret existence (#14865)

6 files changed, 152 insertions, 20 deletions

diff --git a/ydb/library/backup/backup.cpp b/ydb/library/backup/backup.cpp
index 3a3eeb910d7..bc95c348b9c 100644
--- a/ydb/library/backup/backup.cpp
+++ b/ydb/library/backup/backup.cpp
@@ -808,7 +808,7 @@ Ydb::Table::DescribeExternalDataSourceResult DescribeExternalDataSource(TDriver
 
 std::string ToString(std::string_view key, std::string_view value) {
     // indented to follow the default YQL formatting
-    return std::format(R"(    {} = "{}")", key, value);
+    return std::format(R"(  {} = '{}')", key, value);
 }
 
 namespace NExternalDataSource {
diff --git a/ydb/public/lib/ydb_cli/dump/restore_impl.cpp b/ydb/public/lib/ydb_cli/dump/restore_impl.cpp
index fd87a220aa4..5cda5a1f136 100644
--- a/ydb/public/lib/ydb_cli/dump/restore_impl.cpp
+++ b/ydb/public/lib/ydb_cli/dump/restore_impl.cpp
@@ -20,6 +20,7 @@
 #include <library/cpp/threading/future/core/future.h>
 
 #include <util/generic/deque.h>
+#include <util/generic/guid.h>
 #include <util/generic/hash.h>
 #include <util/generic/hash_set.h>
 #include <util/generic/maybe.h>
@@ -1115,6 +1116,32 @@ TRestoreResult TRestoreClient::RestoreTopic(
     return Result<TRestoreResult>(dbPath, std::move(result));
 }
 
+TRestoreResult TRestoreClient::CheckSecretExistence(const TString& secretName) {
+    LOG_D("Check existence of the secret " << secretName.Quote());
+
+    const auto tmpUser = CreateGuidAsString();
+    const auto create = std::format("CREATE OBJECT `{}:{}` (TYPE SECRET_ACCESS);", secretName.c_str(), tmpUser.c_str());
+    const auto drop = std::format("DROP OBJECT `{}:{}` (TYPE SECRET_ACCESS);", secretName.c_str(), tmpUser.c_str());
+
+    auto result = QueryClient.RetryQuerySync([&](NQuery::TSession session) {
+        return session.ExecuteQuery(create, NQuery::TTxControl::NoTx()).ExtractValueSync();
+    });
+    if (!result.IsSuccess()) {
+        return Result<TRestoreResult>(EStatus::PRECONDITION_FAILED, TStringBuilder()
+            << "Secret " << secretName.Quote() << " does not exist or you do not have access permissions");
+    }
+
+    result = QueryClient.RetryQuerySync([&](NQuery::TSession session) {
+        return session.ExecuteQuery(drop, NQuery::TTxControl::NoTx()).ExtractValueSync();
+    });
+    if (!result.IsSuccess()) {
+        return Result<TRestoreResult>(EStatus::INTERNAL_ERROR, TStringBuilder()
+            << "Failed to drop temporary secret access " << secretName << ":" << tmpUser);
+    }
+
+    return result;
+}
+
 TRestoreResult TRestoreClient::RestoreReplication(
     const TFsPath& fsPath,
     const TString& dbRestoreRoot,
@@ -1136,6 +1163,11 @@ TRestoreResult TRestoreClient::RestoreReplication(
     }
 
     auto query = ReadAsyncReplicationQuery(fsPath, Log.get());
+    if (const auto secretName = GetSecretName(query)) {
+        if (auto result = CheckSecretExistence(secretName); !result.IsSuccess()) {
+            return Result<TRestoreResult>(fsPath.GetPath(), std::move(result));
+        }
+    }
 
     NYql::TIssues issues;
     if (!RewriteObjectRefs(query, dbRestoreRoot, issues)) {
@@ -1267,6 +1299,11 @@ TRestoreResult TRestoreClient::RestoreExternalDataSource(
     }
 
     TString query = ReadExternalDataSourceQuery(fsPath, Log.get());
+    if (const auto secretName = GetSecretName(query)) {
+        if (auto result = CheckSecretExistence(secretName); !result.IsSuccess()) {
+            return Result<TRestoreResult>(fsPath.GetPath(), std::move(result));
+        }
+    }
 
     NYql::TIssues issues;
     if (!RewriteCreateQuery(query, "CREATE EXTERNAL DATA SOURCE IF NOT EXISTS `{}`", dbPath, issues)) {
diff --git a/ydb/public/lib/ydb_cli/dump/restore_impl.h b/ydb/public/lib/ydb_cli/dump/restore_impl.h
index f7f9b233fd6..73c35a39c78 100644
--- a/ydb/public/lib/ydb_cli/dump/restore_impl.h
+++ b/ydb/public/lib/ydb_cli/dump/restore_impl.h
@@ -219,6 +219,8 @@ class TRestoreClient {
         const TString& dbPath, const TRestoreSettings& settings, const NTable::TTableDescription& desc,
         ui32 dataFilesCount);
 
+    TRestoreResult CheckSecretExistence(const TString& secretName);
+
 public:
     explicit TRestoreClient(const TDriver& driver, const std::shared_ptr<TLog>& log);
 
diff --git a/ydb/public/lib/ydb_cli/dump/util/query_utils.cpp b/ydb/public/lib/ydb_cli/dump/util/query_utils.cpp
index 0e61b679426..5acea796702 100644
--- a/ydb/public/lib/ydb_cli/dump/util/query_utils.cpp
+++ b/ydb/public/lib/ydb_cli/dump/util/query_utils.cpp
@@ -190,7 +190,7 @@ struct TTableRefValidator {
     NYql::TIssues& Issues;
 };
 
-TString GetOpt(TStringInput query, TStringBuf pattern) {
+TString GetToken(TStringInput query, TStringBuf pattern) {
     TString line;
     while (query.ReadLine(line)) {
         StripInPlace(line);
@@ -205,11 +205,26 @@ TString GetOpt(TStringInput query, TStringBuf pattern) {
 } // anonymous
 
 TString GetBackupRoot(const TString& query) {
-    return GetOpt(query, R"(-- backup root: ")");
+    return GetToken(query, R"(-- backup root: ")");
 }
 
 TString GetDatabase(const TString& query) {
-    return GetOpt(query, R"(-- database: ")");
+    return GetToken(query, R"(-- database: ")");
+}
+
+TString GetSecretName(const TString& query) {
+    TString secretName;
+    if (auto pwd = GetToken(query, R"(PASSWORD_SECRET_NAME = ')")) {
+        secretName = std::move(pwd);
+    } else if (auto token = GetToken(query, R"(TOKEN_SECRET_NAME = ')")) {
+        secretName = std::move(token);
+    }
+
+    if (secretName.EndsWith("'")) {
+        secretName.resize(secretName.size() - 1);
+    }
+
+    return secretName;
 }
 
 bool SqlToProtoAst(const TString& queryStr, TRule_sql_query& queryProto, NYql::TIssues& issues) {
diff --git a/ydb/public/lib/ydb_cli/dump/util/query_utils.h b/ydb/public/lib/ydb_cli/dump/util/query_utils.h
index b029aa1a018..2caa8bebd77 100644
--- a/ydb/public/lib/ydb_cli/dump/util/query_utils.h
+++ b/ydb/public/lib/ydb_cli/dump/util/query_utils.h
@@ -28,5 +28,6 @@ bool RewriteCreateQuery(TString& query, std::string_view pattern, const std::str
 
 TString GetBackupRoot(const TString& query);
 TString GetDatabase(const TString& query);
+TString GetSecretName(const TString& query);
 
 } // NYdb::NDump
diff --git a/ydb/services/ydb/backup_ut/ydb_backup_ut.cpp b/ydb/services/ydb/backup_ut/ydb_backup_ut.cpp
index 105c76f0c5d..94717619efd 100644
--- a/ydb/services/ydb/backup_ut/ydb_backup_ut.cpp
+++ b/ydb/services/ydb/backup_ut/ydb_backup_ut.cpp
@@ -979,6 +979,20 @@ void TestCoordinationNodeResourcesArePreserved(
     }
 }
 
+void WaitReplicationInit(NReplication::TReplicationClient& client, const TString& path) {
+    int retry = 0;
+    do {
+        auto result = client.DescribeReplication(path).ExtractValueSync();
+        const auto& desc = result.GetReplicationDescription();
+        if (desc.GetItems().empty()) {
+            Sleep(TDuration::Seconds(1));
+        } else {
+            break;
+        }
+    } while (++retry < 10);
+    UNIT_ASSERT(retry < 10);
+}
+
 void TestReplicationSettingsArePreserved(
         const TString& endpoint,
         NQuery::TSession& session,
@@ -997,20 +1011,6 @@ void TestReplicationSettingsArePreserved(
         );)", endpoint.c_str()), true
     );
 
-    auto waitReplicationInit = [&client]() {
-        int retry = 0;
-        do {
-            auto result = client.DescribeReplication("/Root/replication").ExtractValueSync();
-            const auto& desc = result.GetReplicationDescription();
-            if (desc.GetItems().empty()) {
-                Sleep(TDuration::Seconds(1));
-            } else {
-                break;
-            }
-        } while (++retry < 10);
-        UNIT_ASSERT(retry < 10);
-    };
-
     auto checkDescription = [&client, &endpoint]() {
         auto result = client.DescribeReplication("/Root/replication").ExtractValueSync();
         const auto& desc = result.GetReplicationDescription();
@@ -1026,12 +1026,12 @@ void TestReplicationSettingsArePreserved(
         UNIT_ASSERT_VALUES_EQUAL(items.at(0).DstPath, "/Root/replica");
     };
 
-    waitReplicationInit();
+    WaitReplicationInit(client, "/Root/replication");
     checkDescription();
     backup();
     ExecuteQuery(session, "DROP ASYNC REPLICATION `/Root/replication` CASCADE;", true);
     restore();
-    waitReplicationInit();
+    WaitReplicationInit(client, "/Root/replication");
     checkDescription();
 }
 
@@ -1545,6 +1545,44 @@ Y_UNIT_TEST_SUITE(BackupRestore) {
         );
     }
 
+    Y_UNIT_TEST(RestoreReplicationWithoutSecret) {
+        TKikimrWithGrpcAndRootSchema server;
+
+        const auto endpoint = Sprintf("localhost:%u", server.GetPort());
+        auto driver = TDriver(TDriverConfig().SetEndpoint(endpoint).SetAuthToken("root@builtin"));
+
+        NQuery::TQueryClient queryClient(driver);
+        auto session = queryClient.GetSession().ExtractValueSync().GetSession();
+        NReplication::TReplicationClient replicationClient(driver);
+
+        TTempDir tempDir;
+        const auto& pathToBackup = tempDir.Path();
+
+        ExecuteQuery(session, "CREATE OBJECT `secret` (TYPE SECRET) WITH (value = 'root@builtin');", true);
+        ExecuteQuery(session, "CREATE TABLE `/Root/table` (k Uint32, v Utf8, PRIMARY KEY (k));", true);
+        ExecuteQuery(session, Sprintf(R"(
+            CREATE ASYNC REPLICATION `/Root/replication` FOR
+                `/Root/table` AS `/Root/replica`
+            WITH (
+                CONNECTION_STRING = 'grpc://%s/?database=/Root',
+                TOKEN_SECRET_NAME = 'secret'
+            );)", endpoint.c_str()), true
+        );
+        WaitReplicationInit(replicationClient, "/Root/replication");
+
+        NDump::TClient backupClient(driver);
+        {
+            const auto result = backupClient.Dump("/Root", pathToBackup, NDump::TDumpSettings().Database("/Root"));
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+        ExecuteQuery(session, "DROP OBJECT `secret` (TYPE SECRET);", true);
+        ExecuteQuery(session, "DROP ASYNC REPLICATION `/Root/replication` CASCADE;", true);
+        {
+            const auto result = backupClient.Restore(pathToBackup, "/Root");
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+        }
+    }
+
     void TestExternalDataSourceBackupRestore() {
         TKikimrWithGrpcAndRootSchema server;
         server.GetRuntime()->GetAppData().FeatureFlags.SetEnableExternalDataSources(true);
@@ -1567,6 +1605,45 @@ Y_UNIT_TEST_SUITE(BackupRestore) {
         );
     }
 
+    Y_UNIT_TEST(RestoreExternalDataSourceWithoutSecret) {
+        TKikimrWithGrpcAndRootSchema server;
+        server.GetRuntime()->GetAppData().FeatureFlags.SetEnableExternalDataSources(true);
+
+        auto driver = TDriver(TDriverConfig().SetEndpoint(Sprintf("localhost:%u", server.GetPort())));
+        TTableClient tableClient(driver);
+        auto tableSession = tableClient.CreateSession().ExtractValueSync().GetSession();
+
+        NQuery::TQueryClient queryClient(driver);
+        auto querySession = queryClient.GetSession().ExtractValueSync().GetSession();
+
+        TTempDir tempDir;
+        const auto& pathToBackup = tempDir.Path();
+
+        ExecuteQuery(querySession, "CREATE OBJECT `secret` (TYPE SECRET) WITH (value = 'secret');", true);
+        ExecuteQuery(querySession, R"(
+            CREATE EXTERNAL DATA SOURCE `/Root/externalDataSource` WITH (
+                SOURCE_TYPE = "PostgreSQL",
+                DATABASE_NAME = "db",
+                LOCATION = "192.168.1.1:8123",
+                AUTH_METHOD = "BASIC",
+                LOGIN = "user",
+                PASSWORD_SECRET_NAME = "secret"
+            );
+        )", true);
+
+        NDump::TClient backupClient(driver);
+        {
+            const auto result = backupClient.Dump("/Root", pathToBackup, NDump::TDumpSettings().Database("/Root"));
+            UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+        }
+        ExecuteQuery(querySession, "DROP OBJECT `secret` (TYPE SECRET);", true);
+        ExecuteQuery(querySession, "DROP EXTERNAL DATA SOURCE `/Root/externalDataSource`;", true);
+        {
+            const auto result = backupClient.Restore(pathToBackup, "/Root");
+            UNIT_ASSERT_VALUES_EQUAL_C(result.GetStatus(), EStatus::PRECONDITION_FAILED, result.GetIssues().ToString());
+        }
+    }
+
     void TestExternalTableBackupRestore() {
         TKikimrWithGrpcAndRootSchema server;
         server.GetRuntime()->GetAppData().FeatureFlags.SetEnableExternalDataSources(true);