Skip to content

Commit 0099ee7

Browse files
zoewanggzoewangg
committed
Create wrapper class for ddb enhanced extension input parameters
1 parent 86b05d4 commit 0099ee7

20 files changed

+458
-202
lines changed

‎services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/DynamoDbEnhancedClientExtension.java

Lines changed: 6 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,9 @@
1515

1616
package software.amazon.awssdk.enhanced.dynamodb;
1717

18-
import java.util.Map;
1918
import software.amazon.awssdk.annotations.SdkPublicApi;
2019
import software.amazon.awssdk.enhanced.dynamodb.extensions.ReadModification;
2120
import software.amazon.awssdk.enhanced.dynamodb.extensions.WriteModification;
22-
import software.amazon.awssdk.enhanced.dynamodb.internal.operations.OperationContext;
23-
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
2421

2522
/**
2623
* Interface for extending the DynamoDb Enhanced client. Two hooks are provided, one that is called just before a record
@@ -37,29 +34,23 @@ public interface DynamoDbEnhancedClientExtension {
3734
* This hook is called just before an operation is going to write data to the database. The extension that
3835
* implements this method can choose to transform the item itself, or add a condition to the write operation
3936
* or both.
40-
* @param item The item that is about to be written.
41-
* @param operationContext The context under which the operation to be modified is taking place.
42-
* @param tableMetadata The structure of the table.
37+
*
38+
* @param context The {@link DynamoDbExtensionContext.BeforeWrite} context containing the state of the execution.
4339
* @return A {@link WriteModification} object that can alter the behavior of the write operation.
4440
*/
45-
default WriteModification beforeWrite(Map<String, AttributeValue> item,
46-
OperationContext operationContext,
47-
TableMetadata tableMetadata) {
41+
default WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
4842
return WriteModification.builder().build();
4943
}
5044

5145
/**
5246
* This hook is called just after an operation that has read data from the database. The extension that
5347
* implements this method can choose to transform the item, and then it is the transformed item that will be
5448
* mapped back to the application instead of the item that was actually read from the database.
55-
* @param item The item that has just been read.
56-
* @param operationContext The context under which the operation to be modified is taking place.
57-
* @param tableMetadata The structure of the table.
49+
*
50+
* @param context The {@link DynamoDbExtensionContext.AfterRead} context containing the state of the execution.
5851
* @return A {@link ReadModification} object that can alter the results of a read operation.
5952
*/
60-
default ReadModification afterRead(Map<String, AttributeValue> item,
61-
OperationContext operationContext,
62-
TableMetadata tableMetadata) {
53+
default ReadModification afterRead(DynamoDbExtensionContext.AfterRead context) {
6354
return ReadModification.builder().build();
6455
}
6556
}

‎services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/DynamoDbExtensionContext.java

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.enhanced.dynamodb;
17+
18+
import java.util.Map;
19+
import software.amazon.awssdk.annotations.SdkPublicApi;
20+
import software.amazon.awssdk.enhanced.dynamodb.internal.operations.OperationContext;
21+
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
22+
23+
/**
24+
* A wrapper for the immutable context objects that are visible to the {@link DynamoDbEnhancedClientExtension}s.
25+
*/
26+
@SdkPublicApi
27+
public final class DynamoDbExtensionContext {
28+
private DynamoDbExtensionContext() {
29+
}
30+
31+
@SdkPublicApi
32+
public interface Context {
33+
/**
34+
* @return The {@link AttributeValue} map of the items that is about to be written or has just been read.
35+
*/
36+
Map<String, AttributeValue> items();
37+
38+
/**
39+
* @return The context under which the operation to be modified is taking place.
40+
*/
41+
OperationContext operationContext();
42+
43+
/**
44+
* @return A {@link TableMetadata} object describing the structure of the modelled table.
45+
*/
46+
TableMetadata tableMetadata();
47+
}
48+
49+
/**
50+
* The state of the execution when the {@link DynamoDbEnhancedClientExtension#beforeWrite} method is invoked.
51+
*/
52+
@SdkPublicApi
53+
public interface BeforeWrite extends Context {
54+
}
55+
56+
/**
57+
* The state of the execution when the {@link DynamoDbEnhancedClientExtension#afterRead} method is invoked.
58+
*/
59+
@SdkPublicApi
60+
public interface AfterRead extends Context {
61+
}
62+
}

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

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,8 @@
2626
import software.amazon.awssdk.annotations.SdkPublicApi;
2727
import software.amazon.awssdk.enhanced.dynamodb.AttributeValueType;
2828
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
29+
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbExtensionContext;
2930
import software.amazon.awssdk.enhanced.dynamodb.Expression;
30-
import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
31-
import software.amazon.awssdk.enhanced.dynamodb.internal.operations.OperationContext;
3231
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
3332
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
3433
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
@@ -90,16 +89,15 @@ public Consumer<StaticTableMetadata.Builder> modifyMetadata(String attributeName
9089
}
9190

9291
@Override
93-
public WriteModification beforeWrite(Map<String, AttributeValue> item,
94-
OperationContext operationContext,
95-
TableMetadata tableMetadata) {
96-
Optional<String> versionAttributeKey = tableMetadata.customMetadataObject(CUSTOM_METADATA_KEY, String.class);
92+
public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
93+
Optional<String> versionAttributeKey = context.tableMetadata()
94+
.customMetadataObject(CUSTOM_METADATA_KEY, String.class);
9795

9896
if (!versionAttributeKey.isPresent()) {
9997
return WriteModification.builder().build();
10098
}
10199

102-
Map<String, AttributeValue> itemToTransform = new HashMap<>(item);
100+
Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
103101
AttributeValue newVersionValue;
104102
Expression condition;
105103
Optional<AttributeValue> existingVersionValue =

‎services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/EnhancedClientUtils.java

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import software.amazon.awssdk.enhanced.dynamodb.Key;
2828
import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
2929
import software.amazon.awssdk.enhanced.dynamodb.extensions.ReadModification;
30+
import software.amazon.awssdk.enhanced.dynamodb.internal.extensions.DefaultDynamoDbExtensionContext;
3031
import software.amazon.awssdk.enhanced.dynamodb.internal.operations.OperationContext;
3132
import software.amazon.awssdk.enhanced.dynamodb.model.Page;
3233
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
@@ -67,10 +68,12 @@ public static <T> T readAndTransformSingleItem(Map<String, AttributeValue> itemM
6768
}
6869

6970
if (dynamoDbEnhancedClientExtension != null) {
70-
ReadModification readModification = dynamoDbEnhancedClientExtension.afterRead(itemMap,
71-
operationContext,
72-
tableSchema.tableMetadata());
73-
71+
ReadModification readModification = dynamoDbEnhancedClientExtension.afterRead(
72+
DefaultDynamoDbExtensionContext.builder()
73+
.items(itemMap)
74+
.operationContext(operationContext)
75+
.tableMetadata(tableSchema.tableMetadata())
76+
.build());
7477
if (readModification != null && readModification.transformedItem() != null) {
7578
return tableSchema.mapToItem(readModification.transformedItem());
7679
}

‎services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/ChainExtension.java

Lines changed: 21 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,10 @@
2323
import java.util.concurrent.atomic.AtomicReference;
2424
import software.amazon.awssdk.annotations.SdkInternalApi;
2525
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
26+
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbExtensionContext;
2627
import software.amazon.awssdk.enhanced.dynamodb.Expression;
27-
import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
2828
import software.amazon.awssdk.enhanced.dynamodb.extensions.ReadModification;
2929
import software.amazon.awssdk.enhanced.dynamodb.extensions.WriteModification;
30-
import software.amazon.awssdk.enhanced.dynamodb.internal.operations.OperationContext;
3130
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
3231

3332
/**
@@ -84,22 +83,23 @@ public static ChainExtension create(List<DynamoDbEnhancedClientExtension> extens
8483
* Multiple conditional statements will be separated by the string " AND ". Expression values will be coalesced
8584
* unless they conflict in which case an exception will be thrown.
8685
*
87-
* @param item The {@link AttributeValue} map of the item to be written.
88-
* @param tableMetadata A {@link TableMetadata} object describing the structure of the modelled table.
86+
* @param context A {@link DynamoDbExtensionContext.BeforeWrite} context
8987
* @return A single {@link WriteModification} representing the coalesced results of all the chained extensions.
9088
*/
9189
@Override
92-
public WriteModification beforeWrite(Map<String, AttributeValue> item,
93-
OperationContext operationContext,
94-
TableMetadata tableMetadata) {
90+
public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
9591
AtomicReference<Map<String, AttributeValue>> transformedItem = new AtomicReference<>();
9692
AtomicReference<Expression> conditionalExpression = new AtomicReference<>();
9793

9894
this.extensionChain.forEach(extension -> {
99-
Map<String, AttributeValue> itemToTransform = transformedItem.get() == null ? item : transformedItem.get();
100-
WriteModification writeModification = extension.beforeWrite(itemToTransform,
101-
operationContext,
102-
tableMetadata);
95+
Map<String, AttributeValue> itemToTransform = transformedItem.get() == null ? context.items() : transformedItem.get();
96+
DynamoDbExtensionContext.BeforeWrite beforeWrite =
97+
DefaultDynamoDbExtensionContext.builder()
98+
.items(itemToTransform)
99+
.operationContext(context.operationContext())
100+
.tableMetadata(context.tableMetadata())
101+
.build();
102+
WriteModification writeModification = extension.beforeWrite(beforeWrite);
103103

104104
if (writeModification.transformedItem() != null) {
105105
transformedItem.set(writeModification.transformedItem());
@@ -127,19 +127,22 @@ public WriteModification beforeWrite(Map<String, AttributeValue> item,
127127
* Implementation of the {@link DynamoDbEnhancedClientExtension} interface that will call all the chained extensions
128128
* in reverse order, passing the results of each one to the next and coalescing the results into a single modification.
129129
*
130-
* @param item The {@link AttributeValue} map of the item that is being read.
131-
* @param tableMetadata A {@link TableMetadata} object describing the structure of the modelled table.
130+
* @param context A {@link DynamoDbExtensionContext.AfterRead} context
132131
* @return A single {@link ReadModification} representing the final transformation of all the chained extensions.
133132
*/
134133
@Override
135-
public ReadModification afterRead(Map<String, AttributeValue> item,
136-
OperationContext operationContext,
137-
TableMetadata tableMetadata) {
134+
public ReadModification afterRead(DynamoDbExtensionContext.AfterRead context) {
138135
AtomicReference<Map<String, AttributeValue>> transformedItem = new AtomicReference<>();
139136

140137
this.extensionChain.descendingIterator().forEachRemaining(extension -> {
141-
Map<String, AttributeValue> itemToTransform = transformedItem.get() == null ? item : transformedItem.get();
142-
ReadModification readModification = extension.afterRead(itemToTransform, operationContext, tableMetadata);
138+
Map<String, AttributeValue> itemToTransform = transformedItem.get() == null ? context.items() : transformedItem.get();
139+
DynamoDbExtensionContext.AfterRead afterRead =
140+
DefaultDynamoDbExtensionContext.builder().items(itemToTransform)
141+
.operationContext(context.operationContext())
142+
.tableMetadata(context.tableMetadata())
143+
.build();
144+
145+
ReadModification readModification = extension.afterRead(afterRead);
143146

144147
if (readModification.transformedItem() != null) {
145148
transformedItem.set(readModification.transformedItem());

‎services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/DefaultDynamoDbExtensionContext.java

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.enhanced.dynamodb.internal.extensions;
17+
18+
import java.util.Map;
19+
import java.util.Objects;
20+
import software.amazon.awssdk.annotations.SdkInternalApi;
21+
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbExtensionContext;
22+
import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
23+
import software.amazon.awssdk.enhanced.dynamodb.internal.operations.OperationContext;
24+
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
25+
26+
/**
27+
* An SDK-internal implementation of {@link DynamoDbExtensionContext.BeforeWrite} and {@link DynamoDbExtensionContext.AfterRead}.
28+
*/
29+
@SdkInternalApi
30+
public final class DefaultDynamoDbExtensionContext implements DynamoDbExtensionContext.BeforeWrite,
31+
DynamoDbExtensionContext.AfterRead {
32+
private Map<String, AttributeValue> items;
33+
private OperationContext operationContext;
34+
private TableMetadata tableMetadata;
35+
36+
private DefaultDynamoDbExtensionContext(Builder builder) {
37+
this.items = builder.items;
38+
this.operationContext = builder.operationContext;
39+
this.tableMetadata = builder.tableMetadata;
40+
}
41+
42+
public static Builder builder() {
43+
return new Builder();
44+
}
45+
46+
@Override
47+
public Map<String, AttributeValue> items() {
48+
return items;
49+
}
50+
51+
@Override
52+
public OperationContext operationContext() {
53+
return operationContext;
54+
}
55+
56+
@Override
57+
public TableMetadata tableMetadata() {
58+
return tableMetadata;
59+
}
60+
61+
@Override
62+
public boolean equals(Object o) {
63+
if (this == o) {
64+
return true;
65+
}
66+
if (o == null || getClass() != o.getClass()) {
67+
return false;
68+
}
69+
70+
DefaultDynamoDbExtensionContext that = (DefaultDynamoDbExtensionContext) o;
71+
72+
if (!Objects.equals(items, that.items)) {
73+
return false;
74+
}
75+
if (!Objects.equals(operationContext, that.operationContext)) {
76+
return false;
77+
}
78+
return Objects.equals(tableMetadata, that.tableMetadata);
79+
}
80+
81+
@Override
82+
public int hashCode() {
83+
int result = items != null ? items.hashCode() : 0;
84+
result = 31 * result + (operationContext != null ? operationContext.hashCode() : 0);
85+
result = 31 * result + (tableMetadata != null ? tableMetadata.hashCode() : 0);
86+
return result;
87+
}
88+
89+
public static final class Builder {
90+
private Map<String, AttributeValue> items;
91+
private OperationContext operationContext;
92+
private TableMetadata tableMetadata;
93+
94+
public Builder items(Map<String, AttributeValue> item) {
95+
this.items = item;
96+
return this;
97+
}
98+
99+
public Builder operationContext(OperationContext operationContext) {
100+
this.operationContext = operationContext;
101+
return this;
102+
}
103+
104+
public Builder tableMetadata(TableMetadata tableMetadata) {
105+
this.tableMetadata = tableMetadata;
106+
return this;
107+
}
108+
109+
public DefaultDynamoDbExtensionContext build() {
110+
return new DefaultDynamoDbExtensionContext(this);
111+
}
112+
}
113+
}

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

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
2525
import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
2626
import software.amazon.awssdk.enhanced.dynamodb.extensions.WriteModification;
27+
import software.amazon.awssdk.enhanced.dynamodb.internal.extensions.DefaultDynamoDbExtensionContext;
2728
import software.amazon.awssdk.enhanced.dynamodb.model.PutItemEnhancedRequest;
2829
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
2930
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
@@ -69,7 +70,13 @@ public PutItemRequest generateRequest(TableSchema<T> tableSchema,
6970
Map<String, AttributeValue> itemMap = tableSchema.itemToMap(this.request.item(), alwaysIgnoreNulls);
7071

7172
WriteModification transformation =
72-
extension != null ? extension.beforeWrite(itemMap, operationContext, tableMetadata) : null;
73+
extension != null ? extension.beforeWrite(
74+
DefaultDynamoDbExtensionContext.builder()
75+
.items(itemMap)
76+
.operationContext(operationContext)
77+
.tableMetadata(tableMetadata)
78+
.build())
79+
: null;
7380

7481
if (transformation != null && transformation.transformedItem() != null) {
7582
itemMap = transformation.transformedItem();

0 commit comments

Comments
 (0)