Skip to content

Commit b4ecd2d

Browse files
prateek-vatsragnar9780zoewangg
authored
feature: Support of ReturnConsumedCapacity in DynamoDBEnhancedClient's BatchWrite operation (#5462)
* Add support for ConsumedCapacity in a BatchWrite request * Added changelog * Add support for consumed capacity in transactWriteItems * Revert "Add support for consumed capacity in transactWriteItems" This reverts commit a7bd547. * Add support for consumed capacity in transactWriteItems (cherry picked from commit a7bd547) * Address review comments * Revert "Add support for consumed capacity in transactWriteItems" This reverts commit c58db98. * Address javadoc comments * Fix build --------- Co-authored-by: prateek.vats <prateek.vats@salesforce.com> Co-authored-by: Zoe Wang <33073555+zoewangg@users.noreply.github.com>
1 parent e89007e commit b4ecd2d

File tree

10 files changed

+668
-55
lines changed

10 files changed

+668
-55
lines changed
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"type": "feature",
3+
"category": "DynamoDBEnhancedClient",
4+
"description": "This commit introduces ConsumedCapacity in the response of a BatchWrite response",
5+
"contributor": "prateek-vats"
6+
}

services-custom/dynamodb-enhanced/pom.xml

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,6 +141,11 @@
141141
<artifactId>utils</artifactId>
142142
<version>${awsjavasdk.version}</version>
143143
</dependency>
144+
<dependency>
145+
<groupId>com.google.guava</groupId>
146+
<artifactId>guava</artifactId>
147+
<scope>test</scope>
148+
</dependency>
144149
<dependency>
145150
<groupId>software.amazon.awssdk</groupId>
146151
<artifactId>service-test-utils</artifactId>

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/operations/BatchWriteItemOperation.java

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,13 +67,19 @@ public BatchWriteItemRequest generateRequest(DynamoDbEnhancedClientExtension ext
6767
return BatchWriteItemRequest.builder()
6868
.requestItems(
6969
Collections.unmodifiableMap(CollectionUtils.deepCopyMap(allRequestItems)))
70+
.returnConsumedCapacity(request.returnConsumedCapacity())
71+
.returnItemCollectionMetrics(request.returnItemCollectionMetrics())
7072
.build();
7173
}
7274

7375
@Override
7476
public BatchWriteResult transformResponse(BatchWriteItemResponse response,
7577
DynamoDbEnhancedClientExtension extension) {
76-
return BatchWriteResult.builder().unprocessedRequests(response.unprocessedItems()).build();
78+
return BatchWriteResult.builder()
79+
.unprocessedRequests(response.unprocessedItems())
80+
.consumedCapacity(response.consumedCapacity())
81+
.itemCollectionMetrics(response.itemCollectionMetrics())
82+
.build();
7783
}
7884

7985
@Override

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

Lines changed: 105 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,26 +20,34 @@
2020
import java.util.Collection;
2121
import java.util.Collections;
2222
import java.util.List;
23+
import java.util.Objects;
2324
import software.amazon.awssdk.annotations.NotThreadSafe;
2425
import software.amazon.awssdk.annotations.SdkPublicApi;
2526
import software.amazon.awssdk.annotations.ThreadSafe;
2627
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient;
28+
import software.amazon.awssdk.services.dynamodb.model.BatchWriteItemRequest;
29+
import software.amazon.awssdk.services.dynamodb.model.ReturnConsumedCapacity;
30+
import software.amazon.awssdk.services.dynamodb.model.ReturnItemCollectionMetrics;
2731

2832
/**
2933
* Defines parameters used for the batchWriteItem() operation (such as
3034
* {@link DynamoDbEnhancedClient#batchWriteItem(BatchWriteItemEnhancedRequest)}).
3135
* <p>
32-
* A request contains references to keys for delete actions and items for put actions,
33-
* organized into one {@link WriteBatch} object per accessed table.
36+
* A request contains references to keys for delete actions and items for put actions, organized into one {@link WriteBatch}
37+
* object per accessed table.
3438
*/
3539
@SdkPublicApi
3640
@ThreadSafe
3741
public final class BatchWriteItemEnhancedRequest {
3842

3943
private final List<WriteBatch> writeBatches;
44+
private final String returnConsumedCapacity;
45+
private final String returnItemCollectionMetrics;
4046

4147
private BatchWriteItemEnhancedRequest(Builder builder) {
4248
this.writeBatches = getListIfExist(builder.writeBatches);
49+
this.returnConsumedCapacity = builder.returnConsumedCapacity;
50+
this.returnItemCollectionMetrics = builder.returnItemCollectionMetrics;
4351
}
4452

4553
/**
@@ -50,10 +58,50 @@ public static Builder builder() {
5058
}
5159

5260
/**
53-
* Returns a builder initialized with all existing values on the request object.
61+
* @return a builder with all existing values set
5462
*/
5563
public Builder toBuilder() {
56-
return new Builder().writeBatches(writeBatches);
64+
return builder().writeBatches(writeBatches)
65+
.returnConsumedCapacity(returnConsumedCapacity)
66+
.returnItemCollectionMetrics(returnItemCollectionMetrics);
67+
}
68+
69+
/**
70+
* Whether to return the capacity consumed by this operation.
71+
*
72+
* @see BatchWriteItemEnhancedRequest#returnConsumedCapacity()
73+
*/
74+
public ReturnConsumedCapacity returnConsumedCapacity() {
75+
return ReturnConsumedCapacity.fromValue(returnConsumedCapacity);
76+
}
77+
78+
/**
79+
* Whether to return the capacity consumed by this operation.
80+
* <p>
81+
* Similar to {@link #returnConsumedCapacity()} but return the value as a string. This is useful in situations where the value
82+
* is not defined in {@link ReturnConsumedCapacity}.
83+
*/
84+
public String returnConsumedCapacityAsString() {
85+
return returnConsumedCapacity;
86+
}
87+
88+
/**
89+
* Whether to return the item collection metrics.
90+
*
91+
* @see BatchWriteItemRequest#returnItemCollectionMetrics()
92+
*/
93+
public ReturnItemCollectionMetrics returnItemCollectionMetrics() {
94+
return ReturnItemCollectionMetrics.fromValue(returnItemCollectionMetrics);
95+
}
96+
97+
/**
98+
* Whether to return the item collection metrics.
99+
* <p>
100+
* Similar to {@link #returnItemCollectionMetrics()} but return the value as a string. This is useful in situations where the
101+
* value is not defined in {@link ReturnItemCollectionMetrics}.
102+
*/
103+
public String returnItemCollectionMetricsAsString() {
104+
return returnItemCollectionMetrics;
57105
}
58106

59107
/**
@@ -71,15 +119,18 @@ public boolean equals(Object o) {
71119
if (o == null || getClass() != o.getClass()) {
72120
return false;
73121
}
74-
75122
BatchWriteItemEnhancedRequest that = (BatchWriteItemEnhancedRequest) o;
76-
77-
return writeBatches != null ? writeBatches.equals(that.writeBatches) : that.writeBatches == null;
123+
return Objects.equals(writeBatches, that.writeBatches) &&
124+
Objects.equals(returnConsumedCapacity, that.returnConsumedCapacity) &&
125+
Objects.equals(returnItemCollectionMetrics, that.returnItemCollectionMetrics);
78126
}
79127

80128
@Override
81129
public int hashCode() {
82-
return writeBatches != null ? writeBatches.hashCode() : 0;
130+
int result = writeBatches != null ? writeBatches.hashCode() : 0;
131+
result = 31 * result + (returnConsumedCapacity != null ? returnConsumedCapacity.hashCode() : 0);
132+
result = 31 * result + (returnItemCollectionMetrics != null ? returnItemCollectionMetrics.hashCode() : 0);
133+
return result;
83134
}
84135

85136
private static List<WriteBatch> getListIfExist(List<WriteBatch> writeBatches) {
@@ -92,10 +143,54 @@ private static List<WriteBatch> getListIfExist(List<WriteBatch> writeBatches) {
92143
@NotThreadSafe
93144
public static final class Builder {
94145
private List<WriteBatch> writeBatches;
146+
private String returnConsumedCapacity;
147+
private String returnItemCollectionMetrics;
95148

96149
private Builder() {
97150
}
98151

152+
/**
153+
* Sets the ConsumedCapacity via ReturnConsumedCapacity object
154+
*
155+
* @see Builder#returnConsumedCapacity(ReturnConsumedCapacity)
156+
*/
157+
public Builder returnConsumedCapacity(ReturnConsumedCapacity returnConsumedCapacity) {
158+
this.returnConsumedCapacity = returnConsumedCapacity == null ? null : returnConsumedCapacity.toString();
159+
return this;
160+
}
161+
162+
/**
163+
* Sets the ConsumedCapacity via String.
164+
*
165+
* @see Builder#returnConsumedCapacity(String)
166+
*/
167+
public Builder returnConsumedCapacity(String returnConsumedCapacity) {
168+
this.returnConsumedCapacity = returnConsumedCapacity;
169+
return this;
170+
}
171+
172+
/**
173+
* Sets the item collection metrics.
174+
*
175+
* @see BatchWriteItemEnhancedRequest.Builder#returnItemCollectionMetrics(ReturnItemCollectionMetrics)
176+
*/
177+
public BatchWriteItemEnhancedRequest.Builder returnItemCollectionMetrics(ReturnItemCollectionMetrics
178+
returnItemCollectionMetrics) {
179+
this.returnItemCollectionMetrics = returnItemCollectionMetrics == null ? null :
180+
returnItemCollectionMetrics.toString();
181+
return this;
182+
}
183+
184+
/**
185+
* Sets the item collection metrics.
186+
*
187+
* @see BatchWriteItemEnhancedRequest.Builder#returnItemCollectionMetrics(String)
188+
*/
189+
public BatchWriteItemEnhancedRequest.Builder returnItemCollectionMetrics(String returnItemCollectionMetrics) {
190+
this.returnItemCollectionMetrics = returnItemCollectionMetrics;
191+
return this;
192+
}
193+
99194
/**
100195
* Sets a collection of write batches to use in the batchWriteItem operation.
101196
*
@@ -119,8 +214,8 @@ public Builder writeBatches(WriteBatch... writeBatches) {
119214
}
120215

121216
/**
122-
* Adds a write batch to the collection of batches on this builder.
123-
* If this is the first batch, the method creates a new list.
217+
* Adds a write batch to the collection of batches on this builder. If this is the first batch, the method creates a new
218+
* list.
124219
*
125220
* @param writeBatch a single write batch
126221
* @return a builder of this type

0 commit comments

Comments
 (0)