Skip to content

Commit 9ac5ce0

Browse files
anasatirbasaanasatirbasa
committed
Added support for DynamoDbAutoGeneratedKey annotation [Fix for #5497 - ConditionalCheckFailedException with version attribute and partition key using auto-generated UUID]
1 parent c497318 commit 9ac5ce0

File tree

6 files changed

+898
-0
lines changed

6 files changed

+898
-0
lines changed

‎.changes/next-release/feature-AmazonDynamoDBEnhancedClient-cbcc2bb.json

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
{
2+
"type": "feature",
3+
"category": "Amazon DynamoDB Enhanced Client",
4+
"contributor": "",
5+
"description": "Added the support for DynamoDbAutoGeneratedKey annotation"
6+
}

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

Lines changed: 143 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,143 @@
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.extensions;
17+
18+
import java.util.Collection;
19+
import java.util.Collections;
20+
import java.util.HashMap;
21+
import java.util.Map;
22+
import java.util.UUID;
23+
import java.util.function.Consumer;
24+
import software.amazon.awssdk.annotations.SdkPublicApi;
25+
import software.amazon.awssdk.annotations.ThreadSafe;
26+
import software.amazon.awssdk.enhanced.dynamodb.AttributeValueType;
27+
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
28+
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbExtensionContext;
29+
import software.amazon.awssdk.enhanced.dynamodb.EnhancedType;
30+
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
31+
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
32+
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
33+
import software.amazon.awssdk.utils.Validate;
34+
35+
/**
36+
* This extension facilitates the automatic generation of a unique UUID for any attribute tagged with
37+
* {@code @DynamoDbAutoGeneratedKey}. The generated value uses {@link UUID#randomUUID()}.
38+
* <p>
39+
* Register this extension explicitly when building the {@link software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient}.
40+
* <p>
41+
* Example:
42+
* <pre>
43+
* DynamoDbEnhancedClient client = DynamoDbEnhancedClient.builder()
44+
* .dynamoDbClient(lowLevelClient)
45+
* .extensions(AutoGeneratedKeyExtension.create())
46+
* .build();
47+
* </pre>
48+
*/
49+
@SdkPublicApi
50+
@ThreadSafe
51+
public final class AutoGeneratedKeyExtension implements DynamoDbEnhancedClientExtension {
52+
53+
private static final String CUSTOM_METADATA_KEY =
54+
"software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension:AutoGeneratedKeyAttribute";
55+
56+
private static final AutoGeneratedKeyAttribute AUTO_GENERATED_KEY_ATTRIBUTE = new AutoGeneratedKeyAttribute();
57+
58+
private AutoGeneratedKeyExtension() {
59+
}
60+
61+
/**
62+
* @return an instance of {@link AutoGeneratedKeyExtension}
63+
*/
64+
public static AutoGeneratedKeyExtension create() {
65+
return new AutoGeneratedKeyExtension();
66+
}
67+
68+
/**
69+
* If this table has attributes tagged for auto-generation, insert a UUID value into the outgoing item for any such attribute
70+
* that is currently missing/empty.
71+
*/
72+
@Override
73+
public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
74+
Collection<String> taggedAttributes =
75+
context.tableMetadata()
76+
.customMetadataObject(CUSTOM_METADATA_KEY, Collection.class)
77+
.orElse(null);
78+
79+
if (taggedAttributes == null || taggedAttributes.isEmpty()) {
80+
return WriteModification.builder().build();
81+
}
82+
83+
Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
84+
taggedAttributes.forEach(attr -> insertUuidIfMissing(itemToTransform, attr));
85+
86+
return WriteModification.builder()
87+
.transformedItem(Collections.unmodifiableMap(itemToTransform))
88+
.build();
89+
}
90+
91+
private void insertUuidIfMissing(Map<String, AttributeValue> itemToTransform, String key) {
92+
AttributeValue existing = itemToTransform.get(key);
93+
boolean missing = existing == null || existing.s() == null || existing.s().isEmpty();
94+
if (missing) {
95+
itemToTransform.put(key, AttributeValue.builder().s(UUID.randomUUID().toString()).build());
96+
}
97+
}
98+
99+
/**
100+
* Static helpers used by the {@code @BeanTableSchemaAttributeTag}-based annotation tag.
101+
*/
102+
public static final class AttributeTags {
103+
private AttributeTags() {
104+
}
105+
106+
/**
107+
* @return a {@link StaticAttributeTag} that marks the attribute for auto-generated key behavior.
108+
*/
109+
public static StaticAttributeTag autoGeneratedKeyAttribute() {
110+
return AUTO_GENERATED_KEY_ATTRIBUTE;
111+
}
112+
}
113+
114+
/**
115+
* Validates the Java type and writes table metadata so {@link #beforeWrite} can find the tagged attributes at runtime.
116+
*/
117+
private static final class AutoGeneratedKeyAttribute implements StaticAttributeTag {
118+
119+
@Override
120+
public <R> void validateType(String attributeName,
121+
EnhancedType<R> type,
122+
AttributeValueType attributeValueType) {
123+
124+
Validate.notNull(type, "type is null");
125+
Validate.notNull(type.rawClass(), "rawClass is null");
126+
Validate.notNull(attributeValueType, "attributeValueType is null");
127+
128+
if (!String.class.equals(type.rawClass())) {
129+
throw new IllegalArgumentException(String.format(
130+
"Attribute '%s' of Java type %s is not valid for @DynamoDbAutoGeneratedKey. Only String is supported.",
131+
attributeName, type.rawClass()));
132+
}
133+
}
134+
135+
@Override
136+
public Consumer<StaticTableMetadata.Builder> modifyMetadata(String attributeName,
137+
AttributeValueType attributeValueType) {
138+
return metadata -> metadata
139+
.addCustomMetadataObject(CUSTOM_METADATA_KEY, Collections.singleton(attributeName))
140+
.markAttributeAsKey(attributeName, attributeValueType);
141+
}
142+
}
143+
}

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

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
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.extensions.annotations;
17+
18+
import java.lang.annotation.Documented;
19+
import java.lang.annotation.ElementType;
20+
import java.lang.annotation.Retention;
21+
import java.lang.annotation.RetentionPolicy;
22+
import java.lang.annotation.Target;
23+
import software.amazon.awssdk.annotations.SdkPublicApi;
24+
import software.amazon.awssdk.enhanced.dynamodb.internal.extensions.AutoGeneratedKeyTag;
25+
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.BeanTableSchemaAttributeTag;
26+
27+
/**
28+
* Annotation that marks a DynamoDB attribute as an auto-generated key.
29+
* <p>
30+
* When applied, the associated attribute will be automatically populated with a randomly generated {@link java.util.UUID} string
31+
* value if it is not explicitly set when performing a {@code putItem} operation. This behavior occurs only when the attribute is
32+
* initially absent; existing values are preserved and will never be overwritten or regenerated during subsequent updates.
33+
* </p>
34+
*
35+
* <p>
36+
* Typical usage is to apply this annotation to a partition key or sort key property of type {@link String}. The annotated element
37+
* can be either a JavaBean-style getter method or a class field.
38+
* </p>
39+
*
40+
* <p>
41+
* This annotation is functionally similar to the legacy V1 {@code @DynamoDBAutoGeneratedKey} and is intended for use with the V2
42+
* {@link software.amazon.awssdk.enhanced.dynamodb.mapper.BeanTableSchema}.
43+
* </p>
44+
*
45+
* @see java.util.UUID
46+
* @see software.amazon.awssdk.enhanced.dynamodb.mapper.BeanTableSchema
47+
*/
48+
@SdkPublicApi
49+
@Documented
50+
@Retention(RetentionPolicy.RUNTIME)
51+
@Target( {ElementType.METHOD, ElementType.FIELD})
52+
@BeanTableSchemaAttributeTag(AutoGeneratedKeyTag.class)
53+
public @interface DynamoDbAutoGeneratedKey {
54+
55+
}

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

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
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 software.amazon.awssdk.annotations.SdkInternalApi;
19+
import software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension;
20+
import software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbAutoGeneratedKey;
21+
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
22+
23+
@SdkInternalApi
24+
public final class AutoGeneratedKeyTag {
25+
26+
private AutoGeneratedKeyTag() {
27+
}
28+
29+
public static StaticAttributeTag attributeTagFor(DynamoDbAutoGeneratedKey annotation) {
30+
return AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute();
31+
}
32+
33+
}

0 commit comments

Comments
 (0)