Skip to content

Commit 6c310a6

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 6c310a6

File tree

7 files changed

+1175
-0
lines changed

7 files changed

+1175
-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: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
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 GSI/LSI, 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/GSI/LSI 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 GSI/LSI. If not, throw IllegalArgumentException
92+
allowedKeys.add(meta.primaryPartitionKey());
93+
meta.primarySortKey().ifPresent(allowedKeys::add);
94+
95+
for (IndexMetadata idx : meta.indices()) {
96+
String indexName = idx.name();
97+
allowedKeys.add(meta.indexPartitionKey(indexName));
98+
meta.indexSortKey(indexName).ifPresent(allowedKeys::add);
99+
}
100+
101+
for (String attr : taggedAttributes) {
102+
if (!allowedKeys.contains(attr)) {
103+
throw new IllegalArgumentException(
104+
"@DynamoDbAutoGeneratedKey can only be applied to key attributes: " +
105+
"primary partition key, primary sort key, or GSI/LSI partition/sort keys." +
106+
"Invalid placement on attribute: " + attr);
107+
108+
}
109+
}
110+
111+
// Generate UUIDs for missing/empty annotated attributes
112+
Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
113+
taggedAttributes.forEach(attr -> insertUuidIfMissing(itemToTransform, attr));
114+
115+
return WriteModification.builder()
116+
.transformedItem(Collections.unmodifiableMap(itemToTransform))
117+
.build();
118+
}
119+
120+
private void insertUuidIfMissing(Map<String, AttributeValue> itemToTransform, String key) {
121+
AttributeValue existing = itemToTransform.get(key);
122+
boolean missing = existing == null || existing.s() == null || existing.s().isEmpty();
123+
if (missing) {
124+
itemToTransform.put(key, AttributeValue.builder().s(UUID.randomUUID().toString()).build());
125+
}
126+
}
127+
128+
/**
129+
* Static helpers used by the {@code @BeanTableSchemaAttributeTag}-based annotation tag.
130+
*/
131+
public static final class AttributeTags {
132+
private AttributeTags() {
133+
}
134+
135+
/**
136+
* @return a {@link StaticAttributeTag} that marks the attribute for auto-generated key behavior.
137+
*/
138+
public static StaticAttributeTag autoGeneratedKeyAttribute() {
139+
return AUTO_GENERATED_KEY_ATTRIBUTE;
140+
}
141+
}
142+
143+
/**
144+
* Stateless builder.
145+
*/
146+
public static final class Builder {
147+
private Builder() {
148+
}
149+
150+
public AutoGeneratedKeyExtension build() {
151+
return new AutoGeneratedKeyExtension();
152+
}
153+
}
154+
155+
/**
156+
* Validates Java type and records the tagged attribute into table metadata so {@link #beforeWrite} can find it at runtime.
157+
*/
158+
private static final class AutoGeneratedKeyAttribute implements StaticAttributeTag {
159+
160+
@Override
161+
public <R> void validateType(String attributeName,
162+
EnhancedType<R> type,
163+
AttributeValueType attributeValueType) {
164+
165+
Validate.notNull(type, "type is null");
166+
Validate.notNull(type.rawClass(), "rawClass is null");
167+
Validate.notNull(attributeValueType, "attributeValueType is null");
168+
169+
if (!String.class.equals(type.rawClass())) {
170+
throw new IllegalArgumentException(String.format(
171+
"Attribute '%s' of Java type %s is not valid for @DynamoDbAutoGeneratedKey. Only String is supported.",
172+
attributeName, type.rawClass()));
173+
}
174+
}
175+
176+
@Override
177+
public Consumer<StaticTableMetadata.Builder> modifyMetadata(String attributeName,
178+
AttributeValueType attributeValueType) {
179+
// Record the names of the attributes annotated with @DynamoDbAutoGeneratedKey for later lookup in beforeWrite()
180+
return metadata -> metadata.addCustomMetadataObject(
181+
CUSTOM_METADATA_KEY, Collections.singleton(attributeName));
182+
}
183+
}
184+
}

‎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)