Changes default behavior for DynamoDb Enhanced atomic counter extensi…

…on (#4314)

* Changes the default behavior for DynamoDb Enhanced atomic counter extension to filter out any counter attributes in the item map if present

.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,12 +39,12 @@
 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
- * {@link DynamoDbEnhancedClient} and only needs to be added to the client if you
- * are adding custom extensions to the client.
+ * 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>
  * To utilize atomic counters, first create a field in your model that will be used to store the counter.
  * This class field should of type {@link Long} and you need to tag it as an atomic counter:
@@ -56,8 +58,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 +87,18 @@
  * }
  * </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 from the record during an
+ *     <i>updateItem</i> operation. Manually editing attributes marked as atomic counters will have <b>NO EFFECT</b>.</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,35 @@ 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");
+
+        retrievedRecord = mappedTable.updateItem(retrievedRecord);
+        assertThat(retrievedRecord).isNotNull();
+        assertThat(retrievedRecord.getDefaultCounter()).isEqualTo(1L);
+        assertThat(retrievedRecord.getCustomCounter()).isEqualTo(15L);
+        assertThat(retrievedRecord.getDecreasingCounter()).isEqualTo(-21L);
     }
 
     @Test