Skip to content

Commit d874530

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 1c0b19f commit d874530

File tree

7 files changed

+1171
-0
lines changed

7 files changed

+1171
-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: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
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.HashSet;
22+
import java.util.Map;
23+
import java.util.Set;
24+
import java.util.UUID;
25+
import java.util.function.Consumer;
26+
import software.amazon.awssdk.annotations.SdkPublicApi;
27+
import software.amazon.awssdk.annotations.ThreadSafe;
28+
import software.amazon.awssdk.enhanced.dynamodb.AttributeValueType;
29+
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
30+
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbExtensionContext;
31+
import software.amazon.awssdk.enhanced.dynamodb.EnhancedType;
32+
import software.amazon.awssdk.enhanced.dynamodb.IndexMetadata;
33+
import software.amazon.awssdk.enhanced.dynamodb.TableMetadata;
34+
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
35+
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
36+
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
37+
import software.amazon.awssdk.utils.Validate;
38+
39+
/**
40+
* Generates a random UUID (via {@link java.util.UUID#randomUUID()}) for any attribute tagged with
41+
* {@code @DynamoDbAutoGeneratedKey} when that attribute is missing or empty on a write (put/update).
42+
* <p>
43+
* The annotation may be placed <b>only</b> on key attributes:
44+
* <ul>
45+
* <li>Primary partition key (PK) or primary sort key (SK)</li>
46+
* <li>Partition key or sort key of any secondary index (GSI or LSI)</li>
47+
* </ul>
48+
*
49+
* <p><b>Validation:</b> The extension enforces this at runtime during {@link #beforeWrite} by comparing the
50+
* annotated attributes against the table's known key attributes. If an annotated attribute
51+
* is not a PK/SK or an LSI/GSI, an {@link IllegalArgumentException} is thrown.</p>
52+
*/
53+
@SdkPublicApi
54+
@ThreadSafe
55+
public final class AutoGeneratedKeyExtension implements DynamoDbEnhancedClientExtension {
56+
57+
/**
58+
* Custom metadata key under which we store the set of annotated attribute names.
59+
*/
60+
private static final String CUSTOM_METADATA_KEY =
61+
"software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension:AutoGeneratedKeyAttribute";
62+
63+
private static final AutoGeneratedKeyAttribute AUTO_GENERATED_KEY_ATTRIBUTE = new AutoGeneratedKeyAttribute();
64+
65+
private AutoGeneratedKeyExtension() {
66+
}
67+
68+
public static Builder builder() {
69+
return new Builder();
70+
}
71+
72+
/**
73+
* If this table has attributes tagged for auto-generation, insert a UUID value into the outgoing item for any such attribute
74+
* that is currently missing/empty.
75+
* <p>
76+
* Also validates that the annotation is only used on PK/SK/LSI/GSI key attributes.
77+
*/
78+
@Override
79+
public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
80+
Collection<String> taggedAttributes = context.tableMetadata()
81+
.customMetadataObject(CUSTOM_METADATA_KEY, Collection.class)
82+
.orElse(null);
83+
84+
if (taggedAttributes == null || taggedAttributes.isEmpty()) {
85+
return WriteModification.builder().build();
86+
}
87+
88+
TableMetadata meta = context.tableMetadata();
89+
Set<String> allowedKeys = new HashSet<>();
90+
91+
// ensure every @DynamoDbAutoGeneratedKey attribute is a PK/SK or LSI/GSI. If not, throw IllegalArgumentException.
92+
allowedKeys.add(meta.primaryPartitionKey()); // get the PK attributes
93+
meta.primarySortKey().ifPresent(allowedKeys::add); // get the SK attributes
94+
95+
for (IndexMetadata idx : meta.indices()) {
96+
String indexName = idx.name();
97+
allowedKeys.add(meta.indexPartitionKey(indexName)); // get the GSI attributes
98+
meta.indexSortKey(indexName).ifPresent(allowedKeys::add); // get the LSI attributes
99+
}
100+
101+
for (String attr : taggedAttributes) {
102+
if (!allowedKeys.contains(attr)) {
103+
// if the annotated attributes is not a PK/SK or an LSI/GSI, throw IllegalArgumentException
104+
throw new IllegalArgumentException(
105+
"@DynamoDbAutoGeneratedKey is only allowed on primary or secondary index keys. "
106+
+ "Invalid placement on attribute: " + attr);
107+
}
108+
}
109+
110+
// Generate UUIDs for missing/empty annotated attributes
111+
Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
112+
taggedAttributes.forEach(attr -> insertUuidIfMissing(itemToTransform, attr));
113+
114+
return WriteModification.builder()
115+
.transformedItem(Collections.unmodifiableMap(itemToTransform))
116+
.build();
117+
}
118+
119+
private void insertUuidIfMissing(Map<String, AttributeValue> itemToTransform, String key) {
120+
AttributeValue existing = itemToTransform.get(key);
121+
boolean missing = existing == null || existing.s() == null || existing.s().isEmpty();
122+
if (missing) {
123+
itemToTransform.put(key, AttributeValue.builder().s(UUID.randomUUID().toString()).build());
124+
}
125+
}
126+
127+
/**
128+
* Static helpers used by the {@code @BeanTableSchemaAttributeTag}-based annotation tag.
129+
*/
130+
public static final class AttributeTags {
131+
private AttributeTags() {
132+
}
133+
134+
/**
135+
* @return a {@link StaticAttributeTag} that marks the attribute for auto-generated key behavior.
136+
*/
137+
public static StaticAttributeTag autoGeneratedKeyAttribute() {
138+
return AUTO_GENERATED_KEY_ATTRIBUTE;
139+
}
140+
}
141+
142+
/**
143+
* Stateless builder.
144+
*/
145+
public static final class Builder {
146+
private Builder() {
147+
}
148+
149+
public AutoGeneratedKeyExtension build() {
150+
return new AutoGeneratedKeyExtension();
151+
}
152+
}
153+
154+
/**
155+
* Validates Java type and records the tagged attribute into table metadata so {@link #beforeWrite} can find it at runtime.
156+
*/
157+
private static final class AutoGeneratedKeyAttribute implements StaticAttributeTag {
158+
159+
@Override
160+
public <R> void validateType(String attributeName,
161+
EnhancedType<R> type,
162+
AttributeValueType attributeValueType) {
163+
164+
Validate.notNull(type, "type is null");
165+
Validate.notNull(type.rawClass(), "rawClass is null");
166+
Validate.notNull(attributeValueType, "attributeValueType is null");
167+
168+
if (!String.class.equals(type.rawClass())) {
169+
throw new IllegalArgumentException(String.format(
170+
"Attribute '%s' of Java type %s is not valid for @DynamoDbAutoGeneratedKey. Only String is supported.",
171+
attributeName, type.rawClass()));
172+
}
173+
}
174+
175+
@Override
176+
public Consumer<StaticTableMetadata.Builder> modifyMetadata(String attributeName,
177+
AttributeValueType attributeValueType) {
178+
// Record the names of the attributes annotated with @DynamoDbAutoGeneratedKey for later lookup in beforeWrite()
179+
return metadata -> metadata.addCustomMetadataObject(
180+
CUSTOM_METADATA_KEY, Collections.singleton(attributeName));
181+
}
182+
}
183+
}

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
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.UpdateBehavior;
26+
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.BeanTableSchemaAttributeTag;
27+
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbUpdateBehavior;
28+
29+
/**
30+
* Annotation that marks a string attribute to be automatically populated with a random UUID if no value is provided during a
31+
* write operation (put or update).
32+
*
33+
* <p>This annotation is designed for use with the V2 {@link software.amazon.awssdk.enhanced.dynamodb.mapper.BeanTableSchema}.
34+
* It is registered via {@link BeanTableSchemaAttributeTag} and its behavior is implemented by
35+
* {@link software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension}.</p>
36+
*
37+
* <h3>Where this annotation can be applied</h3>
38+
* This annotation is only valid on attributes that serve as keys:
39+
* <ul>
40+
* <li>The table’s primary partition key or sort key</li>
41+
* <li>The partition key or sort key of a secondary index (GSI or LSI)</li>
42+
* </ul>
43+
* If applied to any other attribute, the {@code AutoGeneratedKeyExtension} will throw an
44+
* {@link IllegalArgumentException} at runtime.
45+
*
46+
* <h3>How values are generated</h3>
47+
* <ul>
48+
* <li>On writes where the annotated attribute is null or empty, a new UUID value is generated
49+
* using {@link java.util.UUID#randomUUID()}.</li>
50+
* <li>If a value is already set on the attribute, that value is preserved and not replaced.</li>
51+
* </ul>
52+
*
53+
* <h3>Controlling regeneration on update</h3>
54+
* This annotation can be combined with {@link DynamoDbUpdateBehavior} to control whether a new
55+
* UUID should be generated on each update:
56+
* <ul>
57+
* <li>{@link UpdateBehavior#WRITE_ALWAYS} (default) –
58+
* Generate a new UUID whenever the attribute is missing during write.</li>
59+
* <li>{@link UpdateBehavior#WRITE_IF_NOT_EXISTS} –
60+
* Generate a UUID only the first time (on insert), and preserve that value on subsequent updates.</li>
61+
* </ul>
62+
*
63+
* <h3>Type restriction</h3>
64+
* This annotation is only valid on attributes of type {@link String}.
65+
*/
66+
@SdkPublicApi
67+
@Documented
68+
@Retention(RetentionPolicy.RUNTIME)
69+
@Target( {ElementType.METHOD, ElementType.FIELD})
70+
@BeanTableSchemaAttributeTag(AutoGeneratedKeyTag.class)
71+
public @interface DynamoDbAutoGeneratedKey {
72+
}

‎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+
}

‎services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/UuidTestUtils.java

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
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.UUID;
19+
20+
public class UuidTestUtils {
21+
22+
public static boolean isValidUuid(String uuid) {
23+
try {
24+
UUID.fromString(uuid);
25+
return true;
26+
} catch (Exception e) {
27+
return false;
28+
}
29+
}
30+
}

0 commit comments

Comments
 (0)