Add allowDuplicateKeys to ImmutableMap Builder

Author: davidh44

.changes/next-release/bugfix-AWSSDKforJavav2-01696b7.json

@@ -1,6 +1,6 @@
 {
-    "type": "bugfix",
+    "type": "feature",
     "category": "AWS SDK for Java v2",
     "contributor": "",
-    "description": "Log error instead of throwing error for duplicate keys in ImmutableMap"
+    "description": "Add allowDuplicateKeys to ImmutableMap Builder"
 }

codegen-lite/src/main/java/software/amazon/awssdk/codegen/lite/regions/ServiceMetadataGenerator.java

@@ -257,7 +257,7 @@ private CodeBlock dnsSuffixesByPartition(Partitions partitions) {
     private CodeBlock hostnamesByRegion(Partitions partitions) {
         Map<Partition, Service> services = getServiceData(partitions);
 
-        CodeBlock.Builder builder = CodeBlock.builder().add("$T.<$T, $T>builder()",
+        CodeBlock.Builder builder = CodeBlock.builder().add("$T.<$T, $T>builder().allowDuplicateKeys(true)",
                                                             ImmutableMap.class, serviceEndpointKeyClass(), String.class);
 
         services.forEach((partition, service) -> {

codegen-lite/src/test/resources/software/amazon/awssdk/codegen/lite/regions/s3-service-metadata.java

@@ -69,6 +69,7 @@ public final class S3ServiceMetadata implements ServiceMetadata {
 
     private static final Map<ServiceEndpointKey, String> HOSTNAMES_BY_REGION = ImmutableMap
         .<ServiceEndpointKey, String> builder()
+        .allowDuplicateKeys(true)
         .put(ServiceEndpointKey.builder().region(Region.of("af-south-1")).tags(EndpointTag.of("dualstack")).build(),
              "s3.dualstack.af-south-1.amazonaws.com")
         .put(ServiceEndpointKey.builder().region(Region.of("ap-east-1")).tags(EndpointTag.of("dualstack")).build(),

codegen-lite/src/test/resources/software/amazon/awssdk/codegen/lite/regions/sts-service-metadata.java

@@ -60,6 +60,7 @@ public final class StsServiceMetadata implements ServiceMetadata {
 
     private static final Map<ServiceEndpointKey, String> HOSTNAMES_BY_REGION = ImmutableMap
         .<ServiceEndpointKey, String> builder()
+        .allowDuplicateKeys(true)
         .put(ServiceEndpointKey.builder().region(Region.of("aws-global")).build(), "sts.amazonaws.com")
         .put(ServiceEndpointKey.builder().region(Region.of("us-east-1")).tags(EndpointTag.of("fips")).build(),
              "sts-fips.us-east-1.amazonaws.com")

utils/src/main/java/software/amazon/awssdk/utils/ImmutableMap.java

@@ -60,6 +60,7 @@ public final class ImmutableMap<K, V> implements Map<K, V> {
 
     private static final Logger log = Logger.loggerFor(ImmutableMap.class);
     private static final String UNMODIFIABLE_MESSAGE = "This is an immutable map.";
+    private static final String DUPLICATED_KEY_MESSAGE = "Duplicate keys are provided.";
 
     private final Map<K, V> map;
 
@@ -198,10 +199,18 @@ public static <K, V> ImmutableMap<K, V> of(K k0, V v0, K k1, V v1,
 
     private static <K, V> void putAndWarnDuplicateKeys(Map<K, V> map, K key,
                                                        V value) {
+        putAndWarnDuplicateKeys(map, key, value, false);
+    }
+
+    private static <K, V> void putAndWarnDuplicateKeys(Map<K, V> map, K key, V value, boolean allowDuplicateKeys) {
         if (map.containsKey(key)) {
-            log.error(() -> String.format("Duplicate keys are provided for [%s]. The first value [%s] will be kept, and the "
-                                          + "second value [%s] will be ignored.", key, map.get(key), value));
-            return;
+
+            if (allowDuplicateKeys) {
+                log.warn(() -> String.format("Duplicate keys are provided for [%s]. The newer value [%s] will be saved, and the "
+                                             + "existing value [%s] will be overwritten.", key, value, map.get(key)));
+            } else {
+                throw new IllegalArgumentException(DUPLICATED_KEY_MESSAGE);
+            }
         }
         map.put(key, value);
     }
@@ -291,20 +300,38 @@ public String toString() {
     public static class Builder<K, V> {
 
         private final Map<K, V> entries;
+        private boolean allowDuplicateKeys;
 
         public Builder() {
             this.entries = new HashMap<>();
         }
 
         /**
-         * Add a key-value pair into the built map. If duplicate keys are provided, the first value will be kept and the second
-         * value will be ignored, and the SDK will log an error.
+         * Add a key-value pair into the built map. The method will throw
+         * IllegalArgumentException immediately when duplicate keys are
+         * provided and allowDuplicateKeys is false (default value).
+         *
+         *
+         * If duplicate keys are provided and allowDuplicateKeys is true, the latest value will overwrite the existing value, and
+         * the SDK will log a message at WARN level.
          *
          * @return Returns a reference to this object so that method calls can
          *         be chained together.
          */
         public Builder<K, V> put(K key, V value) {
-            putAndWarnDuplicateKeys(entries, key, value);
+            putAndWarnDuplicateKeys(entries, key, value, allowDuplicateKeys);
+            return this;
+        }
+
+        /**
+         * Sets whether duplicate keys are allowed. If true, the latest value will overwrite the existing value. If
+         * false, an error will be thrown when attempting to put an entry with a duplicate key. Default value is false.
+         *
+         * @return Returns a reference to this object so that method calls can
+         *         be chained together.
+         */
+        public Builder<K, V> allowDuplicateKeys(boolean allowDuplicateKeys) {
+            this.allowDuplicateKeys = allowDuplicateKeys;
             return this;
         }
 

utils/src/test/java/software/amazon/awssdk/utils/ImmutableMapTest.java

@@ -71,13 +71,29 @@ public void testOfBuilder() {
     }
 
     @Test
-    public void putDuplicateKeys_keepsFirstOneAndDoesNotThrowError() {
+    public void testErrorOnDuplicateKeys() {
+        try {
+            Map<Integer, String> builtMap = new ImmutableMap.Builder<Integer, String>()
+                .put(1, "one")
+                .put(1, "two")
+                .build();
+            fail("IllegalArgumentException expected.");
+        } catch (IllegalArgumentException iae) {
+            // Ignored or expected.
+        } catch (Exception e) {
+            fail("IllegalArgumentException expected.");
+        }
+    }
+
+    @Test
+    public void putDuplicateKeys_allowDuplicateKeysTrue_doesNotThrowErrorAndKeepsNewestValue() {
         Map<Integer, String> builtMap = new ImmutableMap.Builder<Integer, String>()
+            .allowDuplicateKeys(true)
             .put(1, "one")
             .put(1, "two")
             .build();
 
-        assertEquals("one", builtMap.get(1));
+        assertEquals("two", builtMap.get(1));
     }
 
     @Test