Skip to content

Commit a919686

Browse files
anasatirbasaanasatirbasa
committed
Added support for DynamoDbAutoGeneratedKey annotation
1 parent 3351169 commit a919686

File tree

2 files changed

+80
-13
lines changed

2 files changed

+80
-13
lines changed

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import java.util.HashMap;
2121
import java.util.HashSet;
2222
import java.util.Map;
23+
import java.util.Objects;
2324
import java.util.Set;
2425
import java.util.UUID;
2526
import java.util.function.Consumer;
@@ -34,6 +35,7 @@
3435
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
3536
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
3637
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
38+
import software.amazon.awssdk.utils.StringUtils;
3739
import software.amazon.awssdk.utils.Validate;
3840

3941
/**
@@ -119,8 +121,7 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
119121

120122
private void insertUuidIfMissing(Map<String, AttributeValue> itemToTransform, String key) {
121123
AttributeValue existing = itemToTransform.get(key);
122-
boolean missing = existing == null || existing.s() == null || existing.s().isEmpty();
123-
if (missing) {
124+
if (Objects.isNull(existing) || StringUtils.isBlank(existing.s())) {
124125
itemToTransform.put(key, AttributeValue.builder().s(UUID.randomUUID().toString()).build());
125126
}
126127
}

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

Lines changed: 77 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ public void createTable() {
114114

115115
@After
116116
public void deleteTable() {
117-
getDynamoDbClient().deleteTable(b -> b.tableName(tableName));
117+
mappedTable.deleteTable();
118118
}
119119

120120
@Test
@@ -153,7 +153,7 @@ public void update_preservesKeysWithWriteIfNotExists_and_regeneratesKeysWithWrit
153153

154154
@Test
155155
public void autogenerateKey_setOnGsiPartitionKey_performsUuidGeneration() {
156-
String tn = getConcreteTableName("autogen-gsi-pk");
156+
String tableName = getConcreteTableName("autogen-gsi-pk");
157157
DynamoDbEnhancedClient client = DynamoDbEnhancedClient.builder()
158158
.dynamoDbClient(getDynamoDbClient())
159159
.extensions(AutoGeneratedKeyExtension.builder().build())
@@ -172,7 +172,7 @@ public void autogenerateKey_setOnGsiPartitionKey_performsUuidGeneration() {
172172
AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute()))
173173
.build();
174174

175-
DynamoDbTable<GsiPartitionKeyBean> table = client.table(tn, schema);
175+
DynamoDbTable<GsiPartitionKeyBean> table = client.table(tableName, schema);
176176
try {
177177
table.createTable(r -> r.provisionedThroughput(getDefaultProvisionedThroughput()));
178178
GsiPartitionKeyBean bean = new GsiPartitionKeyBean();
@@ -182,13 +182,13 @@ public void autogenerateKey_setOnGsiPartitionKey_performsUuidGeneration() {
182182
GsiPartitionKeyBean out = table.getItem(r -> r.key(k -> k.partitionValue("id123")));
183183
assertValidUuid(out.getGsiPk());
184184
} finally {
185-
getDynamoDbClient().deleteTable(b -> b.tableName(tn));
185+
deleteTableByName(tableName);
186186
}
187187
}
188188

189189
@Test
190190
public void autogenerateKey_setOnSecondaryIndexSortKey_performsUuidGeneration() {
191-
String tn = getConcreteTableName("autogen-idx-sk");
191+
String tableName = getConcreteTableName("autogen-idx-sk");
192192
DynamoDbEnhancedClient client = DynamoDbEnhancedClient.builder()
193193
.dynamoDbClient(getDynamoDbClient())
194194
.extensions(AutoGeneratedKeyExtension.builder().build())
@@ -208,7 +208,7 @@ public void autogenerateKey_setOnSecondaryIndexSortKey_performsUuidGeneration()
208208
AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute()))
209209
.build();
210210

211-
DynamoDbTable<SecondaryIndexSortKeyBean> table = client.table(tn, schema);
211+
DynamoDbTable<SecondaryIndexSortKeyBean> table = client.table(tableName, schema);
212212
try {
213213
table.createTable(r -> r.provisionedThroughput(getDefaultProvisionedThroughput()));
214214
SecondaryIndexSortKeyBean bean = new SecondaryIndexSortKeyBean();
@@ -218,13 +218,13 @@ public void autogenerateKey_setOnSecondaryIndexSortKey_performsUuidGeneration()
218218
SecondaryIndexSortKeyBean out = table.getItem(r -> r.key(k -> k.partitionValue("id123")));
219219
assertValidUuid(out.getSecondarySortKey());
220220
} finally {
221-
getDynamoDbClient().deleteTable(b -> b.tableName(tn));
221+
deleteTableByName(tableName);
222222
}
223223
}
224224

225225
@Test
226226
public void autogenerateKey_setOnInvalidKeyAttribute_throws_exception() {
227-
String tn = getConcreteTableName("autogen-non-key");
227+
String tableName = getConcreteTableName("autogen-non-key");
228228
DynamoDbEnhancedClient client = DynamoDbEnhancedClient.builder()
229229
.dynamoDbClient(getDynamoDbClient())
230230
.extensions(AutoGeneratedKeyExtension.builder().build())
@@ -242,7 +242,7 @@ public void autogenerateKey_setOnInvalidKeyAttribute_throws_exception() {
242242
.addTag(AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute()))
243243
.build();
244244

245-
DynamoDbTable<NonKeyBean> table = client.table(tn, schema);
245+
DynamoDbTable<NonKeyBean> table = client.table(tableName, schema);
246246
try {
247247
table.createTable(r -> r.provisionedThroughput(getDefaultProvisionedThroughput()));
248248
NonKeyBean bean = new NonKeyBean();
@@ -253,7 +253,7 @@ public void autogenerateKey_setOnInvalidKeyAttribute_throws_exception() {
253253
+ "primary partition key, primary sort key, or GSI/LSI partition/sort keys.")
254254
.withMessageContaining("notAKey");
255255
} finally {
256-
getDynamoDbClient().deleteTable(b -> b.tableName(tn));
256+
deleteTableByName(tableName);
257257
}
258258
}
259259

@@ -288,7 +288,44 @@ public void autogenerateKey_onVersionedRecord_setOnPrimaryKey_performsUuidGenera
288288
assertEquals("new_payload", updated.getPayload());
289289
assertThat(updated.getVersion()).isEqualTo(2L);
290290
} finally {
291-
getDynamoDbClient().deleteTable(b -> b.tableName(versionedTableName));
291+
deleteTableByName(versionedTableName);
292+
}
293+
}
294+
295+
@Test
296+
public void autogenerateKey_withUpdateBehaviorOnPrimaryKey_preservesKeyOnUpdate() {
297+
String tableName = getConcreteTableName("autogen-pk-update-behavior");
298+
DynamoDbEnhancedClient client = DynamoDbEnhancedClient.builder()
299+
.dynamoDbClient(getDynamoDbClient())
300+
.extensions(AutoGeneratedKeyExtension.builder().build())
301+
.build();
302+
303+
TableSchema<PrimaryKeyWithUpdateBehaviorBean> schema = TableSchema.fromBean(PrimaryKeyWithUpdateBehaviorBean.class);
304+
DynamoDbTable<PrimaryKeyWithUpdateBehaviorBean> table = client.table(tableName, schema);
305+
306+
try {
307+
table.createTable(r -> r.provisionedThroughput(getDefaultProvisionedThroughput()));
308+
309+
// First put - should generate UUID for primary key
310+
PrimaryKeyWithUpdateBehaviorBean bean = new PrimaryKeyWithUpdateBehaviorBean();
311+
bean.setPayload("initial");
312+
table.putItem(bean);
313+
314+
PrimaryKeyWithUpdateBehaviorBean afterPut = table.scan().items().stream().findFirst()
315+
.orElseThrow(() -> new AssertionError("No record found"));
316+
String generatedId = afterPut.getId();
317+
assertValidUuid(generatedId);
318+
assertEquals("initial", afterPut.getPayload());
319+
320+
// Update - should preserve the generated primary key due to WRITE_IF_NOT_EXISTS
321+
afterPut.setPayload("updated");
322+
table.updateItem(afterPut);
323+
324+
PrimaryKeyWithUpdateBehaviorBean afterUpdate = table.getItem(r -> r.key(k -> k.partitionValue(generatedId)));
325+
assertThat(afterUpdate.getId()).isEqualTo(generatedId); // Key preserved
326+
assertEquals("updated", afterUpdate.getPayload());
327+
} finally {
328+
deleteTableByName(tableName);
292329
}
293330
}
294331

@@ -532,4 +569,33 @@ public void setPayload(String payload) {
532569
this.payload = payload;
533570
}
534571
}
572+
573+
@DynamoDbBean
574+
public static class PrimaryKeyWithUpdateBehaviorBean {
575+
private String id;
576+
private String payload;
577+
578+
@DynamoDbPartitionKey
579+
@DynamoDbAutoGeneratedKey
580+
@DynamoDbUpdateBehavior(UpdateBehavior.WRITE_IF_NOT_EXISTS)
581+
public String getId() {
582+
return id;
583+
}
584+
585+
public void setId(String id) {
586+
this.id = id;
587+
}
588+
589+
public String getPayload() {
590+
return payload;
591+
}
592+
593+
public void setPayload(String payload) {
594+
this.payload = payload;
595+
}
596+
}
597+
598+
private void deleteTableByName(String tableName) {
599+
getDynamoDbClient().deleteTable(b -> b.tableName(tableName));
600+
}
535601
}

0 commit comments

Comments
 (0)