Use wrapper for atomicallyReplaceFileReferencesWithNewOnes

gchq · Feb 21, 2025 · 9c90d14 · 9c90d14
1 parent 7315580
commit 9c90d14
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 6 deletions.
diff --git a/...n/src/test/java/sleeper/compaction/job/execution/testutils/CompactionRunnerTestUtils.java b/...n/src/test/java/sleeper/compaction/job/execution/testutils/CompactionRunnerTestUtils.java
@@ -26,6 +26,7 @@
 import java.util.stream.Stream;
 
 import static java.util.stream.Collectors.toUnmodifiableList;
+import static sleeper.core.statestore.testutils.StateStoreUpdatesWrapper.update;
 
 public class CompactionRunnerTestUtils {
 
@@ -59,7 +60,7 @@ public static void assignJobIdToInputFiles(StateStore stateStore, CompactionJob
     }
 
     public static void assignJobIdsToInputFiles(StateStore stateStore, CompactionJob... jobs) throws Exception {
-        stateStore.assignJobIds(Stream.of(jobs)
+        update(stateStore).assignJobIds(Stream.of(jobs)
                 .map(CompactionJob::createAssignJobIdRequest)
                 .collect(toUnmodifiableList()));
     }

diff --git a/...em-test-dsl/src/test/java/sleeper/systemtest/dsl/testutil/drivers/InMemoryCompaction.java b/...em-test-dsl/src/test/java/sleeper/systemtest/dsl/testutil/drivers/InMemoryCompaction.java
@@ -31,8 +31,6 @@
 import sleeper.core.record.Record;
 import sleeper.core.schema.Schema;
 import sleeper.core.statestore.StateStore;
-import sleeper.core.statestore.transactionlog.AddTransactionRequest;
-import sleeper.core.statestore.transactionlog.transaction.impl.ReplaceFileReferencesTransaction;
 import sleeper.core.tracker.compaction.job.CompactionJobTracker;
 import sleeper.core.tracker.compaction.job.InMemoryCompactionJobTracker;
 import sleeper.core.tracker.compaction.job.query.CompactionJobStatus;
@@ -63,6 +61,7 @@
 import java.util.UUID;
 
 import static java.util.stream.Collectors.toUnmodifiableList;
+import static sleeper.core.statestore.testutils.StateStoreUpdatesWrapper.update;
 
 public class InMemoryCompaction {
     private final List<CompactionJob> queuedJobs = new ArrayList<>();
@@ -184,10 +183,9 @@ private JobRunSummary compact(CompactionJob job, TableProperties tableProperties
         Schema schema = tableProperties.getSchema();
         Partition partition = getPartitionForJob(stateStore, job);
         RecordsProcessed recordsProcessed = mergeInputFiles(job, partition, schema);
-        stateStore.addTransaction(AddTransactionRequest.withTransaction(new ReplaceFileReferencesTransaction(List.of(
+        update(stateStore).atomicallyReplaceFileReferencesWithNewOnes(List.of(
                 job.replaceFileReferencesRequestBuilder(recordsProcessed.getRecordsWritten())
-                        .taskId(run.getTaskId()).build())))
-                .build());
+                        .taskId(run.getTaskId()).build()));
         Instant finishTime = startTime.plus(Duration.ofMinutes(1));
         return new JobRunSummary(recordsProcessed, startTime, finishTime);
     }