Update wrapper for split partitions nested class

gchq · Feb 20, 2025 · d5c9683 · d5c9683
1 parent fc039e3
commit d5c9683
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 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
@@ -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
@@ -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<Partition> 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/...rc/test/java/sleeper/core/statestore/transactionlog/TransactionLogPartitionStoreTest.java b/...rc/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")))