fix: update object context diff logic to be shallow rather than deep (#3287)

nidhiii-27 · web-flow · commit 2fd15f69e93a · 2025-09-23T14:19:28.000-04:00
Previously, deleting an object context key by setting its value to null
would fail. This change corrects the diffMaps utility to ensure that
null payloads are always treated as deletions.

The produced update mask is now `contexts.custom.&lt;key&gt;` instead of
`contexts.custom.&lt;key&gt;.value`
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/BlobInfo.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/BlobInfo.java
@@ -33,7 +33,6 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Maps;
 import com.google.common.io.BaseEncoding;
 import java.io.Serializable;
 import java.nio.ByteBuffer;
@@ -1273,18 +1272,15 @@ public Builder setContexts(ObjectContexts contexts) {
       // about the timestamps when determining if a value needs to be patched. Create a new map
       // where we remove the timestamps so equals is usable.
       Map<String, ObjectCustomContextPayload> left =
-          this.contexts == null
-              ? null
-              : ignoreCustomContextPayloadTimestamps(this.contexts.getCustom());
+          this.contexts == null ? null : this.contexts.getCustom();
       Map<String, ObjectCustomContextPayload> right =
-          contexts == null ? null : ignoreCustomContextPayloadTimestamps(contexts.getCustom());
+          contexts == null ? null : contexts.getCustom();
       if (!Objects.equals(left, right)) {
         if (right != null) {
           diffMaps(
               NamedField.nested(BlobField.OBJECT_CONTEXTS, NamedField.literal("custom")),
               left,
               right,
-              f -> NamedField.nested(f, NAMED_FIELD_LITERAL_VALUE),
               modifiedFields::add);
           this.contexts = contexts;
         } else {
@@ -1295,20 +1291,6 @@ public Builder setContexts(ObjectContexts contexts) {
       return this;
     }
 
-    private static @Nullable Map<@NonNull String, @Nullable ObjectCustomContextPayload>
-        ignoreCustomContextPayloadTimestamps(
-            @Nullable Map<@NonNull String, @Nullable ObjectCustomContextPayload> orig) {
-      if (orig == null) {
-        return null;
-      }
-      return Maps.transformValues(
-          orig,
-          v ->
-              v == null
-                  ? null
-                  : ObjectCustomContextPayload.newBuilder().setValue(v.getValue()).build());
-    }
-
     @Override
     public BlobInfo build() {
       checkNotNull(blobId);
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/JsonUtils.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/JsonUtils.java
@@ -22,9 +22,6 @@
 import com.google.cloud.storage.UnifiedOpts.NamedField;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.MapDifference;
-import com.google.common.collect.MapDifference.ValueDifference;
-import com.google.common.collect.Maps;
 import com.google.gson.Gson;
 import com.google.gson.GsonBuilder;
 import com.google.gson.JsonArray;
@@ -41,7 +38,6 @@
 import java.util.Set;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
-import java.util.stream.Stream;
 import org.checkerframework.checker.nullness.qual.NonNull;
 
 final class JsonUtils {
@@ -70,7 +66,7 @@ static <T extends GenericJson> T getOutputJsonWithSelectedFields(
             .map(NamedField::getApiaryName)
             .collect(ImmutableSet.toImmutableSet());
     try {
-      // The datamodel of the apiairy json representation doesn't have a common parent for all
+      // The datamodel of the apiary json representation doesn't have a common parent for all
       // field types, rather than writing a significant amount of code to handle all of these types
       // leverage Gson.
       // 1. serialize the object to it's json string
@@ -104,24 +100,25 @@ static <T extends GenericJson> T getOutputJsonWithSelectedFields(
   static @NonNull JsonObject getOutputJson(JsonObject inputJson, Set<String> fieldsInOutput) {
 
     Map<String, String> l = flatten(inputJson);
-    Map<String, String> r = Utils.setToMap(fieldsInOutput, k -> null);
-
-    MapDifference<String, String> diff = Maps.difference(l, r);
 
     // use hashmap so we can have null values
     HashMap<String, String> flat = new HashMap<>();
-    Stream.of(
-            diff.entriesInCommon().entrySet().stream(),
-            diff.entriesOnlyOnRight().entrySet().stream(),
-            // if the key is present in both maps, but has a differing value select the value from
-            // the left side, as that is the value from inputJson
-            Maps.transformValues(diff.entriesDiffering(), ValueDifference::leftValue)
-                .entrySet()
-                .stream())
-        // flatten
-        .flatMap(x -> x)
-        .forEach(e -> flat.put(e.getKey(), e.getValue()));
-
+    for (String fieldToRetain : fieldsInOutput) {
+      boolean keyFound = false;
+      // Check for exact match or prefix match in the flattened source map (l)
+      for (Map.Entry<String, String> sourceEntry : l.entrySet()) {
+        String sourceKey = sourceEntry.getKey();
+        if (sourceKey.equals(fieldToRetain) || sourceKey.startsWith(fieldToRetain + ".")) {
+          flat.put(sourceKey, sourceEntry.getValue());
+          keyFound = true;
+        }
+      }
+      // If the field to retain wasn't found in the source, it means we need to add it
+      // to the output with a null value, signaling a deletion.
+      if (!keyFound) {
+        flat.put(fieldToRetain, null);
+      }
+    }
     return treeify(flat);
   }
 
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/Utils.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/Utils.java
@@ -29,7 +29,6 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.MapDifference;
-import com.google.common.collect.MapDifference.ValueDifference;
 import com.google.common.collect.Maps;
 import com.google.common.io.BaseEncoding;
 import com.google.common.primitives.Ints;
@@ -282,57 +281,28 @@ static <T> T firstNonNull(Supplier<@Nullable T>... ss) {
    */
   static void diffMaps(
       NamedField parent, Map<String, ?> left, Map<String, ?> right, Consumer<NamedField> sink) {
-    diffMaps(parent, left, right, Function.identity(), sink);
-  }
-
-  /**
-   * Diff two maps, and append each differing key to {@code sink} with the parent of {{@code
-   * parent}. Conditionally apply {@code dec} if deeper qualification is necessary.
-   */
-  static void diffMaps(
-      NamedField parent,
-      Map<String, ?> left,
-      Map<String, ?> right,
-      Function<NamedField, NamedField> dec,
-      Consumer<NamedField> sink) {
-    final Stream<NamedField> keys;
+    final Stream<String> keys;
     if (left != null && right == null) {
-      keys = left.keySet().stream().map(NamedField::literal);
+      keys = left.keySet().stream();
     } else if (left == null && right != null) {
-      keys = right.keySet().stream().map(NamedField::literal).map(dec);
+      keys = right.keySet().stream();
     } else if (left != null && right != null) {
       MapDifference<String, ?> difference = Maps.difference(left, right);
       keys =
           Stream.of(
                   // keys with modified values
-                  difference.entriesDiffering().entrySet().stream()
-                      .map(
-                          e -> {
-                            String key = e.getKey();
-                            NamedField literal = NamedField.literal(key);
-                            ValueDifference<?> diff = e.getValue();
-
-                            if (diff.leftValue() != null && diff.rightValue() == null) {
-                              return literal;
-                            } else if (diff.leftValue() == null && diff.rightValue() != null) {
-                              return literal;
-                            } else {
-                              return dec.apply(literal);
-                            }
-                          }),
+                  difference.entriesDiffering().keySet().stream(),
                   // Only include keys to remove if ALL keys were removed
                   right.isEmpty()
-                      ? difference.entriesOnlyOnLeft().keySet().stream().map(NamedField::literal)
-                      : Stream.<NamedField>empty(),
+                      ? difference.entriesOnlyOnLeft().keySet().stream()
+                      : Stream.<String>empty(),
                   // new keys
-                  difference.entriesOnlyOnRight().keySet().stream()
-                      .map(NamedField::literal)
-                      .map(dec))
+                  difference.entriesOnlyOnRight().keySet().stream())
               .flatMap(x -> x);
     } else {
       keys = Stream.empty();
     }
-    keys.map(k -> NamedField.nested(parent, k)).forEach(sink);
+    keys.map(NamedField::literal).map(k -> NamedField.nested(parent, k)).forEach(sink);
   }
 
   static <T> T[] subArray(T[] ts, int offset, int length) {
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/BlobInfoTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/BlobInfoTest.java
@@ -19,6 +19,8 @@
 import static com.google.cloud.storage.Acl.Project.ProjectRole.VIEWERS;
 import static com.google.cloud.storage.Acl.Role.READER;
 import static com.google.cloud.storage.Acl.Role.WRITER;
+import static com.google.cloud.storage.TestUtils.hashMapOf;
+import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNull;
@@ -30,12 +32,16 @@
 import com.google.cloud.storage.BlobInfo.CustomerEncryption;
 import com.google.cloud.storage.BlobInfo.ObjectContexts;
 import com.google.cloud.storage.BlobInfo.ObjectCustomContextPayload;
+import com.google.cloud.storage.UnifiedOpts.NamedField;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.ImmutableSet;
 import java.math.BigInteger;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
 import org.junit.Test;
 
 public class BlobInfoTest {
@@ -352,4 +358,49 @@ public void testToPbAndFromPb() {
   public void testBlobId() {
     assertEquals(BlobId.of("b", "n", GENERATION), BLOB_INFO.getBlobId());
   }
+
+  @Test
+  public void deepFieldDiffDetectionWorksCorrectly_mutateRetrievedObject() {
+    BlobInfo info =
+        BlobInfo.newBuilder("bucket", "object")
+            .setContexts(
+                ObjectContexts.newBuilder()
+                    .setCustom(
+                        hashMapOf(
+                            "c1", ObjectCustomContextPayload.newBuilder().setValue("C1").build(),
+                            "c2", ObjectCustomContextPayload.newBuilder().setValue("C2").build()))
+                    .build())
+            .setMetadata(
+                hashMapOf(
+                    "m1", "M1",
+                    "m2", "M2"))
+            .build();
+
+    BlobInfo modified =
+        info.toBuilder()
+            .setMetadata(hashMapOf("m2", null))
+            .setContexts(ObjectContexts.newBuilder().setCustom(hashMapOf("k2", null)).build())
+            .build();
+    Set<String> modifiedFields =
+        modified.getModifiedFields().stream()
+            .map(NamedField::getGrpcName)
+            .collect(Collectors.toSet());
+
+    assertThat(modifiedFields).isEqualTo(ImmutableSet.of("contexts.custom.k2", "metadata.m2"));
+  }
+
+  @Test
+  public void deepFieldDiffDetectionWorksCorrectly_declaredDiff() {
+    BlobInfo modified =
+        BlobInfo.newBuilder("bucket", "object")
+            .setMetadata(hashMapOf("m2", null))
+            .setContexts(ObjectContexts.newBuilder().setCustom(hashMapOf("k2", null)).build())
+            .build();
+    Set<String> modifiedFields =
+        modified.getModifiedFields().stream()
+            .map(UnifiedOpts.NamedField::getGrpcName)
+            .collect(Collectors.toSet());
+
+    assertThat(modifiedFields).isEqualTo(ImmutableSet.of("contexts.custom.k2", "metadata.m2"));
+  }
 }
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/TestUtils.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/TestUtils.java
@@ -279,20 +279,19 @@ public static void assertAll(ThrowingRunnable... trs) throws Exception {
   }
 
   /** ImmutableMap does not allow null values, this method does */
-  public static Map<@NonNull String, @Nullable String> hashMapOf(
-      @NonNull String k1, @Nullable String v1) {
+  public static <K, V> Map<@NonNull K, @Nullable V> hashMapOf(@NonNull K k1, @Nullable V v1) {
     requireNonNull(k1, "k1 must be non null");
-    HashMap<String, String> map = new HashMap<>();
+    HashMap<K, V> map = new HashMap<>();
     map.put(k1, v1);
     return Collections.unmodifiableMap(map);
   }
 
   /** ImmutableMap does not allow null values, this method does */
-  public static Map<@NonNull String, @Nullable String> hashMapOf(
-      @NonNull String k1, @Nullable String v1, @NonNull String k2, @Nullable String v2) {
+  public static <K, V> Map<@NonNull K, @Nullable V> hashMapOf(
+      @NonNull K k1, @Nullable V v1, @NonNull K k2, @Nullable V v2) {
     requireNonNull(k1, "k1 must be non null");
     requireNonNull(k2, "k2 must be non null");
-    HashMap<String, String> map = new HashMap<>();
+    HashMap<K, V> map = new HashMap<>();
     map.put(k1, v1);
     map.put(k2, v2);
     return Collections.unmodifiableMap(map);
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITNestedUpdateMaskTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITNestedUpdateMaskTest.java
@@ -78,6 +78,7 @@ public static final class NestedUpdateMaskParametersProvider implements Paramete
     private static final Map<String, String> k1a_k2null = hashMapOf("k1", "a", "k2", null);
     private static final Map<String, String> k1null = hashMapOf("k1", null);
     private static final Map<String, String> k2null = hashMapOf("k2", null);
+    private static final Map<String, String> k1null_k2null = hashMapOf("k1", null, "k2", null);
 
     /**
      *
@@ -95,7 +96,7 @@ public static final class NestedUpdateMaskParametersProvider implements Paramete
      * | {"k1":"a","k2":"b"} | {"k2":null}          | {"k1":"a"}          |
      * | {"k1":"a"}          | {}                   | null                |
      * | {"k1":"a"}          | {"k1":null}          | null                |
-     * | {"k1":"a","k2":"b"} | null                 | null                |
+     * | {"k1":"a","k2":"b"} | {"k1:null,"k2":null} | null                |
      * </pre>
      */
     @Override
@@ -111,7 +112,7 @@ public ImmutableList<Param> parameters() {
           new Param("2 keys, modify 1 null (fine)", k1a_k2b, k2null, k1a),
           new Param("1 key, set empty", k1a, empty, null),
           new Param("1 key, null key", k1a, k1null, null),
-          new Param("2 keys, set null", k1a_k2b, null, null));
+          new Param("2 keys, set null", k1a_k2b, k1null_k2null, null));
     }
   }
 
