diff --git a/.changes/next-release/feature-AmazonDynamoDBEnhancedClient-ee18d70.json b/.changes/next-release/feature-AmazonDynamoDBEnhancedClient-ee18d70.json new file mode 100644 index 000000000000..7a7817222744 --- /dev/null +++ b/.changes/next-release/feature-AmazonDynamoDBEnhancedClient-ee18d70.json @@ -0,0 +1,6 @@ +{ + "type": "feature", + "category": "Amazon DynamoDB Enhanced Client", + "contributor": "psnilesh", + "description": "This commit introduces returnConsumedCapacity input to BatchGetItemEnhancedRequest that allows customers to find out exactly how much read units were consumed by the operation." +} diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchGetItemOperation.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchGetItemOperation.java index 7357016d20cc..541f205f6cae 100644 --- a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchGetItemOperation.java +++ b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchGetItemOperation.java @@ -62,6 +62,7 @@ public BatchGetItemRequest generateRequest(DynamoDbEnhancedClientExtension exten return BatchGetItemRequest.builder() .requestItems(Collections.unmodifiableMap(requestItems)) + .returnConsumedCapacity(request.returnConsumedCapacityAsString()) .build(); } diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetItemEnhancedRequest.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetItemEnhancedRequest.java index 9c360aa9f225..ea2bdc8fd1fc 100644 --- a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetItemEnhancedRequest.java +++ b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetItemEnhancedRequest.java @@ -20,10 +20,14 @@ import java.util.Collection; import java.util.Collections; import java.util.List; +import java.util.Objects; import software.amazon.awssdk.annotations.NotThreadSafe; import software.amazon.awssdk.annotations.SdkPublicApi; import software.amazon.awssdk.annotations.ThreadSafe; import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient; +import software.amazon.awssdk.services.dynamodb.model.BatchGetItemRequest; +import software.amazon.awssdk.services.dynamodb.model.GetItemRequest; +import software.amazon.awssdk.services.dynamodb.model.ReturnConsumedCapacity; /** * Defines parameters used for the batchGetItem() operation (such as @@ -36,9 +40,11 @@ public final class BatchGetItemEnhancedRequest { private final List readBatches; + private final String returnConsumedCapacity; private BatchGetItemEnhancedRequest(Builder builder) { this.readBatches = getListIfExist(builder.readBatches); + this.returnConsumedCapacity = builder.returnConsumedCapacity; } /** @@ -52,7 +58,7 @@ public static Builder builder() { * Returns a builder initialized with all existing values on the request object. */ public Builder toBuilder() { - return new Builder().readBatches(readBatches); + return new Builder().readBatches(readBatches).returnConsumedCapacity(this.returnConsumedCapacity); } /** @@ -62,6 +68,25 @@ public Collection readBatches() { return readBatches; } + /** + * Whether to return the capacity consumed by this operation. + * + * @see GetItemRequest#returnConsumedCapacity() + */ + public ReturnConsumedCapacity returnConsumedCapacity() { + return ReturnConsumedCapacity.fromValue(returnConsumedCapacity); + } + + /** + * Whether to return the capacity consumed by this operation. + *

+ * Similar to {@link #returnConsumedCapacity()} but return the value as a string. This is useful in situations where the + * value is not defined in {@link ReturnConsumedCapacity}. + */ + public String returnConsumedCapacityAsString() { + return returnConsumedCapacity; + } + @Override public boolean equals(Object o) { if (this == o) { @@ -73,12 +98,15 @@ public boolean equals(Object o) { BatchGetItemEnhancedRequest that = (BatchGetItemEnhancedRequest) o; - return readBatches != null ? readBatches.equals(that.readBatches) : that.readBatches == null; + return Objects.equals(this.readBatches, that.readBatches) && + Objects.equals(this.returnConsumedCapacity, that.returnConsumedCapacity); } @Override public int hashCode() { - return readBatches != null ? readBatches.hashCode() : 0; + int hc = readBatches != null ? readBatches.hashCode() : 0; + hc = 31 * hc + (returnConsumedCapacity != null ? returnConsumedCapacity.hashCode() : 0); + return hc; } private static List getListIfExist(List readBatches) { @@ -91,6 +119,7 @@ private static List getListIfExist(List readBatches) { @NotThreadSafe public static final class Builder { private List readBatches; + private String returnConsumedCapacity; private Builder() { } @@ -132,9 +161,30 @@ public Builder addReadBatch(ReadBatch readBatch) { return this; } + /** + * Whether to return the capacity consumed by this operation. + * + * @see BatchGetItemRequest.Builder#returnConsumedCapacity(ReturnConsumedCapacity) + * @return a builder of this type + */ + public Builder returnConsumedCapacity(ReturnConsumedCapacity returnConsumedCapacity) { + this.returnConsumedCapacity = returnConsumedCapacity == null ? null : returnConsumedCapacity.toString(); + return this; + } + + /** + * Whether to return the capacity consumed by this operation. + * + * @see BatchGetItemRequest.Builder#returnConsumedCapacity(String) + * @return a builder of this type + */ + public Builder returnConsumedCapacity(String returnConsumedCapacity) { + this.returnConsumedCapacity = returnConsumedCapacity; + return this; + } + public BatchGetItemEnhancedRequest build() { return new BatchGetItemEnhancedRequest(this); } } - } diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetResultPage.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetResultPage.java index fbb5000a74b6..adec895c0cf1 100644 --- a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetResultPage.java +++ b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetResultPage.java @@ -33,6 +33,7 @@ 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.ConsumedCapacity; import software.amazon.awssdk.services.dynamodb.model.KeysAndAttributes; /** @@ -111,6 +112,10 @@ public List unprocessedKeysForTable(MappedTableResource mappedTable) { .collect(Collectors.toList()); } + public List consumedCapacity() { + return this.batchGetItemResponse.consumedCapacity(); + } + /** * A builder that is used to create a result object with the desired parameters. */ diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/AsyncBatchGetItemTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/AsyncBatchGetItemTest.java index 3c90d04562da..c22263f6e4e0 100644 --- a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/AsyncBatchGetItemTest.java +++ b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/AsyncBatchGetItemTest.java @@ -17,9 +17,14 @@ import static org.hamcrest.MatcherAssert.assertThat; import static org.hamcrest.Matchers.containsInAnyOrder; +import static org.hamcrest.Matchers.empty; +import static org.hamcrest.Matchers.hasSize; import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.not; +import static org.hamcrest.Matchers.nullValue; import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.primaryPartitionKey; +import java.util.Arrays; import java.util.List; import java.util.Objects; import java.util.stream.Collectors; @@ -37,58 +42,80 @@ import software.amazon.awssdk.enhanced.dynamodb.model.BatchGetResultPage; import software.amazon.awssdk.enhanced.dynamodb.model.BatchGetResultPagePublisher; import software.amazon.awssdk.enhanced.dynamodb.model.ReadBatch; +import software.amazon.awssdk.services.dynamodb.model.ConsumedCapacity; import software.amazon.awssdk.services.dynamodb.model.DeleteTableRequest; +import software.amazon.awssdk.services.dynamodb.model.ReturnConsumedCapacity; public class AsyncBatchGetItemTest extends LocalDynamoDbAsyncTestBase { private static class Record1 { private Integer id; + private String stringAttr; private Integer getId() { return id; } + private String getStringAttr() { + return stringAttr; + } + private Record1 setId(Integer id) { this.id = id; return this; } + private Record1 setStringAttr(String str) { + this.stringAttr = str; + return this; + } + @Override public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; Record1 record1 = (Record1) o; - return Objects.equals(id, record1.id); + return Objects.equals(id, record1.id) && Objects.equals(stringAttr, record1.stringAttr); } @Override public int hashCode() { - return Objects.hash(id); + return Objects.hash(id, stringAttr); } } private static class Record2 { private Integer id; + private String stringAttr; private Integer getId() { return id; } + private String getStringAttr() { + return stringAttr; + } + private Record2 setId(Integer id) { this.id = id; return this; } + private Record2 setStringAttr(String str) { + this.stringAttr = str; + return this; + } + @Override public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; Record2 record2 = (Record2) o; - return Objects.equals(id, record2.id); + return Objects.equals(id, record2.id) && Objects.equals(stringAttr, record2.stringAttr); } @Override public int hashCode() { - return Objects.hash(id); + return Objects.hash(id, stringAttr); } } @@ -99,6 +126,9 @@ public int hashCode() { .getter(Record1::getId) .setter(Record1::setId) .tags(primaryPartitionKey())) + .addAttribute(String.class, a -> a.name("str_1") + .getter(Record1::getStringAttr) + .setter(Record1::setStringAttr)) .build(); private static final TableSchema TABLE_SCHEMA_2 = @@ -108,26 +138,29 @@ public int hashCode() { .getter(Record2::getId) .setter(Record2::setId) .tags(primaryPartitionKey())) + .addAttribute(String.class, a -> a.name("str_1") + .getter(Record2::getStringAttr) + .setter(Record2::setStringAttr)) .build(); - private DynamoDbEnhancedAsyncClient enhancedAsyncClient = + private final DynamoDbEnhancedAsyncClient enhancedAsyncClient = DefaultDynamoDbEnhancedAsyncClient.builder() .dynamoDbClient(getDynamoDbAsyncClient()) .build(); - private DynamoDbAsyncTable mappedTable1 = enhancedAsyncClient.table(getConcreteTableName("table-name-1"), - TABLE_SCHEMA_1); - private DynamoDbAsyncTable mappedTable2 = enhancedAsyncClient.table(getConcreteTableName("table-name-2"), - TABLE_SCHEMA_2); + private final String tableName1 = getConcreteTableName("table-name-1"); + private final String tableName2 = getConcreteTableName("table-name-2"); + private final DynamoDbAsyncTable mappedTable1 = enhancedAsyncClient.table(tableName1, TABLE_SCHEMA_1); + private final DynamoDbAsyncTable mappedTable2 = enhancedAsyncClient.table(tableName2, TABLE_SCHEMA_2); private static final List RECORDS_1 = IntStream.range(0, 2) - .mapToObj(i -> new Record1().setId(i)) + .mapToObj(i -> new Record1().setId(i).setStringAttr(getStringAttrValue(80_000))) .collect(Collectors.toList()); private static final List RECORDS_2 = IntStream.range(0, 2) - .mapToObj(i -> new Record2().setId(i)) + .mapToObj(i -> new Record2().setId(i).setStringAttr(getStringAttrValue(40_000))) .collect(Collectors.toList()); @Before @@ -139,10 +172,10 @@ public void createTable() { @After public void deleteTable() { getDynamoDbAsyncClient().deleteTable(DeleteTableRequest.builder() - .tableName(getConcreteTableName("table-name-1")) + .tableName(tableName1) .build()).join(); getDynamoDbAsyncClient().deleteTable(DeleteTableRequest.builder() - .tableName(getConcreteTableName("table-name-2")) + .tableName(tableName2) .build()).join(); } @@ -160,6 +193,8 @@ public void getRecordsFromMultipleTables() { List results = drainPublisher(publisher, 1); assertThat(results.size(), is(1)); BatchGetResultPage page = results.get(0); + assertThat(page.consumedCapacity(), empty()); + List record1List = page.resultsForTable(mappedTable1); assertThat(record1List.size(), is(2)); assertThat(record1List, containsInAnyOrder(RECORDS_1.get(0), RECORDS_1.get(1))); @@ -222,6 +257,61 @@ public void notFoundRecordReturnsNull_viaFlattenedItems() { assertThat(record2List, containsInAnyOrder(RECORDS_2.toArray())); } + @Test + public void getRecordsFromMultipleTables_withReturnConsumedCapacity() { + insertRecords(); + + SdkPublisher publisher = batchGetResultPageSdkPublisherForBothTables(ReturnConsumedCapacity.TOTAL); + + List results = drainPublisher(publisher, 1); + assertThat(results.size(), is(1)); + BatchGetResultPage page = results.get(0); + assertThat(page.consumedCapacity(), not(nullValue())); + assertThat(page.consumedCapacity(), hasSize(2)); + + assertThat(page.consumedCapacity(), containsInAnyOrder( + ConsumedCapacity.builder().tableName(tableName1).capacityUnits(20.0).build(), + ConsumedCapacity.builder().tableName(tableName2).capacityUnits(10.0).build() + )); + + List record1List = page.resultsForTable(mappedTable1); + assertThat(record1List.size(), is(2)); + assertThat(record1List, containsInAnyOrder(RECORDS_1.get(0), RECORDS_1.get(1))); + + List record2List = page.resultsForTable(mappedTable2); + assertThat(record2List.size(), is(2)); + assertThat(record2List, containsInAnyOrder(RECORDS_2.get(0), RECORDS_2.get(1))); + } + + @Test + public void notFoundRecords_withReturnConsumedCapacity() { + insertRecords(); + + BatchGetItemEnhancedRequest batchGetItemEnhancedRequest = + requestWithNotFoundRecord().toBuilder() + .returnConsumedCapacity(ReturnConsumedCapacity.TOTAL) + .build(); + + SdkPublisher publisher = enhancedAsyncClient.batchGetItem(batchGetItemEnhancedRequest); + + List results = drainPublisher(publisher, 1); + assertThat(results.size(), is(1)); + BatchGetResultPage page = results.get(0); + assertThat(page.consumedCapacity(), containsInAnyOrder( + ConsumedCapacity.builder().tableName(tableName1).capacityUnits(10.0).build(), + ConsumedCapacity.builder().tableName(tableName2).capacityUnits(10.0).build() + )); + + List record1List = page.resultsForTable(mappedTable1); + assertThat(record1List.size(), is(1)); + assertThat(record1List.get(0).getId(), is(0)); + + List record2List = page.resultsForTable(mappedTable2); + assertThat(record2List.size(), is(2)); + assertThat(record2List, containsInAnyOrder(RECORDS_2.get(0), RECORDS_2.get(1))); + + } + private BatchGetItemEnhancedRequest requestWithNotFoundRecord() { return BatchGetItemEnhancedRequest.builder() .readBatches( @@ -245,23 +335,36 @@ private BatchGetItemEnhancedRequest requestWithNotFoundRecord() { } private BatchGetResultPagePublisher batchGetResultPageSdkPublisherForBothTables() { - return enhancedAsyncClient.batchGetItem(r -> r.readBatches( - ReadBatch.builder(Record1.class) - .mappedTableResource(mappedTable1) - .addGetItem(i -> i.key(k -> k.partitionValue(0))) - .build(), - ReadBatch.builder(Record2.class) - .mappedTableResource(mappedTable2) - .addGetItem(i -> i.key(k -> k.partitionValue(0))) - .build(), - ReadBatch.builder(Record2.class) - .mappedTableResource(mappedTable2) - .addGetItem(i -> i.key(k -> k.partitionValue(1))) - .build(), - ReadBatch.builder(Record1.class) - .mappedTableResource(mappedTable1) - .addGetItem(i -> i.key(k -> k.partitionValue(1))) - .build())); + return batchGetResultPageSdkPublisherForBothTables(null); + } + + private BatchGetResultPagePublisher batchGetResultPageSdkPublisherForBothTables(ReturnConsumedCapacity retConsumedCapacity) { + return enhancedAsyncClient.batchGetItem( + r -> r.returnConsumedCapacity(retConsumedCapacity) + .readBatches( + ReadBatch.builder(Record1.class) + .mappedTableResource(mappedTable1) + .addGetItem(i -> i.key(k -> k.partitionValue(0))) + .build(), + ReadBatch.builder(Record2.class) + .mappedTableResource(mappedTable2) + .addGetItem(i -> i.key(k -> k.partitionValue(0))) + .build(), + ReadBatch.builder(Record2.class) + .mappedTableResource(mappedTable2) + .addGetItem(i -> i.key(k -> k.partitionValue(1))) + .build(), + ReadBatch.builder(Record1.class) + .mappedTableResource(mappedTable1) + .addGetItem(i -> i.key(k -> k.partitionValue(1))) + .build() + ) + ); } -} + private static String getStringAttrValue(int nChars) { + char[] bytes = new char[nChars]; + Arrays.fill(bytes, 'a'); + return new String(bytes); + } +} diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/BatchGetItemTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/BatchGetItemTest.java index 721739157e32..6eb92c2d4409 100644 --- a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/BatchGetItemTest.java +++ b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/BatchGetItemTest.java @@ -17,9 +17,11 @@ 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 software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.primaryPartitionKey; +import java.util.Arrays; import java.util.List; import java.util.Objects; import java.util.stream.Collectors; @@ -36,58 +38,80 @@ import software.amazon.awssdk.enhanced.dynamodb.model.BatchGetResultPage; import software.amazon.awssdk.enhanced.dynamodb.model.BatchGetResultPageIterable; import software.amazon.awssdk.enhanced.dynamodb.model.ReadBatch; +import software.amazon.awssdk.services.dynamodb.model.ConsumedCapacity; import software.amazon.awssdk.services.dynamodb.model.DeleteTableRequest; +import software.amazon.awssdk.services.dynamodb.model.ReturnConsumedCapacity; public class BatchGetItemTest extends LocalDynamoDbSyncTestBase { private static class Record1 { private Integer id; + private String stringAttr; private Integer getId() { return id; } + private String getStringAttr() { + return stringAttr; + } + private Record1 setId(Integer id) { this.id = id; return this; } + private Record1 setStringAttr(String str) { + this.stringAttr = str; + return this; + } + @Override public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; Record1 record1 = (Record1) o; - return Objects.equals(id, record1.id); + return Objects.equals(id, record1.id) && Objects.equals(stringAttr, record1.stringAttr); } @Override public int hashCode() { - return Objects.hash(id); + return Objects.hash(id, stringAttr); } } private static class Record2 { private Integer id; + private String stringAttr; private Integer getId() { return id; } + private String getStringAttr() { + return stringAttr; + } + private Record2 setId(Integer id) { this.id = id; return this; } + private Record2 setStringAttr(String str) { + this.stringAttr = str; + return this; + } + @Override public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; Record2 record2 = (Record2) o; - return Objects.equals(id, record2.id); + return Objects.equals(id, record2.id) && Objects.equals(stringAttr, record2.stringAttr); } @Override public int hashCode() { - return Objects.hash(id); + return Objects.hash(id, stringAttr); } } @@ -98,6 +122,9 @@ public int hashCode() { .getter(Record1::getId) .setter(Record1::setId) .tags(primaryPartitionKey())) + .addAttribute(String.class, a -> a.name("str_1") + .getter(Record1::getStringAttr) + .setter(Record1::setStringAttr)) .build(); private static final TableSchema TABLE_SCHEMA_2 = @@ -107,23 +134,28 @@ public int hashCode() { .getter(Record2::getId) .setter(Record2::setId) .tags(primaryPartitionKey())) + .addAttribute(String.class, a -> a.name("str_1") + .getter(Record2::getStringAttr) + .setter(Record2::setStringAttr)) .build(); private DynamoDbEnhancedClient enhancedClient = DynamoDbEnhancedClient.builder() .dynamoDbClient(getDynamoDbClient()) .build(); - private DynamoDbTable mappedTable1 = enhancedClient.table(getConcreteTableName("table-name-1"), TABLE_SCHEMA_1); - private DynamoDbTable mappedTable2 = enhancedClient.table(getConcreteTableName("table-name-2"), TABLE_SCHEMA_2); + private final String tableName1 = getConcreteTableName("table-name-1"); + private final String tableName2 = getConcreteTableName("table-name-2"); + private final DynamoDbTable mappedTable1 = enhancedClient.table(tableName1, TABLE_SCHEMA_1); + private final DynamoDbTable mappedTable2 = enhancedClient.table(tableName2, TABLE_SCHEMA_2); private static final List RECORDS_1 = IntStream.range(0, 2) - .mapToObj(i -> new Record1().setId(i)) + .mapToObj(i -> new Record1().setId(i).setStringAttr(getStringAttrValue(80_000))) .collect(Collectors.toList()); private static final List RECORDS_2 = IntStream.range(0, 2) - .mapToObj(i -> new Record2().setId(i)) + .mapToObj(i -> new Record2().setId(i).setStringAttr(getStringAttrValue(40_000))) .collect(Collectors.toList()); @Before @@ -135,10 +167,10 @@ public void createTable() { @After public void deleteTable() { getDynamoDbClient().deleteTable(DeleteTableRequest.builder() - .tableName(getConcreteTableName("table-name-1")) + .tableName(tableName1) .build()); getDynamoDbClient().deleteTable(DeleteTableRequest.builder() - .tableName(getConcreteTableName("table-name-2")) + .tableName(tableName2) .build()); } @@ -153,7 +185,8 @@ public void getRecordsFromMultipleTables() { SdkIterable results = getBatchGetResultPagesForBothTables(); assertThat(results.stream().count(), is(1L)); - results.iterator().forEachRemaining((page) -> { + results.iterator().forEachRemaining(page -> { + assertThat(page.consumedCapacity(), empty()); List table1Results = page.resultsForTable(mappedTable1); assertThat(table1Results.size(), is(2)); assertThat(table1Results.get(0).id, is(0)); @@ -186,6 +219,8 @@ public void notFoundRecordIgnored() { assertThat(results.stream().count(), is(1L)); results.iterator().forEachRemaining((page) -> { + assertThat(page.consumedCapacity(), empty()); + List mappedTable1Results = page.resultsForTable(mappedTable1); assertThat(mappedTable1Results.size(), is(1)); assertThat(mappedTable1Results.get(0).id, is(0)); @@ -210,6 +245,51 @@ public void notFoundRecordIgnored_viaFlattenedItems() { assertThat(recordsList2, containsInAnyOrder(RECORDS_2.toArray())); } + @Test + public void getRecordsFromMultipleTables_withReturnConsumedCapacity() { + insertRecords(); + + SdkIterable results = getBatchGetResultPagesForBothTables(ReturnConsumedCapacity.TOTAL); + assertThat(results.stream().count(), is(1L)); + + results.iterator().forEachRemaining(page -> { + assertThat(page.consumedCapacity(), containsInAnyOrder( + ConsumedCapacity.builder().tableName(tableName1).capacityUnits(20.0).build(), + ConsumedCapacity.builder().tableName(tableName2).capacityUnits(10.0).build() + )); + + List table1Results = page.resultsForTable(mappedTable1); + assertThat(table1Results.size(), is(2)); + assertThat(table1Results.get(0).id, is(0)); + assertThat(table1Results.get(1).id, is(1)); + assertThat(page.resultsForTable(mappedTable2).size(), is(2)); + }); + } + + @Test + public void notFoundRecordIgnored_withReturnConsumedCapacity() { + insertRecords(); + + BatchGetItemEnhancedRequest batchGetItemEnhancedRequest = batchGetItemEnhancedRequestWithNotFoundRecord() + .toBuilder() + .returnConsumedCapacity(ReturnConsumedCapacity.TOTAL) + .build(); + SdkIterable results = enhancedClient.batchGetItem(batchGetItemEnhancedRequest); + assertThat(results.stream().count(), is(1L)); + + results.iterator().forEachRemaining(page -> { + assertThat(page.consumedCapacity(), containsInAnyOrder( + ConsumedCapacity.builder().tableName(tableName1).capacityUnits(10.0).build(), + ConsumedCapacity.builder().tableName(tableName2).capacityUnits(10.0).build() + )); + + List mappedTable1Results = page.resultsForTable(mappedTable1); + assertThat(mappedTable1Results.size(), is(1)); + assertThat(mappedTable1Results.get(0).id, is(0)); + assertThat(page.resultsForTable(mappedTable2).size(), is(2)); + }); + } + private BatchGetItemEnhancedRequest batchGetItemEnhancedRequestWithNotFoundRecord() { return BatchGetItemEnhancedRequest.builder() .readBatches( @@ -233,7 +313,11 @@ private BatchGetItemEnhancedRequest batchGetItemEnhancedRequestWithNotFoundRecor } private BatchGetResultPageIterable getBatchGetResultPagesForBothTables() { - return enhancedClient.batchGetItem(r -> r.readBatches( + return getBatchGetResultPagesForBothTables(null); + } + + private BatchGetResultPageIterable getBatchGetResultPagesForBothTables(ReturnConsumedCapacity returnConsumedCapacity) { + return enhancedClient.batchGetItem(r -> r.returnConsumedCapacity(returnConsumedCapacity).readBatches( ReadBatch.builder(Record1.class) .mappedTableResource(mappedTable1) .addGetItem(i -> i.key(k -> k.partitionValue(0))) @@ -251,5 +335,10 @@ private BatchGetResultPageIterable getBatchGetResultPagesForBothTables() { .addGetItem(i -> i.key(k -> k.partitionValue(1))) .build())); } -} + private static String getStringAttrValue(int nChars) { + char[] bytes = new char[nChars]; + Arrays.fill(bytes, 'a'); + return new String(bytes); + } +} diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetItemEnhancedRequestTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetItemEnhancedRequestTest.java index 4fdbc413d7c4..c424363d8465 100644 --- a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetItemEnhancedRequestTest.java +++ b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/model/BatchGetItemEnhancedRequestTest.java @@ -16,7 +16,9 @@ package software.amazon.awssdk.enhanced.dynamodb.model; import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.is; +import static org.hamcrest.Matchers.not; import static org.hamcrest.Matchers.nullValue; import java.util.Collections; @@ -29,6 +31,8 @@ import software.amazon.awssdk.enhanced.dynamodb.DynamoDbTable; import software.amazon.awssdk.enhanced.dynamodb.functionaltests.models.FakeItem; import software.amazon.awssdk.services.dynamodb.DynamoDbClient; +import software.amazon.awssdk.services.dynamodb.model.ConsumedCapacity; +import software.amazon.awssdk.services.dynamodb.model.ReturnConsumedCapacity; @RunWith(MockitoJUnitRunner.class) public class BatchGetItemEnhancedRequestTest { @@ -53,6 +57,8 @@ public void builder_minimal() { BatchGetItemEnhancedRequest builtObject = BatchGetItemEnhancedRequest.builder().build(); assertThat(builtObject.readBatches(), is(nullValue())); + assertThat(builtObject.returnConsumedCapacity(), is(nullValue())); + assertThat(builtObject.returnConsumedCapacityAsString(), is(nullValue())); } @Test @@ -64,9 +70,12 @@ public void builder_maximal() { BatchGetItemEnhancedRequest builtObject = BatchGetItemEnhancedRequest.builder() .readBatches(readBatch) + .returnConsumedCapacity(ReturnConsumedCapacity.TOTAL) .build(); assertThat(builtObject.readBatches(), is(Collections.singletonList(readBatch))); + assertThat(builtObject.returnConsumedCapacity(), equalTo(ReturnConsumedCapacity.TOTAL)); + assertThat(builtObject.returnConsumedCapacityAsString(), equalTo(ReturnConsumedCapacity.TOTAL.toString())); } @Test @@ -90,6 +99,60 @@ public void toBuilder() { BatchGetItemEnhancedRequest copiedObject = builtObject.toBuilder().build(); assertThat(copiedObject, is(builtObject)); + assertThat(copiedObject.hashCode(), equalTo(builtObject.hashCode())); + } + + @Test + public void toBuilder_maximal() { + ReadBatch readBatch = ReadBatch.builder(FakeItem.class) + .mappedTableResource(fakeItemMappedTable) + .addGetItem(r -> r.key(k -> k.partitionValue("key"))) + .build(); + + BatchGetItemEnhancedRequest builtObject = BatchGetItemEnhancedRequest.builder() + .readBatches(readBatch) + .returnConsumedCapacity(ReturnConsumedCapacity.TOTAL) + .build(); + + BatchGetItemEnhancedRequest copiedObject = builtObject.toBuilder().build(); + assertThat(copiedObject, is(builtObject)); + assertThat(copiedObject.hashCode(), equalTo(builtObject.hashCode())); + } + + @Test + public void hashCode_includes_readBatches() { + BatchGetItemEnhancedRequest request1 = BatchGetItemEnhancedRequest.builder().build(); + + ReadBatch readBatch = ReadBatch.builder(FakeItem.class) + .mappedTableResource(fakeItemMappedTable) + .addGetItem(r -> r.key(k -> k.partitionValue("key"))) + .build(); + BatchGetItemEnhancedRequest request2 = BatchGetItemEnhancedRequest.builder() + .readBatches(readBatch) + .build(); + + assertThat(request1, not(equalTo(request2))); + assertThat(request1.hashCode(), not(equalTo(request2))); + } + + @Test + public void hashCode_includes_consumedCapacity() { + BatchGetItemEnhancedRequest request1 = BatchGetItemEnhancedRequest.builder().build(); + BatchGetItemEnhancedRequest request2 = BatchGetItemEnhancedRequest.builder() + .returnConsumedCapacity(ReturnConsumedCapacity.INDEXES) + .build(); + + assertThat(request1, not(equalTo(request2))); + assertThat(request1.hashCode(), not(equalTo(request2))); + } + + @Test + public void returnConsumedCapacity_unknownToSdkVersion() { + BatchGetItemEnhancedRequest builtObject = BatchGetItemEnhancedRequest.builder() + .returnConsumedCapacity("abcdefg") + .build(); + assertThat(builtObject.returnConsumedCapacity(), is(ReturnConsumedCapacity.UNKNOWN_TO_SDK_VERSION)); + assertThat(builtObject.returnConsumedCapacityAsString(), equalTo("abcdefg")); } }