Added support for DynamoDbAutoGeneratedKey annotation

.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"
+}

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).
+ * <p>
+ * The annotation may be placed <b>only</b> on key attributes:
+ * <ul>
+ *   <li>Primary partition key (PK) or primary sort key (SK)</li>
+ *   <li>Partition key or sort key of any secondary index (GSI or LSI)</li>
+ * </ul>
+ *
+ * <p><b>Validation:</b> 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.</p>
+ */
+@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.
+     * <p>
+     * Also validates that the annotation is only used on PK/SK/GSI/LSI key attributes.
+     */
+    @Override
+    public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
+        Collection<String> 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<String> 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<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
+        taggedAttributes.forEach(attr -> insertUuidIfMissing(itemToTransform, attr));
+
+        return WriteModification.builder()
+                                .transformedItem(Collections.unmodifiableMap(itemToTransform))
+                                .build();
+    }
+
+    private void insertUuidIfMissing(Map<String, AttributeValue> 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 <R> void validateType(String attributeName,
+                                     EnhancedType<R> type,
+                                     AttributeValueType attributeValueType) {
+
+            Validate.notNull(type, "type is null");
+            Validate.notNull(type.rawClass(), "rawClass is null");
+            Validate.notNull(attributeValueType, "attributeValueType is null");
+
+            if (!String.class.equals(type.rawClass())) {
+                throw new IllegalArgumentException(String.format(
+                    "Attribute '%s' of Java type %s is not valid for @DynamoDbAutoGeneratedKey. Only String is supported.",
+                    attributeName, type.rawClass()));
+            }
+        }
+
+        @Override
+        public Consumer<StaticTableMetadata.Builder> 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));
+        }
+    }
+}

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).
+ *
+ * <p>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}.</p>
+ *
+ * <h3>Where this annotation can be applied</h3>
+ * This annotation is only valid on attributes that serve as keys:
+ * <ul>
+ *   <li>The table’s primary partition key or sort key</li>
+ *   <li>The partition key or sort key of a secondary index (GSI or LSI)</li>
+ * </ul>
+ * If applied to any other attribute, the {@code AutoGeneratedKeyExtension} will throw an
+ * {@link IllegalArgumentException} at runtime.
+ *
+ * <h3>How values are generated</h3>
+ * <ul>
+ *   <li>On writes where the annotated attribute is null or empty, a new UUID value is generated
+ *       using {@link java.util.UUID#randomUUID()}.</li>
+ *   <li>If a value is already set on the attribute, that value is preserved and not replaced.</li>
+ * </ul>
+ *
+ * <h3>Controlling regeneration on update</h3>
+ * This annotation can be combined with {@link DynamoDbUpdateBehavior} to control whether a new
+ * UUID should be generated on each update:
+ * <ul>
+ *   <li>{@link UpdateBehavior#WRITE_ALWAYS} (default) –
+ *       Generate a new UUID whenever the attribute is missing during write.</li>
+ *   <li>{@link UpdateBehavior#WRITE_IF_NOT_EXISTS} –
+ *       Generate a UUID only the first time (on insert), and preserve that value on subsequent updates.</li>
+ * </ul>
+ *
+ * <h3>Type restriction</h3>
+ * 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 {
+}

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();
+    }
+
+}

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