From d5c96831215f7c194bab10118304a1bd588d6b38 Mon Sep 17 00:00:00 2001 From: rtjd6554 Date: Thu, 20 Feb 2025 11:31:35 +0000 Subject: [PATCH] Update wrapper for split partitions nested class --- .../sleeper/core/partition/PartitionsBuilder.java | 4 +++- .../testutils/StateStoreUpdatesWrapper.java | 14 ++++++++++++++ .../TransactionLogPartitionStoreTest.java | 12 ++++++------ 3 files changed, 23 insertions(+), 7 deletions(-) diff --git a/java/core/src/test/java/sleeper/core/partition/PartitionsBuilder.java b/java/core/src/test/java/sleeper/core/partition/PartitionsBuilder.java index ff62e6a908..f2bccbd45f 100644 --- a/java/core/src/test/java/sleeper/core/partition/PartitionsBuilder.java +++ b/java/core/src/test/java/sleeper/core/partition/PartitionsBuilder.java @@ -25,6 +25,8 @@ import java.util.Optional; import java.util.stream.Collectors; +import static sleeper.core.statestore.testutils.StateStoreUpdatesWrapper.update; + /** * A convenience class for specifying a partition tree. This includes methods to define a tree to be readable in a test, * including shorthand which would not be possible with {@link PartitionFactory}. @@ -122,7 +124,7 @@ public void applySplit(StateStore stateStore, String partitionId) { Partition toSplit = partitionById(partitionId); Partition left = partitionById(toSplit.getChildPartitionIds().get(0)); Partition right = partitionById(toSplit.getChildPartitionIds().get(1)); - stateStore.atomicallyUpdatePartitionAndCreateNewOnes(toSplit, left, right); + update(stateStore).atomicallyUpdatePartitionAndCreateNewOnes(toSplit, left, right); } /** diff --git a/java/core/src/test/java/sleeper/core/statestore/testutils/StateStoreUpdatesWrapper.java b/java/core/src/test/java/sleeper/core/statestore/testutils/StateStoreUpdatesWrapper.java index 77c9feb54c..d9d5707040 100644 --- a/java/core/src/test/java/sleeper/core/statestore/testutils/StateStoreUpdatesWrapper.java +++ b/java/core/src/test/java/sleeper/core/statestore/testutils/StateStoreUpdatesWrapper.java @@ -44,6 +44,7 @@ import sleeper.core.statestore.transactionlog.transaction.impl.InitialisePartitionsTransaction; import sleeper.core.statestore.transactionlog.transaction.impl.ReplaceFileReferencesTransaction; import sleeper.core.statestore.transactionlog.transaction.impl.SplitFileReferencesTransaction; +import sleeper.core.statestore.transactionlog.transaction.impl.SplitPartitionTransaction; import java.util.List; @@ -91,6 +92,19 @@ public void initialise(List partitions) throws StateStoreException { addTransaction(new InitialisePartitionsTransaction(partitions)); } + /** + * Atomically splits a partition to create child partitions. Updates the existing partition to record it as split, + * and creates new leaf partitions. + * + * @param splitPartition The {@link Partition} to be updated (must refer to the new leaves as children). + * @param newPartition1 The first new {@link Partition} (must be a leaf partition). + * @param newPartition2 The second new {@link Partition} (must be a leaf partition). + * @throws StateStoreException if split is not valid or update fails + */ + public void atomicallyUpdatePartitionAndCreateNewOnes(Partition splitPartition, Partition newPartition1, Partition newPartition2) throws StateStoreException { + addTransaction(new SplitPartitionTransaction(splitPartition, List.of(newPartition1, newPartition2))); + } + /** * Adds a file to the table, with one reference. * diff --git a/java/core/src/test/java/sleeper/core/statestore/transactionlog/TransactionLogPartitionStoreTest.java b/java/core/src/test/java/sleeper/core/statestore/transactionlog/TransactionLogPartitionStoreTest.java index 2394b60312..fcfd12f158 100644 --- a/java/core/src/test/java/sleeper/core/statestore/transactionlog/TransactionLogPartitionStoreTest.java +++ b/java/core/src/test/java/sleeper/core/statestore/transactionlog/TransactionLogPartitionStoreTest.java @@ -238,7 +238,7 @@ public void shouldSplitAPartition() { .buildTree(); // When - store.atomicallyUpdatePartitionAndCreateNewOnes( + update(store).atomicallyUpdatePartitionAndCreateNewOnes( tree.getPartition("root"), tree.getPartition("leftChild"), tree.getPartition("rightChild")); @@ -280,13 +280,13 @@ public void shouldFailSplittingAPartitionWhichHasAlreadyBeenSplit() { .rootFirst("root") .splitToNewChildren("root", "leftChild", "rightChild", 0L) .buildTree(); - store.atomicallyUpdatePartitionAndCreateNewOnes( + update(store).atomicallyUpdatePartitionAndCreateNewOnes( tree.getPartition("root"), tree.getPartition("leftChild"), tree.getPartition("rightChild")); // When / Then - assertThatThrownBy(() -> store.atomicallyUpdatePartitionAndCreateNewOnes( + assertThatThrownBy(() -> update(store).atomicallyUpdatePartitionAndCreateNewOnes( tree.getPartition("root"), tree.getPartition("leftChild"), tree.getPartition("rightChild"))) @@ -305,7 +305,7 @@ public void shouldFailSplittingAPartitionWithWrongChildren() { .buildTree(); // When / Then - assertThatThrownBy(() -> store.atomicallyUpdatePartitionAndCreateNewOnes( + assertThatThrownBy(() -> update(store).atomicallyUpdatePartitionAndCreateNewOnes( tree.getPartition("root"), tree.getPartition("LL"), tree.getPartition("LR"))) @@ -328,7 +328,7 @@ public void shouldFailSplittingAPartitionWithChildrenOfWrongParent() { .buildTree(); // When / Then - assertThatThrownBy(() -> store.atomicallyUpdatePartitionAndCreateNewOnes( + assertThatThrownBy(() -> update(store).atomicallyUpdatePartitionAndCreateNewOnes( parentTree.getPartition("root"), childrenTree.getPartition("child1"), childrenTree.getPartition("child2"))) @@ -347,7 +347,7 @@ public void shouldFailSplittingAPartitionWhenNewPartitionIsNotALeaf() { .buildTree(); // When / Then - assertThatThrownBy(() -> store.atomicallyUpdatePartitionAndCreateNewOnes( + assertThatThrownBy(() -> update(store).atomicallyUpdatePartitionAndCreateNewOnes( tree.getPartition("root"), tree.getPartition("L"), // Not a leaf tree.getPartition("R")))