Make it possible to change in-memory setting for tables

snaury · snaury · commit db635acffd1c · 2024-11-28T13:34:05.000Z
diff --git a/ydb/core/ydb_convert/column_families.h b/ydb/core/ydb_convert/column_families.h
@@ -183,13 +183,15 @@ namespace NKikimr {
                 case Ydb::FeatureFlag::STATUS_UNSPECIFIED:
                     break;
                 case Ydb::FeatureFlag::ENABLED:
-                    *code = Ydb::StatusIds::BAD_REQUEST;
-                    *error = TStringBuilder()
-                        << "Setting keep_in_memory to ENABLED is not supported in column family '"
-                        << familySettings.name() << "'";
-                    return false;
+                    if (!AppData()->FeatureFlags.GetEnablePublicApiKeepInMemory()) {
+                        *code = Ydb::StatusIds::BAD_REQUEST;
+                        *error = "Setting keep_in_memory to ENABLED is not allowed";
+                        return false;
+                    }
+                    family->SetColumnCache(NKikimrSchemeOp::ColumnCacheEver);
+                    break;
                 case Ydb::FeatureFlag::DISABLED:
-                    family->ClearColumnCache();
+                    family->SetColumnCache(NKikimrSchemeOp::ColumnCacheNone);
                     break;
                 default:
                     *code = Ydb::StatusIds::BAD_REQUEST;
diff --git a/ydb/public/sdk/cpp/client/ydb_table/table.cpp b/ydb/public/sdk/cpp/client/ydb_table/table.cpp
@@ -1115,6 +1115,11 @@ TColumnFamilyBuilder& TColumnFamilyBuilder::SetCompression(EColumnFamilyCompress
     return *this;
 }
 
+TColumnFamilyBuilder& TColumnFamilyBuilder::SetKeepInMemory(bool enabled) {
+    Impl_->Proto.set_keep_in_memory(enabled ? Ydb::FeatureFlag::ENABLED : Ydb::FeatureFlag::DISABLED);
+    return *this;
+}
+
 TColumnFamilyDescription TColumnFamilyBuilder::Build() const {
     return TColumnFamilyDescription(Impl_->Proto);
 }
diff --git a/ydb/public/sdk/cpp/client/ydb_table/table.h b/ydb/public/sdk/cpp/client/ydb_table/table.h
@@ -806,6 +806,7 @@ class TColumnFamilyBuilder {
 
     TColumnFamilyBuilder& SetData(const TString& media);
     TColumnFamilyBuilder& SetCompression(EColumnFamilyCompression compression);
+    TColumnFamilyBuilder& SetKeepInMemory(bool enabled);
 
     TColumnFamilyDescription Build() const;
 
@@ -868,6 +869,11 @@ class TTableColumnFamilyBuilder {
         return *this;
     }
 
+    TTableColumnFamilyBuilder& SetKeepInMemory(bool enabled) {
+        Builder_.SetKeepInMemory(enabled);
+        return *this;
+    }
+
     TTableBuilder& EndColumnFamily();
 
 private:
@@ -1486,6 +1492,11 @@ class TAlterColumnFamilyBuilder {
         return *this;
     }
 
+    TAlterColumnFamilyBuilder& SetKeepInMemory(bool enabled) {
+        Builder_.SetKeepInMemory(enabled);
+        return *this;
+    }
+
     TAlterTableSettings& EndAddColumnFamily();
     TAlterTableSettings& EndAlterColumnFamily();
 
diff --git a/ydb/services/ydb/ydb_ut.cpp b/ydb/services/ydb/ydb_ut.cpp
@@ -824,6 +824,87 @@ Y_UNIT_TEST_SUITE(TGRpcNewClient) {
         client.CreateSession().Apply(createSessionHandler).Wait();
         UNIT_ASSERT(done);
     }
+
+    Y_UNIT_TEST(InMemoryTables) {
+        TKikimrWithGrpcAndRootSchemaNoSystemViews server;
+        server.Server_->GetRuntime()->GetAppData().FeatureFlags.SetEnablePublicApiKeepInMemory(true);
+
+        ui16 grpc = server.GetPort();
+        TString location = TStringBuilder() << "localhost:" << grpc;
+
+        auto connection = NYdb::TDriver(
+            TDriverConfig()
+                .SetEndpoint(location));
+
+        auto client = NYdb::NTable::TTableClient(connection);
+        auto createSessionResult = client.CreateSession().ExtractValueSync();
+        UNIT_ASSERT(!createSessionResult.IsTransportError());
+        auto session = createSessionResult.GetSession();
+
+        auto createTableResult = session.CreateTable("/Root/Table", client.GetTableBuilder()
+            .AddNullableColumn("Key", EPrimitiveType::Int32)
+            .AddNullableColumn("Value", EPrimitiveType::String)
+            .SetPrimaryKeyColumn("Key")
+            // Note: only needed because this test doesn't initial table profiles
+            .BeginStorageSettings()
+                .SetTabletCommitLog0("ssd")
+                .SetTabletCommitLog1("ssd")
+            .EndStorageSettings()
+            .BeginColumnFamily("default")
+                .SetData("ssd")
+                .SetKeepInMemory(true)
+            .EndColumnFamily()
+            .Build()).ExtractValueSync();
+        UNIT_ASSERT_C(createTableResult.IsSuccess(), (NYdb::TStatus&)createTableResult);
+
+        {
+            auto describeTableResult = session.DescribeTable("/Root/Table").ExtractValueSync();
+            UNIT_ASSERT_C(describeTableResult.IsSuccess(), (NYdb::TStatus&)describeTableResult);
+            auto desc = describeTableResult.GetTableDescription();
+            auto families = desc.GetColumnFamilies();
+            UNIT_ASSERT_VALUES_EQUAL(families.size(), 1u);
+            auto family = families.at(0);
+            UNIT_ASSERT_VALUES_EQUAL(family.GetKeepInMemory(), true);
+        }
+
+        {
+            auto alterTableResult = session.AlterTable("/Root/Table", NYdb::NTable::TAlterTableSettings()
+                .BeginAlterColumnFamily("default")
+                    .SetKeepInMemory(false)
+                .EndAlterColumnFamily()).ExtractValueSync();
+            UNIT_ASSERT_C(alterTableResult.IsSuccess(), (NYdb::TStatus&)alterTableResult);
+        }
+
+        {
+            auto describeTableResult = session.DescribeTable("/Root/Table").ExtractValueSync();
+            UNIT_ASSERT_C(describeTableResult.IsSuccess(), (NYdb::TStatus&)describeTableResult);
+            auto desc = describeTableResult.GetTableDescription();
+            auto families = desc.GetColumnFamilies();
+            UNIT_ASSERT_VALUES_EQUAL(families.size(), 1u);
+            auto family = families.at(0);
+            // Note: server cannot currently distinguish between implicitly
+            // unset and explicitly disabled, so it returns the former.
+            UNIT_ASSERT_VALUES_EQUAL(family.GetKeepInMemory(), Nothing());
+        }
+
+        {
+            auto alterTableResult = session.AlterTable("/Root/Table", NYdb::NTable::TAlterTableSettings()
+                .BeginAlterColumnFamily("default")
+                    .SetKeepInMemory(true)
+                .EndAlterColumnFamily()).ExtractValueSync();
+            UNIT_ASSERT_C(alterTableResult.IsSuccess(), (NYdb::TStatus&)alterTableResult);
+        }
+
+        {
+            auto describeTableResult = session.DescribeTable("/Root/Table").ExtractValueSync();
+            UNIT_ASSERT_C(describeTableResult.IsSuccess(), (NYdb::TStatus&)describeTableResult);
+            auto desc = describeTableResult.GetTableDescription();
+            auto families = desc.GetColumnFamilies();
+            UNIT_ASSERT_VALUES_EQUAL(families.size(), 1u);
+            auto family = families.at(0);
+            UNIT_ASSERT_VALUES_EQUAL(family.GetKeepInMemory(), true);
+        }
+    }
 }
 
 static TString CreateSession(std::shared_ptr<grpc::Channel> channel) {

Original file line number	Diff line number	Diff line change
`@@ -1115,6 +1115,11 @@ TColumnFamilyBuilder& TColumnFamilyBuilder::SetCompression(EColumnFamilyCompress`
`1115`	`1115`	`return *this;`
`1116`	`1116`	`}`
`1117`	`1117`
	`1118`	`+TColumnFamilyBuilder& TColumnFamilyBuilder::SetKeepInMemory(bool enabled) {`
	`1119`	`+ Impl_->Proto.set_keep_in_memory(enabled ? Ydb::FeatureFlag::ENABLED : Ydb::FeatureFlag::DISABLED);`
	`1120`	`+ return *this;`
	`1121`	`+}`
	`1122`	`+`
`1118`	`1123`	`TColumnFamilyDescription TColumnFamilyBuilder::Build() const {`
`1119`	`1124`	`return TColumnFamilyDescription(Impl_->Proto);`
`1120`	`1125`	`}`