diff --git a/.changes/next-release/feature-AmazonDynamoDBEnhancedClient-cbcc2bb.json b/.changes/next-release/feature-AmazonDynamoDBEnhancedClient-cbcc2bb.json new file mode 100644 index 000000000000..b8c7700af796 --- /dev/null +++ b/.changes/next-release/feature-AmazonDynamoDBEnhancedClient-cbcc2bb.json @@ -0,0 +1,6 @@ +{ + "type": "feature", + "category": "Amazon DynamoDB Enhanced Client", + "contributor": "", + "description": "Added the support for DynamoDbAutoGeneratedKey annotation" +} diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/extensions/AutoGeneratedKeyExtension.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/extensions/AutoGeneratedKeyExtension.java new file mode 100644 index 000000000000..52c7fc612149 --- /dev/null +++ b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/extensions/AutoGeneratedKeyExtension.java @@ -0,0 +1,184 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.enhanced.dynamodb.extensions; + +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.UUID; +import java.util.function.Consumer; +import software.amazon.awssdk.annotations.SdkPublicApi; +import software.amazon.awssdk.annotations.ThreadSafe; +import software.amazon.awssdk.enhanced.dynamodb.AttributeValueType; +import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension; +import software.amazon.awssdk.enhanced.dynamodb.DynamoDbExtensionContext; +import software.amazon.awssdk.enhanced.dynamodb.EnhancedType; +import software.amazon.awssdk.enhanced.dynamodb.IndexMetadata; +import software.amazon.awssdk.enhanced.dynamodb.TableMetadata; +import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag; +import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; +import software.amazon.awssdk.utils.Validate; + +/** + * Generates a random UUID (via {@link java.util.UUID#randomUUID()}) for any attribute tagged with + * {@code @DynamoDbAutoGeneratedKey} when that attribute is missing or empty on a write (put/update). + *

+ * The annotation may be placed only on key attributes: + *

+ * + *

Validation: The extension enforces this at runtime during {@link #beforeWrite} by comparing the + * annotated attributes against the table's known key attributes. If an annotated attribute + * is not a PK/SK or an GSI/LSI, an {@link IllegalArgumentException} is thrown.

+ */ +@SdkPublicApi +@ThreadSafe +public final class AutoGeneratedKeyExtension implements DynamoDbEnhancedClientExtension { + + /** + * Custom metadata key under which we store the set of annotated attribute names. + */ + private static final String CUSTOM_METADATA_KEY = + "software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension:AutoGeneratedKeyAttribute"; + + private static final AutoGeneratedKeyAttribute AUTO_GENERATED_KEY_ATTRIBUTE = new AutoGeneratedKeyAttribute(); + + private AutoGeneratedKeyExtension() { + } + + public static Builder builder() { + return new Builder(); + } + + /** + * If this table has attributes tagged for auto-generation, insert a UUID value into the outgoing item for any such attribute + * that is currently missing/empty. + *

+ * Also validates that the annotation is only used on PK/SK/GSI/LSI key attributes. + */ + @Override + public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) { + Collection taggedAttributes = context.tableMetadata() + .customMetadataObject(CUSTOM_METADATA_KEY, Collection.class) + .orElse(null); + + if (taggedAttributes == null || taggedAttributes.isEmpty()) { + return WriteModification.builder().build(); + } + + TableMetadata meta = context.tableMetadata(); + Set allowedKeys = new HashSet<>(); + + // ensure every @DynamoDbAutoGeneratedKey attribute is a PK/SK or GSI/LSI. If not, throw IllegalArgumentException + allowedKeys.add(meta.primaryPartitionKey()); + meta.primarySortKey().ifPresent(allowedKeys::add); + + for (IndexMetadata idx : meta.indices()) { + String indexName = idx.name(); + allowedKeys.add(meta.indexPartitionKey(indexName)); + meta.indexSortKey(indexName).ifPresent(allowedKeys::add); + } + + for (String attr : taggedAttributes) { + if (!allowedKeys.contains(attr)) { + throw new IllegalArgumentException( + "@DynamoDbAutoGeneratedKey can only be applied to key attributes: " + + "primary partition key, primary sort key, or GSI/LSI partition/sort keys." + + "Invalid placement on attribute: " + attr); + + } + } + + // Generate UUIDs for missing/empty annotated attributes + Map itemToTransform = new HashMap<>(context.items()); + taggedAttributes.forEach(attr -> insertUuidIfMissing(itemToTransform, attr)); + + return WriteModification.builder() + .transformedItem(Collections.unmodifiableMap(itemToTransform)) + .build(); + } + + private void insertUuidIfMissing(Map itemToTransform, String key) { + AttributeValue existing = itemToTransform.get(key); + boolean missing = existing == null || existing.s() == null || existing.s().isEmpty(); + if (missing) { + itemToTransform.put(key, AttributeValue.builder().s(UUID.randomUUID().toString()).build()); + } + } + + /** + * Static helpers used by the {@code @BeanTableSchemaAttributeTag}-based annotation tag. + */ + public static final class AttributeTags { + private AttributeTags() { + } + + /** + * @return a {@link StaticAttributeTag} that marks the attribute for auto-generated key behavior. + */ + public static StaticAttributeTag autoGeneratedKeyAttribute() { + return AUTO_GENERATED_KEY_ATTRIBUTE; + } + } + + /** + * Stateless builder. + */ + public static final class Builder { + private Builder() { + } + + public AutoGeneratedKeyExtension build() { + return new AutoGeneratedKeyExtension(); + } + } + + /** + * Validates Java type and records the tagged attribute into table metadata so {@link #beforeWrite} can find it at runtime. + */ + private static final class AutoGeneratedKeyAttribute implements StaticAttributeTag { + + @Override + public void validateType(String attributeName, + EnhancedType type, + AttributeValueType attributeValueType) { + + Validate.notNull(type, "type is null"); + Validate.notNull(type.rawClass(), "rawClass is null"); + Validate.notNull(attributeValueType, "attributeValueType is null"); + + if (!type.rawClass().equals(String.class)) { + throw new IllegalArgumentException(String.format( + "Attribute '%s' of Class type %s is not a suitable Java Class type to be used as a Auto Generated " + + "Key attribute. Only String Class type is supported.", attributeName, type.rawClass())); + } + } + + @Override + public Consumer modifyMetadata(String attributeName, + AttributeValueType attributeValueType) { + // Record the names of the attributes annotated with @DynamoDbAutoGeneratedKey for later lookup in beforeWrite() + return metadata -> metadata.addCustomMetadataObject( + CUSTOM_METADATA_KEY, Collections.singleton(attributeName)); + } + } +} diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/extensions/annotations/DynamoDbAutoGeneratedKey.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/extensions/annotations/DynamoDbAutoGeneratedKey.java new file mode 100644 index 000000000000..6f9892035751 --- /dev/null +++ b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/extensions/annotations/DynamoDbAutoGeneratedKey.java @@ -0,0 +1,72 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.enhanced.dynamodb.extensions.annotations; + +import java.lang.annotation.Documented; +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; +import software.amazon.awssdk.annotations.SdkPublicApi; +import software.amazon.awssdk.enhanced.dynamodb.internal.extensions.AutoGeneratedKeyTag; +import software.amazon.awssdk.enhanced.dynamodb.mapper.UpdateBehavior; +import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.BeanTableSchemaAttributeTag; +import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbUpdateBehavior; + +/** + * Annotation that marks a string attribute to be automatically populated with a random UUID if no value is provided during a + * write operation (put or update). + * + *

This annotation is designed for use with the V2 {@link software.amazon.awssdk.enhanced.dynamodb.mapper.BeanTableSchema}. + * It is registered via {@link BeanTableSchemaAttributeTag} and its behavior is implemented by + * {@link software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension}.

+ * + *

Where this annotation can be applied

+ * This annotation is only valid on attributes that serve as keys: + * + * If applied to any other attribute, the {@code AutoGeneratedKeyExtension} will throw an + * {@link IllegalArgumentException} at runtime. + * + *

How values are generated

+ * + * + *

Controlling regeneration on update

+ * This annotation can be combined with {@link DynamoDbUpdateBehavior} to control whether a new + * UUID should be generated on each update: + * + * + *

Type restriction

+ * This annotation is only valid on attributes of type {@link String}. + */ +@SdkPublicApi +@Documented +@Retention(RetentionPolicy.RUNTIME) +@Target({ElementType.METHOD, ElementType.FIELD}) +@BeanTableSchemaAttributeTag(AutoGeneratedKeyTag.class) +public @interface DynamoDbAutoGeneratedKey { +} \ No newline at end of file diff --git a/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/AutoGeneratedKeyTag.java b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/AutoGeneratedKeyTag.java new file mode 100644 index 000000000000..815e15bccd7a --- /dev/null +++ b/services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/AutoGeneratedKeyTag.java @@ -0,0 +1,33 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.enhanced.dynamodb.internal.extensions; + +import software.amazon.awssdk.annotations.SdkInternalApi; +import software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension; +import software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbAutoGeneratedKey; +import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag; + +@SdkInternalApi +public final class AutoGeneratedKeyTag { + + private AutoGeneratedKeyTag() { + } + + public static StaticAttributeTag attributeTagFor(DynamoDbAutoGeneratedKey annotation) { + return AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute(); + } + +} \ No newline at end of file diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/UuidTestUtils.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/UuidTestUtils.java new file mode 100644 index 000000000000..7276bf5639da --- /dev/null +++ b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/UuidTestUtils.java @@ -0,0 +1,30 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://aws.amazon.com/apache2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + */ + +package software.amazon.awssdk.enhanced.dynamodb; + +import java.util.UUID; + +public class UuidTestUtils { + + public static boolean isValidUuid(String uuid) { + try { + UUID.fromString(uuid); + return true; + } catch (Exception e) { + return false; + } + } +} diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/extensions/AutoGeneratedKeyExtensionTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/extensions/AutoGeneratedKeyExtensionTest.java new file mode 100644 index 000000000000..51a48fd223a0 --- /dev/null +++ b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/extensions/AutoGeneratedKeyExtensionTest.java @@ -0,0 +1,315 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * Licensed under the Apache License, Version 2.0 (the "License"). + */ + +package software.amazon.awssdk.enhanced.dynamodb.extensions; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatExceptionOfType; +import static software.amazon.awssdk.enhanced.dynamodb.UuidTestUtils.isValidUuid; +import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.primaryPartitionKey; +import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.secondaryPartitionKey; +import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.secondarySortKey; + +import java.util.Map; +import java.util.Objects; +import java.util.UUID; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.enhanced.dynamodb.OperationContext; +import software.amazon.awssdk.enhanced.dynamodb.TableMetadata; +import software.amazon.awssdk.enhanced.dynamodb.internal.extensions.DefaultDynamoDbExtensionContext; +import software.amazon.awssdk.enhanced.dynamodb.internal.operations.DefaultOperationContext; +import software.amazon.awssdk.enhanced.dynamodb.internal.operations.OperationName; +import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableSchema; +import software.amazon.awssdk.services.dynamodb.model.AttributeValue; + +public class AutoGeneratedKeyExtensionTest { + + private static final String RECORD_ID = "id123"; + private static final String TABLE_NAME = "table-name"; + + private static final OperationContext PRIMARY_CONTEXT = + DefaultOperationContext.create(TABLE_NAME, TableMetadata.primaryIndexName()); + + private final AutoGeneratedKeyExtension extension = AutoGeneratedKeyExtension.builder().build(); + + /** + * Schema that places @DynamoDbAutoGeneratedKey on GSI key ("keyAttribute") so the validation passes. + */ + private static final StaticTableSchema ITEM_WITH_KEY_SCHEMA = + StaticTableSchema.builder(ItemWithKey.class) + .newItemSupplier(ItemWithKey::new) + .addAttribute(String.class, a -> a.name("id") + .getter(ItemWithKey::getId) + .setter(ItemWithKey::setId) + .addTag(primaryPartitionKey()) // PK + .addTag(AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())) + .addAttribute(String.class, a -> a.name("keyAttribute") + .getter(ItemWithKey::getKeyAttribute) + .setter(ItemWithKey::setKeyAttribute) + .tags(secondaryPartitionKey("gsi_keys_only"), // GSI + AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())) + .addAttribute(String.class, a -> a.name("simpleString") + .getter(ItemWithKey::getSimpleString) + .setter(ItemWithKey::setSimpleString)) + .build(); + + /** + * Schema that places @DynamoDbAutoGeneratedKey on a NON-KEY attribute to trigger the exception. + */ + private static final StaticTableSchema INVALID_NONKEY_AUTOGEN_SCHEMA = + StaticTableSchema.builder(ItemWithKey.class) + .newItemSupplier(ItemWithKey::new) + .addAttribute(String.class, a -> a.name("id") + .getter(ItemWithKey::getId) + .setter(ItemWithKey::setId) + .addTag(primaryPartitionKey())) + .addAttribute(String.class, a -> a.name("keyAttribute") + .getter(ItemWithKey::getKeyAttribute) + .setter(ItemWithKey::setKeyAttribute) + // No index tags here — autogen on non-key fails at beforeWrite() + .addTag(AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())) + .addAttribute(String.class, a -> a.name("simpleString") + .getter(ItemWithKey::getSimpleString) + .setter(ItemWithKey::setSimpleString)) + .build(); + + /** + * Schema that places @DynamoDbAutoGeneratedKey on LSI key ("simpleString") so the validation passes. + */ + private static final StaticTableSchema LSI_SK_AUTOGEN_SCHEMA = + StaticTableSchema.builder(ItemWithKey.class) + .newItemSupplier(ItemWithKey::new) + .addAttribute(String.class, a -> a.name("id") + .getter(ItemWithKey::getId) + .setter(ItemWithKey::setId) + .addTag(primaryPartitionKey())) + .addAttribute(String.class, a -> a.name("keyAttribute") + .getter(ItemWithKey::getKeyAttribute) + .setter(ItemWithKey::setKeyAttribute)) + .addAttribute(String.class, a -> a.name("simpleString") + .getter(ItemWithKey::getSimpleString) + .setter(ItemWithKey::setSimpleString) + .tags(secondarySortKey("lsi1"), // LSI + AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())) + .build(); + + @Test + public void updateItem_withExistingKey_preservesValueAndDoesNotGenerateNew() { + ItemWithKey item = new ItemWithKey(); + item.setId(RECORD_ID); + String preset = UUID.randomUUID().toString(); + item.setKeyAttribute(preset); + + Map items = ITEM_WITH_KEY_SCHEMA.itemToMap(item, true); + assertThat(items).hasSize(2); + + WriteModification result = + extension.beforeWrite(DefaultDynamoDbExtensionContext.builder() + .items(items) + .tableMetadata(ITEM_WITH_KEY_SCHEMA.tableMetadata()) + .operationName(OperationName.UPDATE_ITEM) + .operationContext(PRIMARY_CONTEXT) + .build()); + + Map transformed = result.transformedItem(); + assertThat(transformed).isNotNull().hasSize(2); + assertThat(transformed).containsEntry("id", AttributeValue.fromS(RECORD_ID)); + + // Ensures the attribute remains a valid UUID without altering the preset value + assertThat(isValidUuid(transformed.get("keyAttribute").s())).isTrue(); + assertThat(result.updateExpression()).isNull(); + } + + @Test + public void updateItem_withoutKey_generatesNewUuid() { + ItemWithKey item = new ItemWithKey(); + item.setId(RECORD_ID); + + Map items = ITEM_WITH_KEY_SCHEMA.itemToMap(item, true); + assertThat(items).hasSize(1); + + WriteModification result = + extension.beforeWrite(DefaultDynamoDbExtensionContext.builder() + .items(items) + .tableMetadata(ITEM_WITH_KEY_SCHEMA.tableMetadata()) + .operationName(OperationName.UPDATE_ITEM) + .operationContext(PRIMARY_CONTEXT) + .build()); + + Map transformed = result.transformedItem(); + assertThat(transformed).isNotNull().hasSize(2); + assertThat(transformed).containsEntry("id", AttributeValue.fromS(RECORD_ID)); + assertThat(isValidUuid(transformed.get("keyAttribute").s())).isTrue(); + assertThat(result.updateExpression()).isNull(); + } + + @Test + public void updateItem_withMissingKeyAttribute_insertsGeneratedUuid() { + ItemWithKey item = new ItemWithKey(); + item.setId(RECORD_ID); + + Map items = ITEM_WITH_KEY_SCHEMA.itemToMap(item, true); + assertThat(items).hasSize(1); + + WriteModification result = + extension.beforeWrite(DefaultDynamoDbExtensionContext.builder() + .items(items) + .tableMetadata(ITEM_WITH_KEY_SCHEMA.tableMetadata()) + .operationName(OperationName.UPDATE_ITEM) + .operationContext(PRIMARY_CONTEXT) + .build()); + + assertThat(result.transformedItem()).isNotNull(); + assertThat(result.updateExpression()).isNull(); + assertThat(result.transformedItem()).hasSize(2); + assertThat(isValidUuid(result.transformedItem().get("keyAttribute").s())).isTrue(); + } + + @Test + public void nonStringTypeAnnotatedWithAutoGeneratedKey_throwsIllegalArgumentException() { + assertThatExceptionOfType(IllegalArgumentException.class) + .isThrownBy(() -> + StaticTableSchema.builder(ItemWithKey.class) + .newItemSupplier(ItemWithKey::new) + .addAttribute(String.class, a -> a.name("id") + .getter(ItemWithKey::getId) + .setter(ItemWithKey::setId) + .addTag(primaryPartitionKey())) + .addAttribute(Integer.class, a -> a.name("intAttribute") + .getter(ItemWithKey::getIntAttribute) + .setter(ItemWithKey::setIntAttribute) + .addTag(AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())) + .addAttribute(String.class, a -> a.name("simpleString") + .getter(ItemWithKey::getSimpleString) + .setter(ItemWithKey::setSimpleString)) + .build() + ) + .withMessage("Attribute 'intAttribute' of Class type class java.lang.Integer is not a suitable Java Class type " + + "to be used as a Auto Generated Key attribute. Only String Class type is supported."); + } + + @Test + public void autoGeneratedKey_onSecondaryPartitionKey_generatesUuid() { + ItemWithKey item = new ItemWithKey(); + item.setId(RECORD_ID); // keyAttribute (GSI PK) is missing → should be generated + + Map items = ITEM_WITH_KEY_SCHEMA.itemToMap(item, true); + + WriteModification result = + extension.beforeWrite(DefaultDynamoDbExtensionContext.builder() + .items(items) + .tableMetadata(ITEM_WITH_KEY_SCHEMA.tableMetadata()) + .operationName(OperationName.PUT_ITEM) + .operationContext(PRIMARY_CONTEXT) + .build()); + + Map transformed = result.transformedItem(); + assertThat(transformed).isNotNull(); + assertThat(isValidUuid(transformed.get("keyAttribute").s())).isTrue(); // generated for GSI PK + } + + @Test + public void autoGeneratedKey_onSecondarySortKey_generatesUuid() { + ItemWithKey item = new ItemWithKey(); + item.setId(RECORD_ID); // simpleString (GSI/LSI) is missing → should be generated + + Map items = LSI_SK_AUTOGEN_SCHEMA.itemToMap(item, true); + + WriteModification result = + extension.beforeWrite(DefaultDynamoDbExtensionContext.builder() + .items(items) + .tableMetadata(LSI_SK_AUTOGEN_SCHEMA.tableMetadata()) + .operationName(OperationName.PUT_ITEM) + .operationContext(PRIMARY_CONTEXT) + .build()); + + Map transformed = result.transformedItem(); + assertThat(transformed).isNotNull(); + assertThat(isValidUuid(transformed.get("simpleString").s())).isTrue(); // generated for index SK + } + + @Test + public void autoGeneratedKey_onNonKey_throwsIllegalArgumentException() { + ItemWithKey item = new ItemWithKey(); + item.setId(RECORD_ID); // keyAttribute is annotated but NOT a key in this schema → should fail at beforeWrite() + + Map items = INVALID_NONKEY_AUTOGEN_SCHEMA.itemToMap(item, true); + + assertThatExceptionOfType(IllegalArgumentException.class) + .isThrownBy(() -> + extension.beforeWrite(DefaultDynamoDbExtensionContext.builder() + .items(items) + .tableMetadata(INVALID_NONKEY_AUTOGEN_SCHEMA.tableMetadata()) + .operationName(OperationName.PUT_ITEM) + .operationContext(PRIMARY_CONTEXT) + .build()) + ) + .withMessageContaining("@DynamoDbAutoGeneratedKey can only be applied to key attributes: " + + "primary partition key, primary sort key, or GSI/LSI partition/sort keys.") + .withMessageContaining("keyAttribute"); + } + + private static class ItemWithKey { + + private String id; + private String keyAttribute; + private String simpleString; + private Integer intAttribute; + + ItemWithKey() { + } + + public Integer getIntAttribute() { + return intAttribute; + } + + public void setIntAttribute(Integer intAttribute) { + this.intAttribute = intAttribute; + } + + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public String getKeyAttribute() { + return keyAttribute; + } + + public void setKeyAttribute(String keyAttribute) { + this.keyAttribute = keyAttribute; + } + + public String getSimpleString() { + return simpleString; + } + + public void setSimpleString(String simpleString) { + this.simpleString = simpleString; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof ItemWithKey)) { + return false; + } + ItemWithKey that = (ItemWithKey) o; + return Objects.equals(id, that.id) + && Objects.equals(keyAttribute, that.keyAttribute) + && Objects.equals(simpleString, that.simpleString); + } + + @Override + public int hashCode() { + return Objects.hash(id, keyAttribute, simpleString); + } + } +} diff --git a/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/AutoGeneratedKeyRecordTest.java b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/AutoGeneratedKeyRecordTest.java new file mode 100644 index 000000000000..cc238b18ea25 --- /dev/null +++ b/services-custom/dynamodb-enhanced/src/test/java/software/amazon/awssdk/enhanced/dynamodb/functionaltests/AutoGeneratedKeyRecordTest.java @@ -0,0 +1,535 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. + * Licensed under the Apache License, Version 2.0. + */ + +package software.amazon.awssdk.enhanced.dynamodb.functionaltests; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.junit.Assert.assertEquals; +import static software.amazon.awssdk.enhanced.dynamodb.functionaltests.AutoGeneratedUuidRecordTest.assertValidUuid; +import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.primaryPartitionKey; +import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.secondaryPartitionKey; +import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.secondarySortKey; +import static software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags.updateBehavior; + +import java.util.Arrays; +import java.util.Collection; +import java.util.Objects; +import org.assertj.core.api.Assertions; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.Parameterized; +import org.junit.runners.Parameterized.Parameters; +import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient; +import software.amazon.awssdk.enhanced.dynamodb.DynamoDbTable; +import software.amazon.awssdk.enhanced.dynamodb.TableSchema; +import software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedKeyExtension; +import software.amazon.awssdk.enhanced.dynamodb.extensions.VersionedRecordExtension; +import software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbAutoGeneratedKey; +import software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbVersionAttribute; +import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableSchema; +import software.amazon.awssdk.enhanced.dynamodb.mapper.UpdateBehavior; +import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbBean; +import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbFlatten; +import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbPartitionKey; +import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSecondaryPartitionKey; +import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbSecondarySortKey; +import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbUpdateBehavior; + +@RunWith(Parameterized.class) +public class AutoGeneratedKeyRecordTest extends LocalDynamoDbSyncTestBase { + + private final String tableName = getConcreteTableName("AutoGenKey-table"); + private final String versionedTableName = getConcreteTableName("AutoGenKey-versioned-table"); + private final DynamoDbTable mappedTable; + + public AutoGeneratedKeyRecordTest(String testName, + TableSchema schema) { + this.mappedTable = DynamoDbEnhancedClient.builder() + .dynamoDbClient(getDynamoDbClient()) + .extensions(AutoGeneratedKeyExtension.builder().build()) + .build() + .table(tableName, schema); + } + + /* + Flattened bean schema: "generated" attribute is a GSI PK + @DynamoDbAutoGeneratedKey annotation + */ + private static final TableSchema FLATTENED = + StaticTableSchema.builder(FlattenedRecord.class) + .newItemSupplier(FlattenedRecord::new) + .addAttribute(String.class, a -> a.name("generated") + .getter(FlattenedRecord::getGenerated) + .setter(FlattenedRecord::generated) + .tags(secondaryPartitionKey("gsi_flat"), + AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())) + .build(); + + /** + * - id: PK + * - lastUpdatedKey: GSI PK + @DynamoDbAutoGeneratedKey annotation (UpdateBehaviour is WRITE_ALWAYS - default value) + * - createdKey: GSI SK + @DynamoDbAutoGeneratedKey annotation (UpdateBehaviour is WRITE_IF_NOT_EXISTS) + */ + private static final TableSchema STATIC_SCHEMA = + StaticTableSchema.builder(RecordWithMixedUpdateBehaviours.class) + .newItemSupplier(RecordWithMixedUpdateBehaviours::new) + .addAttribute(String.class, a -> a.name("id") + .getter(RecordWithMixedUpdateBehaviours::getId) + .setter(RecordWithMixedUpdateBehaviours::id) + .addTag(primaryPartitionKey())) + .addAttribute(String.class, a -> a.name("attribute") + .getter(RecordWithMixedUpdateBehaviours::getAttribute) + .setter(RecordWithMixedUpdateBehaviours::attribute)) + .addAttribute(String.class, a -> a.name("lastUpdatedKey") + .getter(RecordWithMixedUpdateBehaviours::getLastUpdatedKey) + .setter(RecordWithMixedUpdateBehaviours::lastUpdatedKey) + .tags(secondaryPartitionKey("gsi_auto"), + AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())) + .addAttribute(String.class, a -> a.name("createdKey") + .getter(RecordWithMixedUpdateBehaviours::getCreatedKey) + .setter(RecordWithMixedUpdateBehaviours::createdKey) + .tags(secondarySortKey("gsi_auto"), + AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute(), + updateBehavior(UpdateBehavior.WRITE_IF_NOT_EXISTS))) + .flatten(FLATTENED, + RecordWithMixedUpdateBehaviours::getFlattened, + RecordWithMixedUpdateBehaviours::flattened) + .build(); + + @Parameters(name = "{index}: {0}") + public static Collection data() { + return Arrays.asList(new Object[][] { + {"StaticTableSchema", STATIC_SCHEMA}, + {"BeanTableSchema", TableSchema.fromBean(RecordWithMixedUpdateBehaviours.class)} + }); + } + + @Before + public void createTable() { + mappedTable.createTable(r -> r.provisionedThroughput(getDefaultProvisionedThroughput())); + } + + @After + public void deleteTable() { + getDynamoDbClient().deleteTable(b -> b.tableName(tableName)); + } + + @Test + public void put_generatesAllAnnotatedKeys() { + mappedTable.putItem(r -> r.item(new RecordWithMixedUpdateBehaviours() + .id("id").attribute("one") + .flattened(new FlattenedRecord()))); + + RecordWithMixedUpdateBehaviours out = mappedTable.getItem(r -> r.key(k -> k.partitionValue("id"))); + assertValidUuid(out.getCreatedKey()); + assertValidUuid(out.getLastUpdatedKey()); + assertValidUuid(out.getFlattened().getGenerated()); + assertThat(out.getCreatedKey()).isNotEqualTo(out.getLastUpdatedKey()); + } + + @Test + public void update_preservesKeysWithWriteIfNotExists_and_regeneratesKeysWithWriteAlways() { + mappedTable.putItem(r -> r.item(new RecordWithMixedUpdateBehaviours() + .id("id").attribute("one") + .flattened(new FlattenedRecord()))); + + RecordWithMixedUpdateBehaviours afterPut = mappedTable.getItem(r -> r.key(k -> k.partitionValue("id"))); + String createdAtPut = afterPut.getCreatedKey(); // WRITE_IF_NOT_EXISTS → should keep + String lastUpdatedAtPut = afterPut.getLastUpdatedKey(); // WRITE_ALWAYS (default) → should change + String flatGeneratedAtPut = afterPut.getFlattened().getGenerated(); // WRITE_ALWAYS (default) → should change + + mappedTable.updateItem(r -> r.item(new RecordWithMixedUpdateBehaviours() + .id("id").attribute("two") + .flattened(new FlattenedRecord()))); + + RecordWithMixedUpdateBehaviours afterUpdate = mappedTable.getItem(r -> r.key(k -> k.partitionValue("id"))); + assertThat(afterUpdate.getCreatedKey()).isEqualTo(createdAtPut); // preserved + assertThat(afterUpdate.getLastUpdatedKey()).isNotEqualTo(lastUpdatedAtPut); + assertThat(afterUpdate.getFlattened().getGenerated()).isNotEqualTo(flatGeneratedAtPut); + } + + @Test + public void autogenerateKey_setOnGsiPartitionKey_performsUuidGeneration() { + String tn = getConcreteTableName("autogen-gsi-pk"); + DynamoDbEnhancedClient client = DynamoDbEnhancedClient.builder() + .dynamoDbClient(getDynamoDbClient()) + .extensions(AutoGeneratedKeyExtension.builder().build()) + .build(); + + TableSchema schema = StaticTableSchema.builder(GsiPartitionKeyBean.class) + .newItemSupplier(GsiPartitionKeyBean::new) + .addAttribute(String.class, a -> a.name("id") + .getter(GsiPartitionKeyBean::getId) + .setter(GsiPartitionKeyBean::setId) + .addTag(primaryPartitionKey())) + .addAttribute(String.class, a -> a.name("gsiPk") + .getter(GsiPartitionKeyBean::getGsiPk) + .setter(GsiPartitionKeyBean::setGsiPk) + .tags(secondaryPartitionKey("gsi1"), + AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())) + .build(); + + DynamoDbTable table = client.table(tn, schema); + try { + table.createTable(r -> r.provisionedThroughput(getDefaultProvisionedThroughput())); + GsiPartitionKeyBean bean = new GsiPartitionKeyBean(); + bean.setId("id123"); // gsiPk missing → should be generated + table.putItem(bean); + + GsiPartitionKeyBean out = table.getItem(r -> r.key(k -> k.partitionValue("id123"))); + assertValidUuid(out.getGsiPk()); + } finally { + getDynamoDbClient().deleteTable(b -> b.tableName(tn)); + } + } + + @Test + public void autogenerateKey_setOnSecondaryIndexSortKey_performsUuidGeneration() { + String tn = getConcreteTableName("autogen-idx-sk"); + DynamoDbEnhancedClient client = DynamoDbEnhancedClient.builder() + .dynamoDbClient(getDynamoDbClient()) + .extensions(AutoGeneratedKeyExtension.builder().build()) + .build(); + + TableSchema schema = StaticTableSchema.builder(SecondaryIndexSortKeyBean.class) + .newItemSupplier(SecondaryIndexSortKeyBean::new) + .addAttribute(String.class, a -> a.name("id") + .getter(SecondaryIndexSortKeyBean::getId) + .setter(SecondaryIndexSortKeyBean::setId) + .addTag(primaryPartitionKey())) + .addAttribute(String.class, a -> a.name( + "secondarySortKey") + .getter(SecondaryIndexSortKeyBean::getSecondarySortKey) + .setter(SecondaryIndexSortKeyBean::setSecondarySortKey) + .tags(secondarySortKey("lsi1"), + AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())) + .build(); + + DynamoDbTable table = client.table(tn, schema); + try { + table.createTable(r -> r.provisionedThroughput(getDefaultProvisionedThroughput())); + SecondaryIndexSortKeyBean bean = new SecondaryIndexSortKeyBean(); + bean.setId("id123"); // secondarySortKey missing → should be generated + table.putItem(bean); + + SecondaryIndexSortKeyBean out = table.getItem(r -> r.key(k -> k.partitionValue("id123"))); + assertValidUuid(out.getSecondarySortKey()); + } finally { + getDynamoDbClient().deleteTable(b -> b.tableName(tn)); + } + } + + @Test + public void autogenerateKey_setOnInvalidKeyAttribute_throws_exception() { + String tn = getConcreteTableName("autogen-non-key"); + DynamoDbEnhancedClient client = DynamoDbEnhancedClient.builder() + .dynamoDbClient(getDynamoDbClient()) + .extensions(AutoGeneratedKeyExtension.builder().build()) + .build(); + + TableSchema schema = StaticTableSchema.builder(NonKeyBean.class) + .newItemSupplier(NonKeyBean::new) + .addAttribute(String.class, a -> a.name("id") + .getter(NonKeyBean::getId) + .setter(NonKeyBean::setId) + .addTag(primaryPartitionKey())) + .addAttribute(String.class, a -> a.name("notAKey") + .getter(NonKeyBean::getNotAKey) + .setter(NonKeyBean::setNotAKey) + .addTag(AutoGeneratedKeyExtension.AttributeTags.autoGeneratedKeyAttribute())) + .build(); + + DynamoDbTable table = client.table(tn, schema); + try { + table.createTable(r -> r.provisionedThroughput(getDefaultProvisionedThroughput())); + NonKeyBean bean = new NonKeyBean(); + bean.setId("id123"); + Assertions.assertThatExceptionOfType(IllegalArgumentException.class) + .isThrownBy(() -> table.putItem(bean)) + .withMessageContaining("@DynamoDbAutoGeneratedKey can only be applied to key attributes: " + + "primary partition key, primary sort key, or GSI/LSI partition/sort keys.") + .withMessageContaining("notAKey"); + } finally { + getDynamoDbClient().deleteTable(b -> b.tableName(tn)); + } + } + + @Test + public void autogenerateKey_onVersionedRecord_setOnPrimaryKey_performsUuidGeneration() { + DynamoDbEnhancedClient client = DynamoDbEnhancedClient.builder() + .dynamoDbClient(getDynamoDbClient()) + .extensions(AutoGeneratedKeyExtension.builder().build(), + VersionedRecordExtension.builder().build()) + .build(); + + TableSchema schema = TableSchema.fromBean(VersionedRecord.class); + DynamoDbTable table = client.table(versionedTableName, schema); + + try { + table.createTable(r -> r.provisionedThroughput(getDefaultProvisionedThroughput())); + + VersionedRecord rec = new VersionedRecord(); + rec.setPayload("payload"); + table.putItem(rec); + + VersionedRecord existing = + table.scan().items().stream().findFirst() + .orElseThrow(() -> new AssertionError("No record found in table")); + + assertEquals("payload", existing.getPayload()); + assertThat(existing.getVersion()).isEqualTo(1L); + + existing.setPayload("new_payload"); + VersionedRecord updated = table.updateItem(existing); + + assertEquals("new_payload", updated.getPayload()); + assertThat(updated.getVersion()).isEqualTo(2L); + } finally { + getDynamoDbClient().deleteTable(b -> b.tableName(versionedTableName)); + } + } + + /** + * - createdKey: GSI SK + @DynamoDbAutoGeneratedKey annotation (UpdateBehaviour is WRITE_IF_NOT_EXISTS) + * - lastUpdatedKey: GSI PK + @DynamoDbAutoGeneratedKey annotation (UpdateBehaviour is WRITE_ALWAYS - default value) + * - flattened.generated: GSI PK + @DynamoDbAutoGeneratedKey + */ + @DynamoDbBean + public static class RecordWithMixedUpdateBehaviours { + private String id; + private String attribute; + private String createdKey; + private String lastUpdatedKey; + private FlattenedRecord flattened; + + @DynamoDbPartitionKey + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public RecordWithMixedUpdateBehaviours id(String v) { + this.id = v; + return this; + } + + public String getAttribute() { + return attribute; + } + + public void setAttribute(String attribute) { + this.attribute = attribute; + } + + public RecordWithMixedUpdateBehaviours attribute(String v) { + this.attribute = v; + return this; + } + + @DynamoDbAutoGeneratedKey + @DynamoDbSecondarySortKey(indexNames = "gsi_auto") + @DynamoDbUpdateBehavior(UpdateBehavior.WRITE_IF_NOT_EXISTS) + public String getCreatedKey() { + return createdKey; + } + + public void setCreatedKey(String createdKey) { + this.createdKey = createdKey; + } + + public RecordWithMixedUpdateBehaviours createdKey(String v) { + this.createdKey = v; + return this; + } + + @DynamoDbAutoGeneratedKey + @DynamoDbSecondaryPartitionKey(indexNames = "gsi_auto") + public String getLastUpdatedKey() { + return lastUpdatedKey; + } + + public void setLastUpdatedKey(String lastUpdatedKey) { + this.lastUpdatedKey = lastUpdatedKey; + } + + public RecordWithMixedUpdateBehaviours lastUpdatedKey(String v) { + this.lastUpdatedKey = v; + return this; + } + + @DynamoDbFlatten + public FlattenedRecord getFlattened() { + return flattened; + } + + public void setFlattened(FlattenedRecord fr) { + this.flattened = fr; + } + + public RecordWithMixedUpdateBehaviours flattened(FlattenedRecord fr) { + this.flattened = fr; + return this; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof RecordWithMixedUpdateBehaviours)) { + return false; + } + RecordWithMixedUpdateBehaviours that = (RecordWithMixedUpdateBehaviours) o; + return Objects.equals(id, that.id) + && Objects.equals(attribute, that.attribute) + && Objects.equals(createdKey, that.createdKey) + && Objects.equals(lastUpdatedKey, that.lastUpdatedKey) + && Objects.equals(flattened, that.flattened); + } + + @Override + public int hashCode() { + return Objects.hash(id, attribute, createdKey, lastUpdatedKey, flattened); + } + } + + @DynamoDbBean + public static class FlattenedRecord { + private String generated; + + @DynamoDbAutoGeneratedKey + @DynamoDbSecondaryPartitionKey(indexNames = "gsi_flat") + public String getGenerated() { + return generated; + } + + public void setGenerated(String generated) { + this.generated = generated; + } + + public FlattenedRecord generated(String g) { + this.generated = g; + return this; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof FlattenedRecord)) { + return false; + } + FlattenedRecord that = (FlattenedRecord) o; + return Objects.equals(generated, that.generated); + } + + @Override + public int hashCode() { + return Objects.hash(generated); + } + } + + public static class GsiPartitionKeyBean { + private String id; + private String gsiPk; + + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public String getGsiPk() { + return gsiPk; + } + + public void setGsiPk(String gsiPk) { + this.gsiPk = gsiPk; + } + } + + public static class SecondaryIndexSortKeyBean { + private String id; + private String secondarySortKey; + + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public String getSecondarySortKey() { + return secondarySortKey; + } + + public void setSecondarySortKey(String secondarySortKey) { + this.secondarySortKey = secondarySortKey; + } + } + + public static class NonKeyBean { + private String id; + private String notAKey; + + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + public String getNotAKey() { + return notAKey; + } + + public void setNotAKey(String notAKey) { + this.notAKey = notAKey; + } + } + + @DynamoDbBean + public static class VersionedRecord { + private String id; + private Long version; + private String payload; + + @DynamoDbPartitionKey + @DynamoDbAutoGeneratedKey + public String getId() { + return id; + } + + public void setId(String id) { + this.id = id; + } + + @DynamoDbVersionAttribute + public Long getVersion() { + return version; + } + + public void setVersion(Long version) { + this.version = version; + } + + public String getPayload() { + return payload; + } + + public void setPayload(String payload) { + this.payload = payload; + } + } +}