Adding unprocessedKeysForTable to BatchGetItem results object

.changes/next-release/feature-DynamoDBEnhancedClient-ddec09e.json

@@ -0,0 +1,6 @@
+{
+    "category": "DynamoDB Enhanced Client", 
+    "contributor": "", 
+    "type": "feature", 
+    "description": "Added method to BatchGetItem results to retrieve unprocessed keys for a given table."
+}

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetResultPage.java

@@ -18,16 +18,20 @@
 import static java.util.Collections.emptyList;
 import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.readAndTransformSingleItem;
 
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.stream.Collectors;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
+import software.amazon.awssdk.enhanced.dynamodb.Key;
 import software.amazon.awssdk.enhanced.dynamodb.MappedTableResource;
 import software.amazon.awssdk.enhanced.dynamodb.internal.operations.DefaultOperationContext;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
 import software.amazon.awssdk.services.dynamodb.model.BatchGetItemResponse;
+import software.amazon.awssdk.services.dynamodb.model.KeysAndAttributes;
 
 /**
  * Defines one result page with retrieved items in the result of a batchGetItem() operation, such as
@@ -74,6 +78,36 @@ public <T> List<T> resultsForTable(MappedTableResource<T> mappedTable) {
                       .collect(Collectors.toList());
     }
 
+    /**
+     * Returns a list of keys associated with a given table that were not processed during the operation, typically
+     * because the total size of the request is too large or exceeds the provisioned throughput of the table. If an item
+     * was attempted to be retrieved but not found in the table, it will not appear in this list or the results list.
+     *
+     * @param mappedTable the table to retrieve the unprocessed keys for
+     * @return a list of unprocessed keys
+     */
+    public List<Key> unprocessedKeysForTable(MappedTableResource<?> mappedTable) {
+        KeysAndAttributes keysAndAttributes = this.batchGetItemResponse.unprocessedKeys().get(mappedTable.tableName());
+
+        if (keysAndAttributes == null) {
+            return Collections.emptyList();
+        }
+
+        String partitionKey = mappedTable.tableSchema().tableMetadata().primaryPartitionKey();
+        Optional<String> sortKey = mappedTable.tableSchema().tableMetadata().primarySortKey();
+
+        return keysAndAttributes.keys()
+                                .stream()
+                                .map(keyMap -> {
+                                    AttributeValue partitionValue = keyMap.get(partitionKey);
+                                    AttributeValue sortValue = sortKey.map(keyMap::get).orElse(null);
+                                    return Key.builder()
+                                              .partitionValue(partitionValue)
+                                              .sortValue(sortValue)
+                                              .build(); })
+                                .collect(Collectors.toList());
+    }
+
     /**
      * A builder that is used to create a result object with the desired parameters.
      */

services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchGetItemOperationTest.java

@@ -22,6 +22,7 @@
 import static java.util.stream.Collectors.toList;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsInAnyOrder;
+import static org.hamcrest.Matchers.empty;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.nullValue;
 import static org.hamcrest.Matchers.sameInstance;
@@ -36,6 +37,7 @@
 import static software.amazon.awssdk.enhanced.dynamodb.functionaltests.models.FakeItemWithSort.createUniqueFakeItemWithSort;
 
 import java.util.Arrays;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -94,13 +96,13 @@ public class BatchGetItemOperationTest {
     @Mock
     private DynamoDbEnhancedClientExtension mockExtension;
 
-    private DynamoDbEnhancedClient enhancedClient;
     private DynamoDbTable<FakeItem> fakeItemMappedTable;
     private DynamoDbTable<FakeItemWithSort> fakeItemWithSortMappedTable;
 
     @Before
     public void setupMappedTables() {
-        enhancedClient = DynamoDbEnhancedClient.builder().dynamoDbClient(mockDynamoDbClient).extensions().build();
+        DynamoDbEnhancedClient enhancedClient =
+                DynamoDbEnhancedClient.builder().dynamoDbClient(mockDynamoDbClient).extensions().build();
         fakeItemMappedTable = enhancedClient.table(TABLE_NAME, FakeItem.getTableSchema());
         fakeItemWithSortMappedTable = enhancedClient.table(TABLE_NAME_2, FakeItemWithSort.getTableSchema());
     }
@@ -296,6 +298,75 @@ public void transformResponse_multipleTables_multipleItems_noExtension() {
         assertThat(fakeItemWithSortResultsPage, containsInAnyOrder(FAKESORT_ITEMS.get(0)));
     }
 
+    @Test
+    public void transformResponse_multipleTables_multipleItems_unprocessedKeys() {
+        KeysAndAttributes keysAndAttributes1 =
+                KeysAndAttributes.builder()
+                                 .keys(Arrays.asList(FAKE_ITEM_MAPS.get(0), FAKE_ITEM_MAPS.get(1)))
+                                 .build();
+
+        KeysAndAttributes keysAndAttributes2 =
+                KeysAndAttributes.builder()
+                                 .keys(Collections.singletonList(FAKESORT_ITEM_MAPS.get(0)))
+                                 .build();
+
+        Map<String, KeysAndAttributes> keysAndAttributesMap = new HashMap<>();
+        keysAndAttributesMap.put(TABLE_NAME, keysAndAttributes1);
+        keysAndAttributesMap.put(TABLE_NAME_2, keysAndAttributes2);
+
+        BatchGetItemResponse fakeResponse =
+                BatchGetItemResponse.builder()
+                                    .unprocessedKeys(keysAndAttributesMap)
+                                    .build();
+
+        BatchGetItemOperation operation = BatchGetItemOperation.create(emptyRequest());
+
+        BatchGetResultPage resultsPage = operation.transformResponse(fakeResponse, null);
+
+        List<Key> fakeItemResults1 = resultsPage.unprocessedKeysForTable(fakeItemMappedTable);
+        List<Key> fakeItemResults2 = resultsPage.unprocessedKeysForTable(fakeItemWithSortMappedTable);
+        assertThat(fakeItemResults1, containsInAnyOrder(FAKE_ITEM_KEYS.get(0), FAKE_ITEM_KEYS.get(1)));
+        assertThat(fakeItemResults2, containsInAnyOrder(FAKESORT_ITEM_KEYS.get(0)));
+    }
+
+    @Test
+    public void transformResponse_multipleTables_multipleItems_no_unprocessedKeys() {
+        BatchGetItemResponse fakeResponse = BatchGetItemResponse.builder().build();
+
+        BatchGetItemOperation operation = BatchGetItemOperation.create(emptyRequest());
+
+        BatchGetResultPage resultsPage = operation.transformResponse(fakeResponse, null);
+
+        List<Key> fakeItemResults1 = resultsPage.unprocessedKeysForTable(fakeItemMappedTable);
+        List<Key> fakeItemResults2 = resultsPage.unprocessedKeysForTable(fakeItemWithSortMappedTable);
+        assertThat(fakeItemResults1, empty());
+        assertThat(fakeItemResults2, empty());
+    }
+
+    @Test
+    public void transformResponse_multipleTables_multipleItems_unprocessedKeys_tableNotExists() {
+        KeysAndAttributes keysAndAttributes1 =
+                KeysAndAttributes.builder()
+                                 .keys(Arrays.asList(FAKE_ITEM_MAPS.get(0), FAKE_ITEM_MAPS.get(1)))
+                                 .build();
+
+
+        Map<String, KeysAndAttributes> keysAndAttributesMap = new HashMap<>();
+        keysAndAttributesMap.put(TABLE_NAME, keysAndAttributes1);
+
+        BatchGetItemResponse fakeResponse =
+                BatchGetItemResponse.builder()
+                                    .unprocessedKeys(keysAndAttributesMap)
+                                    .build();
+
+        BatchGetItemOperation operation = BatchGetItemOperation.create(emptyRequest());
+
+        BatchGetResultPage resultsPage = operation.transformResponse(fakeResponse, null);
+
+        List<Key> fakeItemResults = resultsPage.unprocessedKeysForTable(fakeItemWithSortMappedTable);
+        assertThat(fakeItemResults, empty());
+    }
+
     @Test
     public void transformResponse_multipleTables_multipleItems_extensionWithTransformation() {
         Map<String, List<Map<String, AttributeValue>>> page = new HashMap<>();