chore(spanner): update logic

googleapis · Oct 23, 2024 · 5907ad5 · 5907ad5
1 parent 65b3234
commit 5907ad5
Showing 1 changed file with 24 additions and 13 deletions.
diff --git a/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Mutation.java b/google-cloud-spanner/src/main/java/com/google/cloud/spanner/Mutation.java
@@ -433,7 +433,12 @@ static com.google.spanner.v1.Mutation toProtoGetRandomMutation(
           if (proto != null) {
             com.google.spanner.v1.Mutation builtMutation = proto.build();
             out.add(builtMutation);
-            allMutationsExcludingInsert.add(builtMutation);
+            if (checkIfInsertMutationWithLargeValue(builtMutation, largeInsertMutation)) {
+              largeInsertMutation = builtMutation;
+            }
+            if (!builtMutation.hasInsert()) {
+              allMutationsExcludingInsert.add(builtMutation);
+            }
           }
           proto = com.google.spanner.v1.Mutation.newBuilder();
           com.google.spanner.v1.Mutation.Delete.Builder delete =
@@ -457,12 +462,10 @@ static com.google.spanner.v1.Mutation toProtoGetRandomMutation(
           if (proto != null) {
             com.google.spanner.v1.Mutation builtMutation = proto.build();
             out.add(builtMutation);
-            if (builtMutation.hasInsert()) {
-              if (builtMutation.getInsert().getValuesCount()
-                  > largeInsertMutation.getInsert().getValuesCount()) {
-                largeInsertMutation = builtMutation;
-              }
-            } else {
+            if (checkIfInsertMutationWithLargeValue(builtMutation, largeInsertMutation)) {
+              largeInsertMutation = builtMutation;
+            }
+            if (!builtMutation.hasInsert()) {
               allMutationsExcludingInsert.add(builtMutation);
             }
           }
@@ -493,12 +496,10 @@ static com.google.spanner.v1.Mutation toProtoGetRandomMutation(
     if (proto != null) {
       com.google.spanner.v1.Mutation builtMutation = proto.build();
       out.add(proto.build());
-      if (builtMutation.hasInsert()) {
-        if (builtMutation.getInsert().getValuesCount()
-            > largeInsertMutation.getInsert().getValuesCount()) {
-          largeInsertMutation = builtMutation;
-        }
-      } else {
+      if (checkIfInsertMutationWithLargeValue(builtMutation, largeInsertMutation)) {
+        largeInsertMutation = builtMutation;
+      }
+      if (!builtMutation.hasInsert()) {
         allMutationsExcludingInsert.add(builtMutation);
       }
     }
@@ -514,4 +515,14 @@ static com.google.spanner.v1.Mutation toProtoGetRandomMutation(
       return largeInsertMutation;
     }
   }
+
+  // Returns true if the input mutation is of type INSERT and has more values than the current
+  // largest insert mutation.
+  private static boolean checkIfInsertMutationWithLargeValue(
+      com.google.spanner.v1.Mutation mutation,
+      com.google.spanner.v1.Mutation largestInsertMutation) {
+    return mutation.hasInsert()
+        && mutation.getInsert().getValuesCount()
+            > largestInsertMutation.getInsert().getValuesCount();
+  }
 }