Address RollbackMutationComposer as well

Author: komamitsu

core/src/main/java/com/scalar/db/transaction/consensuscommit/RollbackMutationComposer.java

@@ -8,8 +8,10 @@
 import com.scalar.db.api.ConditionBuilder;
 import com.scalar.db.api.Consistency;
 import com.scalar.db.api.Delete;
+import com.scalar.db.api.DeleteBuilder;
 import com.scalar.db.api.DistributedStorage;
 import com.scalar.db.api.Get;
+import com.scalar.db.api.GetBuilder;
 import com.scalar.db.api.Mutation;
 import com.scalar.db.api.Operation;
 import com.scalar.db.api.Put;
@@ -121,17 +123,19 @@ private Delete composeDelete(Operation base, TransactionResult result) throws Ex
     Key partitionKey = ScalarDbUtils.getPartitionKey(result, tableMetadata);
     Optional<Key> clusteringKey = ScalarDbUtils.getClusteringKey(result, tableMetadata);
 
-    return Delete.newBuilder()
-        .namespace(base.forNamespace().get())
-        .table(base.forTable().get())
-        .partitionKey(partitionKey)
-        .clusteringKey(clusteringKey.orElse(null))
-        .condition(
-            ConditionBuilder.deleteIf(ConditionBuilder.column(ID).isEqualToText(id))
-                .and(ConditionBuilder.column(STATE).isEqualToInt(result.getState().get()))
-                .build())
-        .consistency(Consistency.LINEARIZABLE)
-        .build();
+    DeleteBuilder.Buildable deleteBuilder =
+        Delete.newBuilder()
+            .namespace(base.forNamespace().get())
+            .table(base.forTable().get())
+            .partitionKey(partitionKey)
+            .condition(
+                ConditionBuilder.deleteIf(ConditionBuilder.column(ID).isEqualToText(id))
+                    .and(ConditionBuilder.column(STATE).isEqualToInt(result.getState().get()))
+                    .build())
+            .consistency(Consistency.LINEARIZABLE);
+    clusteringKey.ifPresent(deleteBuilder::clusteringKey);
+
+    return deleteBuilder.build();
   }
 
   private Optional<TransactionResult> getLatestResult(
@@ -157,15 +161,16 @@ private Optional<TransactionResult> getLatestResult(
       }
     }
 
-    Get get =
+    GetBuilder.BuildableGetWithPartitionKey getBuilder =
         Get.newBuilder()
             .namespace(operation.forNamespace().get())
             .table(operation.forTable().get())
             .partitionKey(partitionKey)
-            .clusteringKey(clusteringKey)
-            .consistency(Consistency.LINEARIZABLE)
-            .build();
+            .consistency(Consistency.LINEARIZABLE);
+    if (clusteringKey != null) {
+      getBuilder.clusteringKey(clusteringKey);
+    }
 
-    return storage.get(get).map(TransactionResult::new);
+    return storage.get(getBuilder.build()).map(TransactionResult::new);
   }
 }

core/src/test/java/com/scalar/db/transaction/consensuscommit/RollbackMutationComposerTest.java

@@ -188,6 +188,25 @@ private Put preparePut() {
         .build();
   }
 
+  private Put preparePutWithoutCk() {
+    return Put.newBuilder()
+        .namespace(ANY_NAMESPACE_NAME)
+        .table(ANY_TABLE_NAME)
+        .partitionKey(Key.ofText(ANY_NAME_1, ANY_TEXT_1))
+        .intValue(ANY_NAME_3, ANY_INT_3)
+        .booleanValue(ANY_NAME_4, false)
+        .bigIntValue(ANY_NAME_5, ANY_BIGINT_3)
+        .floatValue(ANY_NAME_6, ANY_FLOAT_3)
+        .doubleValue(ANY_NAME_7, ANY_DOUBLE_3)
+        .textValue(ANY_NAME_8, ANY_TEXT_4)
+        .blobValue(ANY_NAME_9, ANY_BLOB_3)
+        .dateValue(ANY_NAME_10, ANY_DATE_3)
+        .timeValue(ANY_NAME_11, ANY_TIME_3)
+        .timestampValue(ANY_NAME_12, ANY_TIMESTAMP_3)
+        .timestampTZValue(ANY_NAME_13, ANY_TIMESTAMPTZ_3)
+        .build();
+  }
+
   private TransactionResult prepareResult(TransactionState state) {
     ImmutableMap<String, Column<?>> columns =
         ImmutableMap.<String, Column<?>>builder()
@@ -548,6 +567,22 @@ public void add_PutAndNullResultGivenAndOldResultGivenFromStorage_ShouldDoNothin
     verify(storage).get(any(Get.class));
   }
 
+  @Test
+  public void add_PutWithoutCkAndNullResultGivenAndOldResultGivenFromStorage_ShouldDoNothing()
+      throws ExecutionException {
+    // Arrange
+    TransactionResult result = prepareInitialResult(ANY_ID_1, TransactionState.PREPARED);
+    when(storage.get(any(Get.class))).thenReturn(Optional.of(result));
+    Put put = preparePutWithoutCk();
+
+    // Act
+    composer.add(put, null);
+
+    // Assert
+    assertThat(composer.get().size()).isEqualTo(0);
+    verify(storage).get(any(Get.class));
+  }
+
   @Test
   public void add_PutAndNullResultGivenAndEmptyResultGivenFromStorage_ShouldDoNothing()
       throws ExecutionException {