add new tests

Author: jepett0

ydb/services/ydb/backup_ut/ydb_backup_ut.cpp

@@ -16,10 +16,27 @@
 #include <library/cpp/testing/unittest/registar.h>
 
 #include <aws/core/Aws.h>
+#include <google/protobuf/util/message_differencer.h>
 
 using namespace NYdb;
 using namespace NYdb::NTable;
 
+namespace NYdb::NTable {
+
+bool operator==(const TValue& lhs, const TValue& rhs) {
+    return google::protobuf::util::MessageDifferencer::Equals(lhs.GetProto(), rhs.GetProto());
+}
+
+bool operator==(const TKeyBound& lhs, const TKeyBound& rhs) {
+    return lhs.GetValue() == rhs.GetValue() && lhs.IsInclusive() == rhs.IsInclusive();
+}
+
+bool operator==(const TKeyRange& lhs, const TKeyRange& rhs) {
+    return lhs.From() == lhs.From() && lhs.To() == rhs.To();
+}
+
+}
+
 namespace {
 
 #define DEBUG_HINT (TStringBuilder() << "at line " << __LINE__)
@@ -194,6 +211,106 @@ void TestIndexTablePartitioningSettingsArePreserved(
     CheckTableDescription(session, indexTablePath, CreateMinPartitionsChecker(minIndexPartitions, DEBUG_HINT));
 }
 
+void TestTableSplitBoundariesArePreserved(
+    const char* table, ui64 partitions, TSession& session, TBackupFunction&& backup, TRestoreFunction&& restore
+) {
+    ExecuteDataDefinitionQuery(session, Sprintf(R"(
+            CREATE TABLE `%s` (
+                Key Uint32,
+                Value Utf8,
+                PRIMARY KEY (Key)
+            )
+            WITH (
+                PARTITION_AT_KEYS = (1, 2, 4, 8, 16, 32, 64, 128, 256)
+            );
+        )",
+        table
+    ));
+    const auto describeSettings = TDescribeTableSettings()
+            .WithTableStatistics(true)
+            .WithKeyShardBoundary(true);
+    const auto originalTableDescription = GetTableDescription(session, table, describeSettings);
+    UNIT_ASSERT_VALUES_EQUAL(originalTableDescription.GetPartitionsCount(), partitions);
+    const auto& originalKeyRanges = originalTableDescription.GetKeyRanges();
+    UNIT_ASSERT_VALUES_EQUAL(originalKeyRanges.size(), partitions);
+
+    backup(table);
+
+    ExecuteDataDefinitionQuery(session, Sprintf(R"(
+            DROP TABLE `%s`;
+        )", table
+    ));
+
+    restore(table);
+    const auto restoredTableDescription = GetTableDescription(session, table, describeSettings);
+    UNIT_ASSERT_VALUES_EQUAL(restoredTableDescription.GetPartitionsCount(), partitions);
+    const auto& restoredKeyRanges = restoredTableDescription.GetKeyRanges();
+    UNIT_ASSERT_VALUES_EQUAL(restoredKeyRanges.size(), partitions);
+    UNIT_ASSERT_EQUAL(restoredTableDescription.GetKeyRanges(), originalKeyRanges);
+}
+
+void TestIndexTableSplitBoundariesArePreserved(
+    const char* table, const char* index, ui64 indexPartitions, TSession& session,
+    TBackupFunction&& backup, TRestoreFunction&& restore
+) {
+    const TString indexTablePath = JoinFsPaths(table, index, "indexImplTable");
+
+    {
+        TExplicitPartitions indexPartitionBoundaries;
+        for (ui32 boundary : {1, 2, 4, 8, 16, 32, 64, 128, 256}) {
+            indexPartitionBoundaries.AppendSplitPoints(
+                // split boundary is technically always a tuple
+                TValueBuilder().BeginTuple().AddElement().OptionalUint32(boundary).EndTuple().Build()
+            );
+        }
+        // By default indexImplTable has auto partitioning by size enabled,
+        // so you must set min partition count for partitions to not merge immediately after indexImplTable is built.
+        TPartitioningSettingsBuilder partitioningSettingsBuilder;
+        partitioningSettingsBuilder
+            .SetMinPartitionsCount(indexPartitions)
+            .SetMaxPartitionsCount(indexPartitions);
+
+        const auto indexSettings = TGlobalIndexSettings{
+            .PartitioningSettings = partitioningSettingsBuilder.Build(),
+            .Partitions = std::move(indexPartitionBoundaries)
+        };
+
+        auto tableBuilder = TTableBuilder()
+            .AddNullableColumn("Key", EPrimitiveType::Uint32)
+            .AddNullableColumn("Value", EPrimitiveType::Uint32)
+            .SetPrimaryKeyColumn("Key")
+            .AddSecondaryIndex("byValue", EIndexType::GlobalSync, { "Value" }, {}, indexSettings);
+
+        const auto result = session.CreateTable(table, tableBuilder.Build()).ExtractValueSync();
+        UNIT_ASSERT_C(result.IsSuccess(), result.GetIssues().ToString());
+    }
+    const auto describeSettings = TDescribeTableSettings()
+            .WithTableStatistics(true)
+            .WithKeyShardBoundary(true);
+    const auto originalIndexTableDescription = GetTableDescription(
+        session, indexTablePath, describeSettings
+    );
+    UNIT_ASSERT_VALUES_EQUAL(originalIndexTableDescription.GetPartitionsCount(), indexPartitions);
+    const auto& originalKeyRanges = originalIndexTableDescription.GetKeyRanges();
+    UNIT_ASSERT_VALUES_EQUAL(originalKeyRanges.size(), indexPartitions);
+
+    backup(table);
+
+    ExecuteDataDefinitionQuery(session, Sprintf(R"(
+            DROP TABLE `%s`;
+        )", table
+    ));
+
+    restore(table);
+    const auto restoredIndexTableDescription = GetTableDescription(
+        session, indexTablePath, describeSettings
+    );
+    UNIT_ASSERT_VALUES_EQUAL(restoredIndexTableDescription.GetPartitionsCount(), indexPartitions);
+    const auto& restoredKeyRanges = restoredIndexTableDescription.GetKeyRanges();
+    UNIT_ASSERT_VALUES_EQUAL(restoredKeyRanges.size(), indexPartitions);
+    UNIT_ASSERT_EQUAL(restoredKeyRanges, originalKeyRanges);
+}
+
 }
 
 Y_UNIT_TEST_SUITE(BackupRestore) {
@@ -279,6 +396,48 @@ Y_UNIT_TEST_SUITE(BackupRestore) {
         );
     }
 
+    Y_UNIT_TEST(RestoreTableSplitBoundaries) {
+        TKikimrWithGrpcAndRootSchema server;
+        auto driver = TDriver(TDriverConfig().SetEndpoint(Sprintf("localhost:%u", server.GetPort())));
+        TTableClient tableClient(driver);
+        auto session = tableClient.GetSession().ExtractValueSync().GetSession();
+        TTempDir tempDir;
+        const auto& pathToBackup = tempDir.Path();
+
+        constexpr const char* table = "/Root/table";
+        constexpr ui64 partitions = 10;
+
+        TestTableSplitBoundariesArePreserved(
+            table,
+            partitions,
+            session,
+            CreateBackupLambda(driver, pathToBackup, true),
+            CreateRestoreLambda(driver, pathToBackup)
+        );
+    }
+
+    Y_UNIT_TEST(RestoreIndexTableSplitBoundaries) {
+        TKikimrWithGrpcAndRootSchema server;
+        auto driver = TDriver(TDriverConfig().SetEndpoint(Sprintf("localhost:%u", server.GetPort())));
+        TTableClient tableClient(driver);
+        auto session = tableClient.GetSession().ExtractValueSync().GetSession();
+        TTempDir tempDir;
+        const auto& pathToBackup = tempDir.Path();
+
+        constexpr const char* table = "/Root/table";
+        constexpr const char* index = "byValue";
+        constexpr ui64 indexPartitions = 10;
+
+        TestIndexTableSplitBoundariesArePreserved(
+            table,
+            index,
+            indexPartitions,
+            session,
+            CreateBackupLambda(driver, pathToBackup, true),
+            CreateRestoreLambda(driver, pathToBackup)
+        );
+    }
+
 }
 
 Y_UNIT_TEST_SUITE(BackupRestoreS3) {
@@ -455,4 +614,34 @@ Y_UNIT_TEST_SUITE(BackupRestoreS3) {
         );
     }
 
+    Y_UNIT_TEST(RestoreTableSplitBoundaries) {
+        TS3TestEnv testEnv;
+        constexpr const char* table = "/Root/table";
+        constexpr ui64 partitions = 10;
+
+        TestTableSplitBoundariesArePreserved(
+            table,
+            partitions,
+            testEnv.GetSession(),
+            CreateBackupLambda(testEnv.GetDriver(), testEnv.GetS3Port()),
+            CreateRestoreLambda(testEnv.GetDriver(), testEnv.GetS3Port())
+        );
+    }
+
+    Y_UNIT_TEST(RestoreIndexTableSplitBoundaries) {
+        TS3TestEnv testEnv;
+        constexpr const char* table = "/Root/table";
+        constexpr const char* index = "byValue";
+        constexpr ui64 indexPartitions = 10;
+
+        TestIndexTableSplitBoundariesArePreserved(
+            table,
+            index,
+            indexPartitions,
+            testEnv.GetSession(),
+            CreateBackupLambda(testEnv.GetDriver(), testEnv.GetS3Port()),
+            CreateRestoreLambda(testEnv.GetDriver(), testEnv.GetS3Port())
+        );
+    }
+
 }