Changes default behavior for DynamoDb Enhanced atomic counter extension

.changes/next-release/bugfix-AWSSDKforJavav2DynamoDbEnhanced-05cdcee.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "AWS SDK for Java v2 - DynamoDb Enhanced",
+    "contributor": "",
+    "description": "Changes the default behavior of the DynamoDb Enhanced atomic counter extension to automatically filter out any counter attributes in the item to be updated. This allows users to read and update items without DynamoDb collision errors."
+}

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/extensions/AtomicCounterExtension.java

@@ -19,8 +19,10 @@
 import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.valueRef;
 import static software.amazon.awssdk.enhanced.dynamodb.internal.update.UpdateExpressionUtils.ifNotExists;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.stream.Collectors;
 import software.amazon.awssdk.annotations.SdkPublicApi;
@@ -37,10 +39,11 @@
 import software.amazon.awssdk.enhanced.dynamodb.update.UpdateExpression;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.utils.CollectionUtils;
+import software.amazon.awssdk.utils.Logger;
 
 /**
- * This extension enables atomic counter attributes to be written to the database.
- * The extension is loaded by default when you instantiate a
+ * This extension enables atomic counter attributes to be changed in DynamoDb by creating instructions for modifying
+ * an existing value or setting a start value. The extension is loaded by default when you instantiate a
  * {@link DynamoDbEnhancedClient} and only needs to be added to the client if you
  * are adding custom extensions to the client.
  * <p>
@@ -56,8 +59,7 @@
  * <p>
  * Every time a new update of the record is successfully written to the database, the counter will be updated automatically.
  * By default, the counter starts at 0 and increments by 1 for each update. The tags provide the capability of adjusting
- * the counter start and increment/decrement values such as described in
- * {@link DynamoDbAtomicCounter}.
+ * the counter start and increment/decrement values such as described in {@link DynamoDbAtomicCounter}.
  * <p>
  * Example 1: Using a bean based table schema
  * <pre>
@@ -86,10 +88,17 @@
  * }
  * </pre>
  * <p>
- * <b>NOTE: </b>When using putItem, the counter will be reset to its start value.
+ * <b>NOTES: </b>
+ * <ul>
+ *     <li>When using putItem, the counter will be reset to its start value.</li>
+ *     <li>The extension will remove any existing occurrences of the atomic counter attributes.</li>
+ * </ul>
  */
 @SdkPublicApi
 public final class AtomicCounterExtension implements DynamoDbEnhancedClientExtension {
+
+    private static final Logger log = Logger.loggerFor(AtomicCounterExtension.class);
+
     private AtomicCounterExtension() {
     }
 
@@ -118,6 +127,7 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
                 break;
             case UPDATE_ITEM:
                 modificationBuilder.updateExpression(createUpdateExpression(counters));
+                modificationBuilder.transformedItem(filterFromItem(counters, context.items()));
                 break;
             default: break;
         }
@@ -136,6 +146,22 @@ private Map<String, AttributeValue> addToItem(Map<String, AtomicCounter> counter
         return Collections.unmodifiableMap(itemToTransform);
     }
 
+    private Map<String, AttributeValue> filterFromItem(Map<String, AtomicCounter> counters, Map<String, AttributeValue> items) {
+        Map<String, AttributeValue> itemToTransform = new HashMap<>(items);
+        List<String> removedAttributes = new ArrayList<>();
+        for (String attributeName : counters.keySet()) {
+            if (itemToTransform.containsKey(attributeName)) {
+                itemToTransform.remove(attributeName);
+                removedAttributes.add(attributeName);
+            }
+        }
+        if (!removedAttributes.isEmpty()) {
+            log.debug(() -> String.format("Filtered atomic counter attributes from existing update item to avoid collisions: %s",
+                                          String.join(",", removedAttributes)));
+        }
+        return Collections.unmodifiableMap(itemToTransform);
+    }
+
     private SetAction counterAction(Map.Entry<String, AtomicCounter> e) {
         String attributeName = e.getKey();
         AtomicCounter counter = e.getValue();

services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/extensions/AtomicCounterExtensionTest.java

@@ -87,7 +87,10 @@ public void beforeWrite_updateItemOperation_hasCounters_createsUpdateExpression(
                                                                               .operationName(OperationName.UPDATE_ITEM)
                                                                               .operationContext(PRIMARY_CONTEXT).build());
 
-        assertThat(result.transformedItem()).isNull();
+        Map<String, AttributeValue> transformedItem = result.transformedItem();
+        assertThat(transformedItem).isNotNull().hasSize(1);
+        assertThat(transformedItem).containsEntry("id", AttributeValue.fromS(RECORD_ID));
+
         assertThat(result.updateExpression()).isNotNull();
 
         List<SetAction> setActions = result.updateExpression().setActions();
@@ -112,11 +115,39 @@ public void beforeWrite_updateItemOperation_noCounters_noChanges() {
                                                                               .tableMetadata(SIMPLE_ITEM_MAPPER.tableMetadata())
                                                                               .operationName(OperationName.UPDATE_ITEM)
                                                                               .operationContext(PRIMARY_CONTEXT).build());
-
         assertThat(result.transformedItem()).isNull();
         assertThat(result.updateExpression()).isNull();
     }
 
+    @Test
+    public void beforeWrite_updateItemOperation_hasCountersInItem_createsUpdateExpressionAndFilters() {
+        AtomicCounterItem atomicCounterItem = new AtomicCounterItem();
+        atomicCounterItem.setId(RECORD_ID);
+        atomicCounterItem.setCustomCounter(255L);
+
+        Map<String, AttributeValue> items = ITEM_MAPPER.itemToMap(atomicCounterItem, true);
+        assertThat(items).hasSize(2);
+
+        WriteModification result =
+            atomicCounterExtension.beforeWrite(DefaultDynamoDbExtensionContext.builder()
+                                                                              .items(items)
+                                                                              .tableMetadata(ITEM_MAPPER.tableMetadata())
+                                                                              .operationName(OperationName.UPDATE_ITEM)
+                                                                              .operationContext(PRIMARY_CONTEXT).build());
+
+        Map<String, AttributeValue> transformedItem = result.transformedItem();
+        assertThat(transformedItem).isNotNull().hasSize(1);
+        assertThat(transformedItem).containsEntry("id", AttributeValue.fromS(RECORD_ID));
+
+        assertThat(result.updateExpression()).isNotNull();
+
+        List<SetAction> setActions = result.updateExpression().setActions();
+        assertThat(setActions).hasSize(2);
+
+        verifyAction(setActions, "customCounter", "5", "5");
+        verifyAction(setActions, "defaultCounter", "-1", "1");
+    }
+
     @Test
     public void beforeWrite_putItemOperation_hasCounters_createsItemTransform() {
         AtomicCounterItem atomicCounterItem = new AtomicCounterItem();

services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/AtomicCounterTest.java

@@ -112,15 +112,31 @@ public void createViaPut_incrementsCorrectly() {
     }
 
     @Test
-    public void createViaUpdate_settingCounterInPojo_throwsException() {
+    public void createViaUpdate_settingCounterInPojo_hasNoEffect() {
         AtomicCounterRecord record = new AtomicCounterRecord();
         record.setId(RECORD_ID);
         record.setDefaultCounter(10L);
         record.setAttribute1(STRING_VALUE);
 
-        assertThatThrownBy(() -> mappedTable.updateItem(record))
-            .isInstanceOf(DynamoDbException.class)
-            .hasMessageContaining("Two document paths");
+        mappedTable.updateItem(record);
+        AtomicCounterRecord persistedRecord = mappedTable.getItem(record);
+        assertThat(persistedRecord.getAttribute1()).isEqualTo(STRING_VALUE);
+        assertThat(persistedRecord.getDefaultCounter()).isEqualTo(0L);
+        assertThat(persistedRecord.getCustomCounter()).isEqualTo(10L);
+        assertThat(persistedRecord.getDecreasingCounter()).isEqualTo(-20L);
+    }
+
+    @Test
+    public void updateItem_retrievedFromDb_shouldNotThrowException() {
+        AtomicCounterRecord record = new AtomicCounterRecord();
+        record.setId(RECORD_ID);
+        record.setAttribute1(STRING_VALUE);
+        mappedTable.updateItem(record);
+
+        AtomicCounterRecord retrievedRecord = mappedTable.getItem(record);
+        retrievedRecord.setAttribute1("ChangingThisAttribute");
+
+        assertThat(mappedTable.updateItem(retrievedRecord)).isNotNull();
     }
 
     @Test