Skip to content

Commit 14ff258

Browse files
anasatirbasaanasatirbasa
committed
Added support for DynamoDbAutoGeneratedKey annotation
1 parent d3aed1c commit 14ff258

File tree

2 files changed

+142
-5
lines changed

2 files changed

+142
-5
lines changed

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

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@
2727
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbUpdateBehavior;
2828

2929
/**
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).
30+
* Annotation that marks a key attribute to be automatically populated with a random UUID if no value is provided during a
31+
* write operation (put or update). This annotation is intended to work specifically with key attributes.
3232
*
3333
* <p>This annotation is designed for use with the V2 {@link software.amazon.awssdk.enhanced.dynamodb.mapper.BeanTableSchema}.
3434
* It is registered via {@link BeanTableSchemaAttributeTag} and its behavior is implemented by
@@ -37,7 +37,7 @@
3737
* <h3>Where this annotation can be applied</h3>
3838
* This annotation is only valid on attributes that serve as keys:
3939
* <ul>
40-
* <li>The table’s primary partition key or sort key</li>
40+
* <li>The table's primary partition key or sort key</li>
4141
* <li>The partition key or sort key of a secondary index (GSI or LSI)</li>
4242
* </ul>
4343
* If applied to any other attribute, the {@code AutoGeneratedKeyExtension} will throw an
@@ -51,14 +51,16 @@
5151
* </ul>
5252
*
5353
* <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:
54+
* This annotation can be combined with {@link DynamoDbUpdateBehavior} to control regeneration behavior:
5655
* <ul>
5756
* <li>{@link UpdateBehavior#WRITE_ALWAYS} (default) –
5857
* Generate a new UUID whenever the attribute is missing during write.</li>
5958
* <li>{@link UpdateBehavior#WRITE_IF_NOT_EXISTS} –
6059
* Generate a UUID only the first time (on insert), and preserve that value on subsequent updates.</li>
6160
* </ul>
61+
* <p><strong>Important:</strong> {@link DynamoDbUpdateBehavior} only affects secondary index keys.
62+
* Primary keys cannot be null in DynamoDB and are required for UpdateItem operations, so UpdateBehavior
63+
* has no practical effect on primary keys.</p>
6264
*
6365
* <h3>Type restriction</h3>
6466
* This annotation is only valid on attributes of type {@link String}.

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

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
import static org.junit.Assert.assertEquals;
1010
import static software.amazon.awssdk.enhanced.dynamodb.functionaltests.AutoGeneratedUuidRecordTest.assertValidUuid;
1111
import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.primaryPartitionKey;
12+
import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.primarySortKey;
1213
import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.secondaryPartitionKey;
1314
import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.secondarySortKey;
1415
import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.updateBehavior;
@@ -37,6 +38,7 @@
3738
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbPartitionKey;
3839
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSecondaryPartitionKey;
3940
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSecondarySortKey;
41+
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSortKey;
4042
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbUpdateBehavior;
4143

4244
@RunWith(Parameterized.class)
@@ -292,6 +294,81 @@ public void autogenerateKey_onVersionedRecord_setOnPrimaryKey_performsUuidGenera
292294
}
293295
}
294296

297+
@Test
298+
public void autogenerateKey_onPrimarySortKey_performsUuidGeneration() {
299+
String tableName = getConcreteTableName("autogen-primary-sk");
300+
DynamoDbEnhancedClient client = DynamoDbEnhancedClient.builder()
301+
.dynamoDbClient(getDynamoDbClient())
302+
.extensions(AutoGeneratedKeyExtension.builder().build())
303+
.build();
304+
305+
TableSchema<PrimarySortKeyBean> schema = StaticTableSchema.builder(PrimarySortKeyBean.class)
306+
.newItemSupplier(PrimarySortKeyBean::new)
307+
.addAttribute(String.class, a -> a.name("id")
308+
.getter(PrimarySortKeyBean::getId)
309+
.setter(PrimarySortKeyBean::setId)
310+
.addTag(primaryPartitionKey()))
311+
.addAttribute(String.class, a -> a.name("sortKey")
312+
.getter(PrimarySortKeyBean::getSortKey)
313+
.setter(PrimarySortKeyBean::setSortKey)
314+
.tags(primarySortKey(),
315+
AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute()))
316+
.build();
317+
318+
DynamoDbTable<PrimarySortKeyBean> table = client.table(tableName, schema);
319+
try {
320+
table.createTable(r -> r.provisionedThroughput(getDefaultProvisionedThroughput()));
321+
PrimarySortKeyBean bean = new PrimarySortKeyBean();
322+
bean.setId("id123"); // sortKey missing → should be generated
323+
table.putItem(bean);
324+
325+
PrimarySortKeyBean out = table.scan().items().stream().findFirst()
326+
.orElseThrow(() -> new AssertionError("No record found"));
327+
assertValidUuid(out.getSortKey());
328+
assertEquals("id123", out.getId());
329+
} finally {
330+
deleteTableByName(tableName);
331+
}
332+
}
333+
334+
@Test
335+
public void autogenerateKey_onPrimarySortKeyWithUpdateBehavior_generatesOnPutOnly() {
336+
String tableName = getConcreteTableName("autogen-primary-sk-update");
337+
DynamoDbEnhancedClient client = DynamoDbEnhancedClient.builder()
338+
.dynamoDbClient(getDynamoDbClient())
339+
.extensions(AutoGeneratedKeyExtension.builder().build())
340+
.build();
341+
342+
TableSchema<PrimarySortKeyWithUpdateBehaviorBean> schema = TableSchema.fromBean(PrimarySortKeyWithUpdateBehaviorBean.class);
343+
DynamoDbTable<PrimarySortKeyWithUpdateBehaviorBean> table = client.table(tableName, schema);
344+
345+
try {
346+
table.createTable(r -> r.provisionedThroughput(getDefaultProvisionedThroughput()));
347+
348+
// First put - should generate UUID for primary sort key
349+
PrimarySortKeyWithUpdateBehaviorBean bean = new PrimarySortKeyWithUpdateBehaviorBean();
350+
bean.setId("id123");
351+
bean.setPayload("initial");
352+
table.putItem(bean);
353+
354+
PrimarySortKeyWithUpdateBehaviorBean afterPut = table.scan().items().stream().findFirst()
355+
.orElseThrow(() -> new AssertionError("No record found"));
356+
String generatedSortKey = afterPut.getSortKey();
357+
assertValidUuid(generatedSortKey);
358+
assertEquals("initial", afterPut.getPayload());
359+
360+
// Update - sort key should remain the same (UpdateBehavior has no effect on primary keys)
361+
afterPut.setPayload("updated");
362+
table.updateItem(afterPut);
363+
364+
PrimarySortKeyWithUpdateBehaviorBean afterUpdate = table.getItem(r -> r.key(k -> k.partitionValue("id123").sortValue(generatedSortKey)));
365+
assertThat(afterUpdate.getSortKey()).isEqualTo(generatedSortKey); // Key preserved (UpdateBehavior irrelevant for primary keys)
366+
assertEquals("updated", afterUpdate.getPayload());
367+
} finally {
368+
deleteTableByName(tableName);
369+
}
370+
}
371+
295372
/**
296373
* - createdKey: GSI SK + @DynamoDbAutoGeneratedKey annotation (UpdateBehaviour is WRITE_IF_NOT_EXISTS)
297374
* - lastUpdatedKey: GSI PK + @DynamoDbAutoGeneratedKey annotation (UpdateBehaviour is WRITE_ALWAYS - default value)
@@ -533,6 +610,64 @@ public void setPayload(String payload) {
533610
}
534611
}
535612

613+
public static class PrimarySortKeyBean {
614+
private String id;
615+
private String sortKey;
616+
617+
public String getId() {
618+
return id;
619+
}
620+
621+
public void setId(String id) {
622+
this.id = id;
623+
}
624+
625+
public String getSortKey() {
626+
return sortKey;
627+
}
628+
629+
public void setSortKey(String sortKey) {
630+
this.sortKey = sortKey;
631+
}
632+
}
633+
634+
@DynamoDbBean
635+
public static class PrimarySortKeyWithUpdateBehaviorBean {
636+
private String id;
637+
private String sortKey;
638+
private String payload;
639+
640+
@DynamoDbPartitionKey
641+
public String getId() {
642+
return id;
643+
}
644+
645+
public void setId(String id) {
646+
this.id = id;
647+
}
648+
649+
// NOTE: @DynamoDbUpdateBehavior on primary sort key has no practical effect
650+
// since primary keys cannot be null in DynamoDB and are required for UpdateItem
651+
@DynamoDbSortKey
652+
@DynamoDbAutoGeneratedKey
653+
@DynamoDbUpdateBehavior(UpdateBehavior.WRITE_IF_NOT_EXISTS)
654+
public String getSortKey() {
655+
return sortKey;
656+
}
657+
658+
public void setSortKey(String sortKey) {
659+
this.sortKey = sortKey;
660+
}
661+
662+
public String getPayload() {
663+
return payload;
664+
}
665+
666+
public void setPayload(String payload) {
667+
this.payload = payload;
668+
}
669+
}
670+
536671
private void deleteTableByName(String tableName) {
537672
getDynamoDbClient().deleteTable(b -> b.tableName(tableName));
538673
}

0 commit comments

Comments
 (0)