YQ-4092 KqpRun improve tenants flags (ydb-platform#14279)

pixcc · Feb 7, 2025 · d93a17b · d93a17b
1 parent 3d05c5c
commit d93a17b
Show file tree

Hide file tree

Showing 4 changed files with 192 additions and 74 deletions.
diff --git a/ydb/tests/tools/kqprun/kqprun.cpp b/ydb/tests/tools/kqprun/kqprun.cpp
@@ -95,7 +95,7 @@ struct TExecutionOptions {
             .TraceId = DefaultTraceId,
             .PoolId = "",
             .UserSID = BUILTIN_ACL_ROOT,
-            .Database = "",
+            .Database = GetValue(0, Databases, TString()),
             .Timeout = TDuration::Zero()
         };
     }
@@ -135,15 +135,15 @@ struct TExecutionOptions {
 
 private:
     void ValidateOptionsSizes(const TRunnerOptions& runnerOptions) const {
-        const auto checker = [numberQueries = ScriptQueries.size()](size_t checkSize, const TString& optionName) {
-            if (checkSize > numberQueries) {
+        const auto checker = [numberQueries = ScriptQueries.size()](size_t checkSize, const TString& optionName, bool useInSchemeQuery = false) {
+            if (checkSize > std::max(numberQueries, static_cast<size_t>(useInSchemeQuery ? 1 : 0))) {
                 ythrow yexception() << "Too many " << optionName << ". Specified " << checkSize << ", when number of script queries is " << numberQueries;
             }
         };
 
         checker(ExecutionCases.size(), "execution cases");
         checker(ScriptQueryActions.size(), "script query actions");
-        checker(Databases.size(), "databases");
+        checker(Databases.size(), "databases", true);
         checker(TraceIds.size(), "trace ids");
         checker(PoolIds.size(), "pool ids");
         checker(UserSIDs.size(), "user SIDs");
@@ -257,7 +257,7 @@ struct TExecutionOptions {
 
     static void ValidateStorageSettings(const TYdbSetupSettings& ydbSettings) {
         if (ydbSettings.DisableDiskMock) {
-            if (ydbSettings.NodeCount + ydbSettings.SharedTenants.size() + ydbSettings.DedicatedTenants.size() > 1) {
+            if (ydbSettings.NodeCount + ydbSettings.Tenants.size() > 1) {
                 ythrow yexception() << "Disable disk mock cannot be used for multi node clusters (already disabled)";
             } else if (ydbSettings.PDisksPath) {
                 ythrow yexception() << "Disable disk mock cannot be used with real PDisks (already disabled)";
@@ -876,17 +876,50 @@ class TMain : public TMainClassArgs {
             .DefaultValue(RunnerOptions.YdbSettings.DomainName)
             .StoreResult(&RunnerOptions.YdbSettings.DomainName);
 
-        options.AddLongOption("dedicated", "Dedicated tenant path, relative inside domain")
+        const auto addTenant = [this](const TString& type, TStorageMeta::TTenant::EType protoType, const NLastGetopt::TOptsParser* option) {
+            TStringBuf tenant;
+            TStringBuf nodesCountStr;
+            TStringBuf(option->CurVal()).Split(':', tenant, nodesCountStr);
+            if (tenant.empty()) {
+                ythrow yexception() << type << " tenant name should not be empty";
+            }
+
+            TStorageMeta::TTenant tenantInfo;
+            tenantInfo.SetType(protoType);
+            tenantInfo.SetNodesCount(nodesCountStr ? FromString<ui32>(nodesCountStr) : 1);
+            if (tenantInfo.GetNodesCount() == 0) {
+                ythrow yexception() << type << " tenant should have at least one node";
+            }
+
+            if (!RunnerOptions.YdbSettings.Tenants.emplace(tenant, tenantInfo).second) {
+                ythrow yexception() << "Got duplicated tenant name: " << tenant;
+            }
+        };
+        options.AddLongOption("dedicated", "Dedicated tenant path, relative inside domain (for node count use dedicated-name:node-count)")
             .RequiredArgument("path")
-            .InsertTo(&RunnerOptions.YdbSettings.DedicatedTenants);
+            .Handler1(std::bind(addTenant, "Dedicated", TStorageMeta::TTenant::DEDICATED, std::placeholders::_1));
 
-        options.AddLongOption("shared", "Shared tenant path, relative inside domain")
+        options.AddLongOption("shared", "Shared tenant path, relative inside domain (for node count use dedicated-name:node-count)")
             .RequiredArgument("path")
-            .InsertTo(&RunnerOptions.YdbSettings.SharedTenants);
+            .Handler1(std::bind(addTenant, "Shared", TStorageMeta::TTenant::SHARED, std::placeholders::_1));
 
         options.AddLongOption("serverless", "Serverless tenant path, relative inside domain (use string serverless-name@shared-name to specify shared database)")
             .RequiredArgument("path")
-            .InsertTo(&RunnerOptions.YdbSettings.ServerlessTenants);
+            .Handler1([this](const NLastGetopt::TOptsParser* option) {
+                TStringBuf serverless;
+                TStringBuf shared;
+                TStringBuf(option->CurVal()).Split('@', serverless, shared);
+                if (serverless.empty()) {
+                    ythrow yexception() << "Serverless tenant name should not be empty";
+                }
+
+                TStorageMeta::TTenant tenantInfo;
+                tenantInfo.SetType(TStorageMeta::TTenant::SERVERLESS);
+                tenantInfo.SetSharedTenant(TString(shared));
+                if (!RunnerOptions.YdbSettings.Tenants.emplace(serverless, tenantInfo).second) {
+                    ythrow yexception() << "Got duplicated tenant name: " << serverless;
+                }
+            });
 
         options.AddLongOption("storage-size", TStringBuilder() << "Domain storage size in gigabytes (" << NKikimr::NBlobDepot::FormatByteSize(DEFAULT_STORAGE_SIZE) << " by default)")
             .RequiredArgument("uint")

diff --git a/ydb/tests/tools/kqprun/src/common.h b/ydb/tests/tools/kqprun/src/common.h
@@ -9,6 +9,8 @@
 #include <ydb/public/api/protos/ydb_cms.pb.h>
 #include <ydb/public/lib/ydb_cli/common/formats.h>
 
+#include <ydb/tests/tools/kqprun/src/proto/storage_meta.pb.h>
+
 #include <yql/essentials/minikql/computation/mkql_computation_node.h>
 #include <yql/essentials/minikql/mkql_function_registry.h>
 
@@ -48,9 +50,7 @@ struct TYdbSetupSettings {
 
     ui32 NodeCount = 1;
     TString DomainName = "Root";
-    std::unordered_set<TString> DedicatedTenants;
-    std::unordered_set<TString> SharedTenants;
-    std::unordered_set<TString> ServerlessTenants;
+    std::map<TString, TStorageMeta::TTenant> Tenants;
     TDuration HealthCheckTimeout = TDuration::Seconds(10);
     EHealthCheck HealthCheckLevel = EHealthCheck::NodesCount;
     bool SameSession = false;

diff --git a/ydb/tests/tools/kqprun/src/proto/storage_meta.proto b/ydb/tests/tools/kqprun/src/proto/storage_meta.proto
@@ -3,6 +3,20 @@ syntax = "proto3";
 package NKqpRun;
 
 message TStorageMeta {
+    message TTenant {
+        enum EType {
+            DEDICATED = 0;
+            SHARED = 1;
+            SERVERLESS = 2;
+        }
+
+        EType Type = 1;
+        uint32 NodesCount = 2;
+        string SharedTenant = 3;  // Only for serverless tenants
+    }
+
     uint64 StorageGeneration = 1;
     uint64 StorageSize = 2;
+    string DomainName = 3;
+    map<string, TTenant> Tenants = 4;
 }