Enable quoting disk space usage by storage pool kind

ydb/core/cms/console/console__create_tenant.cpp

@@ -270,7 +270,26 @@ class TTenantsManager::TTxCreateTenant : public TTransactionBase<TTenantsManager
             auto hardQuota = quotas.data_size_hard_quota();
             auto softQuota = quotas.data_size_soft_quota();
             if (hardQuota && softQuota && hardQuota < softQuota) {
-                return Error(Ydb::StatusIds::BAD_REQUEST, "Data size soft quota cannot be larger than hard quota", ctx);
+                return Error(Ydb::StatusIds::BAD_REQUEST,
+                    TStringBuilder() << "Overall data size soft quota (" << softQuota << ")"
+                                     << " of the database " << path
+                                     << " must be smaller than the hard quota (" << hardQuota << ")",
+                    ctx
+                );
+            }
+            for (const auto& storageQuota : quotas.storage_quotas()) {
+                const auto unitHardQuota = storageQuota.data_size_hard_quota();
+                const auto unitSoftQuota = storageQuota.data_size_soft_quota();
+                if (unitHardQuota && unitSoftQuota && unitHardQuota < unitSoftQuota) {
+                    return Error(Ydb::StatusIds::BAD_REQUEST,
+                        TStringBuilder() << "Data size soft quota (" << unitSoftQuota << ")"
+                                         << " for a " << storageQuota.unit_kind() << " storage unit "
+                                         << " of the database " << path
+                                         << " must be smaller than the corresponding hard quota (" << unitHardQuota << ")",
+                        ctx
+                    );
+                }
+
             }
             Tenant->DatabaseQuotas.ConstructInPlace(quotas);
         }

ydb/core/cms/console/console_ut_tenants.cpp

@@ -90,6 +90,20 @@ TTenantTestConfig DefaultConsoleTestConfig()
     return res;
 }
 
+TString DefaultDatabaseQuotas() {
+    return R"(
+        data_size_hard_quota: 3000
+        storage_quotas {
+            unit_kind: "hdd"
+            data_size_hard_quota: 2000
+        }
+        storage_quotas {
+            unit_kind: "hdd-1"
+            data_size_hard_quota: 1000
+        }
+    )";
+}
+
 void CheckAlterTenantSlots(TTenantTestRuntime &runtime, const TString &path,
                            ui64 generation, Ydb::StatusIds::StatusCode code,
                            TVector<TSlotRequest> add,
@@ -2027,6 +2041,58 @@ Y_UNIT_TEST_SUITE(TConsoleTests) {
         TTenantTestRuntime runtime(DefaultConsoleTestConfig(), {}, true);
         RunTestAlterTenantTooManyStorageResourcesForRunning(runtime);
     }
+
+    void RunTestDatabaseQuotas(TTenantTestRuntime& runtime, const TString& quotas, bool shared = false) {
+        using EType = TCreateTenantRequest::EType;
+
+        CheckCreateTenant(runtime, Ydb::StatusIds::SUCCESS,
+            TCreateTenantRequest(TENANT1_1_NAME, shared ? EType::Shared : EType::Common)
+                .WithPools({{"hdd", 1}, {"hdd-1", 1}})
+                .WithDatabaseQuotas(quotas)
+        );
+
+        RestartTenantPool(runtime);
+
+        CheckTenantStatus(runtime, TENANT1_1_NAME, shared, Ydb::StatusIds::SUCCESS,
+                          Ydb::Cms::GetDatabaseStatusResult::RUNNING,
+                          {{"hdd", 1, 1}, {"hdd-1", 1, 1}}, {});
+    }
+
+    Y_UNIT_TEST(TestDatabaseQuotas) {
+        TTenantTestRuntime runtime(DefaultConsoleTestConfig());
+        RunTestDatabaseQuotas(runtime, DefaultDatabaseQuotas());
+    }
+
+    Y_UNIT_TEST(TestDatabaseQuotasBadOverallQuota) {
+        TTenantTestRuntime runtime(DefaultConsoleTestConfig());
+
+        CheckCreateTenant(runtime, Ydb::StatusIds::BAD_REQUEST,
+            TCreateTenantRequest(TENANT1_1_NAME, TCreateTenantRequest::EType::Common)
+                .WithPools({{"hdd", 1}})
+                .WithDatabaseQuotas(R"(
+                        data_size_hard_quota: 1
+                        data_size_soft_quota: 1000
+                    )"
+                )
+        );
+    }
+
+    Y_UNIT_TEST(TestDatabaseQuotasBadStorageQuota) {
+        TTenantTestRuntime runtime(DefaultConsoleTestConfig());
+
+        CheckCreateTenant(runtime, Ydb::StatusIds::BAD_REQUEST,
+            TCreateTenantRequest(TENANT1_1_NAME, TCreateTenantRequest::EType::Common)
+                .WithPools({{"hdd", 1}})
+                .WithDatabaseQuotas(R"(
+                        storage_quotas {
+                            unit_kind: "hdd"
+                            data_size_hard_quota: 1
+                            data_size_soft_quota: 1000
+                        }
+                    )"
+                )
+        );
+    }
 }
 
 } // namespace NKikimr

ydb/core/protos/subdomains.proto

@@ -61,8 +61,17 @@ message TDiskSpaceUsage {
         optional uint64 UsedReserveSize = 4;
     }
 
+    message TStoragePoolUsage {
+        // in bytes
+        optional string PoolKind = 1;
+        optional uint64 TotalSize = 2;
+        optional uint64 DataSize = 3;
+        optional uint64 IndexSize = 4;
+    }
+
     optional TTables Tables = 1;
     optional TTopics Topics = 2;
+    repeated TStoragePoolUsage StoragePoolsUsage = 3;
 }
 
 message TDomainState {

ydb/core/protos/table_stats.proto

@@ -16,6 +16,15 @@ message TChannelStats {
     optional uint64 IndexSize = 3;
 }
 
+message TStoragePoolsStats {
+    message TPoolUsage {
+        optional string PoolKind = 1;
+        optional uint64 DataSize = 2;
+        optional uint64 IndexSize = 3;
+    }
+    repeated TPoolUsage PoolsUsage = 1;
+}
+
 message TTableStats {
     optional uint64 DataSize = 1; // both inMem and ondisk
     optional uint64 RowCount = 2; // both inMem and ondisk
@@ -55,4 +64,6 @@ message TTableStats {
     optional bool HasLoanedParts = 29;
 
     repeated TChannelStats Channels = 30;
+
+    optional TStoragePoolsStats StoragePools = 31;
 }

ydb/core/testlib/tenant_helpers.h

@@ -96,6 +96,7 @@ struct TCreateTenantRequest {
     TString Path;
     EType Type;
     TAttrsCont Attrs;
+    Ydb::Cms::DatabaseQuotas DatabaseQuotas;
     // Common & Shared
     TPoolsCont Pools;
     TSlotsCont Slots;
@@ -116,6 +117,18 @@ struct TCreateTenantRequest {
         return *this;
     }
 
+    TSelf& WithDatabaseQuotas(const Ydb::Cms::DatabaseQuotas& quotas) {
+        DatabaseQuotas = quotas;
+        return *this;
+    }
+
+    TSelf& WithDatabaseQuotas(const TString& quotas) {
+        Ydb::Cms::DatabaseQuotas parsedQuotas;
+        UNIT_ASSERT_C(NProtoBuf::TextFormat::ParseFromString(quotas, &parsedQuotas), quotas);
+        DatabaseQuotas = std::move(parsedQuotas);
+        return *this;
+    }
+
     TSelf& WithPools(const TPoolsCont& pools) {
         if (Type == EType::Unspecified) {
             Type = EType::Common;
@@ -340,14 +353,16 @@ inline void CheckCreateTenant(TTenantTestRuntime &runtime,
     if (request.PlanResolution) {
         event->Record.MutableRequest()->mutable_options()->set_plan_resolution(request.PlanResolution);
     }
+
+    event->Record.MutableRequest()->mutable_database_quotas()->CopyFrom(request.DatabaseQuotas);
 
     TAutoPtr<IEventHandle> handle;
     runtime.SendToConsole(event);
     auto reply = runtime.GrabEdgeEventRethrow<NConsole::TEvConsole::TEvCreateTenantResponse>(handle);
     auto &operation = reply->Record.GetResponse().operation();
 
     if (operation.ready()) {
-        UNIT_ASSERT_VALUES_EQUAL(operation.status(), code);
+        UNIT_ASSERT_VALUES_EQUAL_C(operation.status(), code, operation.DebugString());
     } else {
         TString id = operation.id();
         auto *request = new NConsole::TEvConsole::TEvNotifyOperationCompletionRequest;

ydb/core/tx/datashard/ut_common/datashard_ut_common.cpp

@@ -1831,7 +1831,12 @@ void ExecSQL(Tests::TServer::TPtr server,
     auto request = MakeSQLRequest(sql, dml);
     runtime.Send(new IEventHandle(NKqp::MakeKqpProxyID(runtime.GetNodeId()), sender, request.Release(), 0, 0, nullptr));
     auto ev = runtime.GrabEdgeEventRethrow<NKqp::TEvKqp::TEvQueryResponse>(sender);
-    UNIT_ASSERT_VALUES_EQUAL(ev->Get()->Record.GetRef().GetYdbStatus(), code);
+    auto& response = ev->Get()->Record.GetRef();
+    auto& issues = response.GetResponse().GetQueryIssues();
+    UNIT_ASSERT_VALUES_EQUAL_C(response.GetYdbStatus(),
+                               code,
+                               issues.empty() ? response.DebugString() : issues.Get(0).DebugString()
+    );
 }
 
 std::unique_ptr<NEvents::TDataEvents::TEvWrite> MakeWriteRequest(ui64 txId, NKikimrDataEvents::TEvWrite::ETxMode txMode, NKikimrDataEvents::TEvWrite_TOperation::EOperationType operationType, const TTableId& tableId, const TVector<TShardedTableOptions::TColumn>& columns, ui32 rowCount, ui64 seed) {

ydb/core/tx/schemeshard/schemeshard__init.cpp

@@ -2259,6 +2259,22 @@ struct TSchemeShard::TTxInit : public TTransactionBase<TSchemeShard> {
                 stats.RowCount = rowSet.GetValue<Schema::TablePartitionStats::RowCount>();
                 stats.DataSize = rowSet.GetValue<Schema::TablePartitionStats::DataSize>();
                 stats.IndexSize = rowSet.GetValue<Schema::TablePartitionStats::IndexSize>();
+                if (rowSet.HaveValue<Schema::TablePartitionStats::StoragePoolsStats>()) {
+                    NKikimrTableStats::TStoragePoolsStats protobufRepresentation;
+                    Y_ABORT_UNLESS(ParseFromStringNoSizeLimit(
+                            protobufRepresentation,
+                            rowSet.GetValue<Schema::TablePartitionStats::StoragePoolsStats>()
+                        )
+                    );
+                    for (const auto& poolUsage : protobufRepresentation.GetPoolsUsage()) {
+                        stats.StoragePoolsStats.emplace(
+                            poolUsage.GetPoolKind(),
+                            TPartitionStats::TStoragePoolStats{poolUsage.GetDataSize(),
+                                                               poolUsage.GetIndexSize()
+                            }
+                        );
+                    }
+                }
 
                 stats.LastAccessTime = TInstant::FromValue(rowSet.GetValue<Schema::TablePartitionStats::LastAccessTime>());
                 stats.LastUpdateTime = TInstant::FromValue(rowSet.GetValue<Schema::TablePartitionStats::LastUpdateTime>());

ydb/core/tx/schemeshard/schemeshard__operation_alter_extsubdomain.cpp

@@ -258,6 +258,16 @@ VerifyParams(TParamsDelta* delta, const TPathId pathId, const TSubDomainInfo::TP
             }
         }
 
+        // storage pools quotas check
+        TString error;
+        if (const auto& effectivePools = requestedPools.empty()
+                ? actualPools
+                : requestedPools;
+            !CheckStorageQuotasKinds(input.GetDatabaseQuotas(), effectivePools, pathId.ToString(), error)
+        ) {
+            return paramError(error);
+        }
+
         std::set_difference(requestedPools.begin(), requestedPools.end(),
                             actualPools.begin(), actualPools.end(),
                             std::back_inserter(storagePoolsAdded));

ydb/core/tx/schemeshard/schemeshard__operation_alter_subdomain.cpp

@@ -287,6 +287,14 @@ class TAlterSubDomain: public TSubOperation {
         }
 
         if (settings.HasDatabaseQuotas()) {
+            if (const auto& effectivePools = requestedPools.empty()
+                    ? actualPools
+                    : requestedPools;
+                !CheckStorageQuotasKinds(settings.GetDatabaseQuotas(), effectivePools, path.PathString(), errStr)
+            ) {
+                result->SetError(NKikimrScheme::StatusInvalidParameter, errStr);
+                return result;
+            }
             alterData->SetDatabaseQuotas(settings.GetDatabaseQuotas());
         }
 

ydb/core/tx/schemeshard/schemeshard__operation_common_subdomain.h

@@ -6,6 +6,37 @@
 namespace NKikimr {
 namespace NSchemeShard {
 
+inline bool CheckStorageQuotasKinds(const Ydb::Cms::DatabaseQuotas& quotas,
+                                    const TVector<TStoragePool>& pools,
+                                    const TString& path,
+                                    TString& error
+) {
+    TVector<TString> quotedKinds;
+    for (const auto& storageQuota : quotas.storage_quotas()) {
+        quotedKinds.emplace_back(storageQuota.unit_kind());
+    }
+    Sort(quotedKinds);
+    const auto uniqueEnd = Unique(quotedKinds.begin(), quotedKinds.end());
+    if (uniqueEnd != quotedKinds.end()) {
+        error = TStringBuilder()
+            << "Malformed subdomain request: storage quotas' unit kinds must be unique, but "
+            << *uniqueEnd << " appears twice in the storage quotas definition of the " << path << " subdomain.";
+        return false;
+    }
+
+    for (const auto& quotedKind : quotedKinds) {
+        if (!AnyOf(pools, [&quotedKind](const TStoragePool& pool) {
+            return pool.GetKind() == quotedKind;
+        })) {
+            error = TStringBuilder()
+                << "Malformed subdomain request: cannot set a " << quotedKind << " storage quota, "
+                << "because no storage pool in the subdomain " << path << " has the specified kind.";
+            return false;
+        }
+    }
+    return true;
+}
+
 namespace NSubDomainState {
 
 class TConfigureParts: public TSubOperationState {

ydb/core/tx/schemeshard/schemeshard__operation_create_subdomain.cpp

@@ -295,6 +295,12 @@ class TCreateSubDomain: public TSubOperation {
         }
 
         if (settings.HasDatabaseQuotas()) {
+            if (!requestedStoragePools.empty()
+                    && !CheckStorageQuotasKinds(settings.GetDatabaseQuotas(), requestedStoragePools, dstPath.PathString(), errStr)
+            ) {
+                result->SetError(NKikimrScheme::StatusInvalidParameter, errStr);
+                return result;
+            }
             alter->SetDatabaseQuotas(settings.GetDatabaseQuotas());
         }