From c22699e502a4580601ad183dc161eada1ccfad7f Mon Sep 17 00:00:00 2001
From: Pisarenko Grigoriy <grigoriypisar@ydb.tech>
Date: Wed, 12 Feb 2025 16:20:12 +0500
Subject: [PATCH] YQ-4101 KqpRun fixed storage and tenants creation (#14472)

---
 ydb/core/testlib/test_client.cpp                  |  4 ++--
 ydb/core/testlib/test_client.h                    |  2 +-
 ydb/tests/tools/kqprun/src/common.h               |  1 +
 .../tools/kqprun/src/proto/storage_meta.proto     |  1 +
 ydb/tests/tools/kqprun/src/ydb_setup.cpp          | 15 ++++++++++++---
 5 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/ydb/core/testlib/test_client.cpp b/ydb/core/testlib/test_client.cpp
index b9a1d6f934a7..b1b5ec913880 100644
--- a/ydb/core/testlib/test_client.cpp
+++ b/ydb/core/testlib/test_client.cpp
@@ -3219,7 +3219,7 @@ namespace Tests {
         return Server->DynamicNodes();
     }
 
-    void TTenants::CreateTenant(Ydb::Cms::CreateDatabaseRequest request, ui32 nodes, TDuration timeout) {
+    void TTenants::CreateTenant(Ydb::Cms::CreateDatabaseRequest request, ui32 nodes, TDuration timeout, bool acceptAlreadyExist) {
         const TString path = request.path();
         const bool serverless = request.has_serverless_resources();
 
@@ -3229,7 +3229,7 @@ namespace Tests {
             std::move(request), "", "", runtime.GetActorSystem(0), true
         ).ExtractValueSync();
 
-        if (result.operation().status() != Ydb::StatusIds::SUCCESS) {
+        if (result.operation().status() != Ydb::StatusIds::SUCCESS && (!acceptAlreadyExist || result.operation().status() != Ydb::StatusIds::ALREADY_EXISTS)) {
             NYql::TIssues issues;
             NYql::IssuesFromMessage(result.operation().issues(), issues);
             ythrow yexception() << "Failed to create tenant " << path << ", " << result.operation().status() << ", reason:\n" << issues.ToString();
diff --git a/ydb/core/testlib/test_client.h b/ydb/core/testlib/test_client.h
index 24de2596bd54..13d320e1228b 100644
--- a/ydb/core/testlib/test_client.h
+++ b/ydb/core/testlib/test_client.h
@@ -709,7 +709,7 @@ namespace Tests {
         ui32 Availabe() const;
         ui32 Capacity() const;
 
-        void CreateTenant(Ydb::Cms::CreateDatabaseRequest request, ui32 nodes = 1, TDuration timeout = TDuration::Seconds(30));
+        void CreateTenant(Ydb::Cms::CreateDatabaseRequest request, ui32 nodes = 1, TDuration timeout = TDuration::Seconds(30), bool acceptAlreadyExist = false);
 
     private:
         TVector<ui32>& Nodes(const TString &name);
diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h
index c1b3645a2558..f401b8d8c864 100644
--- a/ydb/tests/tools/kqprun/src/common.h
+++ b/ydb/tests/tools/kqprun/src/common.h
@@ -20,6 +20,7 @@ namespace NKqpRun {
 
 constexpr char YQL_TOKEN_VARIABLE[] = "YQL_TOKEN";
 constexpr ui64 DEFAULT_STORAGE_SIZE = 32_GB;
+constexpr TDuration TENANT_CREATION_TIMEOUT = TDuration::Seconds(30);
 
 struct TAsyncQueriesSettings {
     enum class EVerbose {
diff --git a/ydb/tests/tools/kqprun/src/proto/storage_meta.proto b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto
index ee4a7c4162a2..2436f37d85d0 100644
--- a/ydb/tests/tools/kqprun/src/proto/storage_meta.proto
+++ b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto
@@ -13,6 +13,7 @@ message TStorageMeta {
         EType Type = 1;
         uint32 NodesCount = 2;
         string SharedTenant = 3;  // Only for serverless tenants
+        bool CreationInProgress = 4;
     }
 
     uint64 StorageGeneration = 1;
diff --git a/ydb/tests/tools/kqprun/src/ydb_setup.cpp b/ydb/tests/tools/kqprun/src/ydb_setup.cpp
index 07a33a266a8c..d70b727c7805 100644
--- a/ydb/tests/tools/kqprun/src/ydb_setup.cpp
+++ b/ydb/tests/tools/kqprun/src/ydb_setup.cpp
@@ -192,7 +192,6 @@ class TYdbSetup::TImpl {
                 if (!google::protobuf::TextFormat::ParseFromString(TFileInput(StorageMetaPath_.GetPath()).ReadAll(), &StorageMeta_)) {
                     ythrow yexception() << "Storage meta is corrupted, please use --format-storage";
                 }
-                StorageMeta_.SetStorageGeneration(StorageMeta_.GetStorageGeneration() + 1);
                 formatDisk = false;
             }
 
@@ -289,11 +288,17 @@ class TYdbSetup::TImpl {
     void CreateTenant(Ydb::Cms::CreateDatabaseRequest&& request, const TString& relativePath, const TString& type, TStorageMeta::TTenant tenantInfo) {
         const auto absolutePath = request.path();
         const auto [it, inserted] = StorageMeta_.MutableTenants()->emplace(relativePath, tenantInfo);
-        if (inserted) {
+        if (inserted || it->second.GetCreationInProgress()) {
             if (Settings_.VerboseLevel >= EVerbose::Info) {
                 Cout << CoutColors_.Yellow() << TInstant::Now().ToIsoStringLocal() << " Creating " << type << " tenant " << absolutePath << "..." << CoutColors_.Default() << Endl;
             }
-            Tenants_->CreateTenant(std::move(request), tenantInfo.GetNodesCount());
+
+            it->second.SetCreationInProgress(true);
+            UpdateStorageMeta();
+
+            Tenants_->CreateTenant(std::move(request), tenantInfo.GetNodesCount(), TENANT_CREATION_TIMEOUT, true);
+
+            it->second.SetCreationInProgress(false);
             UpdateStorageMeta();
         } else {
             if (it->second.GetType() != tenantInfo.GetType()) {
@@ -378,6 +383,10 @@ class TYdbSetup::TImpl {
         NKikimr::Tests::TServerSettings serverSettings = GetServerSettings(grpcPort);
 
         Server_ = MakeIntrusive<NKikimr::Tests::TServer>(serverSettings);
+
+        StorageMeta_.SetStorageGeneration(StorageMeta_.GetStorageGeneration() + 1);
+        UpdateStorageMeta();
+
         Server_->GetRuntime()->SetDispatchTimeout(TDuration::Max());
 
         if (Settings_.GrpcEnabled) {