2121import java .util .Map ;
2222import java .util .UUID ;
2323import java .util .function .Consumer ;
24+ import java .util .function .Supplier ;
25+ import software .amazon .awssdk .annotations .NotThreadSafe ;
2426import software .amazon .awssdk .annotations .SdkPublicApi ;
2527import software .amazon .awssdk .annotations .ThreadSafe ;
2628import software .amazon .awssdk .enhanced .dynamodb .AttributeValueType ;
3335import software .amazon .awssdk .services .dynamodb .model .AttributeValue ;
3436import software .amazon .awssdk .utils .Validate ;
3537
36-
3738/**
3839 * This extension facilitates the automatic generation of a unique UUID (Universally Unique Identifier) for a specified attribute
3940 * every time a new record is written to the database. The generated UUID is obtained using the
40- * {@link java.util.UUID#randomUUID()} method.
41+ * {@link java.util.UUID#randomUUID()} method by default or a custom UUID supplier instance provided by the user .
4142 * <p>
4243 * This extension is not loaded by default when you instantiate a
4344 * {@link software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient}. Therefore, you need to specify it in a custom
4445 * extension when creating the enhanced client.
4546 * <p>
4647 * Example to add AutoGeneratedUuidExtension along with default extensions is
47- * {@snippet :
48- * DynamoDbEnhancedClient.builder().extensions(Stream.concat(ExtensionResolver.defaultExtensions().stream(),
49- * Stream.of(AutoGeneratedUuidExtension.create())).collect(Collectors.toList())).build();
50- *}
48+ * {@code DynamoDbEnhancedClient.builder().extensions(Stream.concat(ExtensionResolver.defaultExtensions().stream(),
49+ * Stream.of(AutoGeneratedUuidExtension.create())).collect(Collectors.toList())).build();}
5150 * </p>
5251 * <p>
5352 * Example to just add AutoGeneratedUuidExtension without default extensions is
54- * {@snippet :
55- * DynamoDbEnhancedClient.builder().extensions(AutoGeneratedUuidExtension.create()).build();
56- *}
53+ * {@code DynamoDbEnhancedClient.builder().extensions(AutoGeneratedUuidExtension.create())).build();}
5754 * </p>
5855 * <p>
5956 * To utilize the auto-generated UUID feature, first, create a field in your model that will store the UUID for the attribute.
6057 * This class field must be of type {@link java.lang.String}, and you need to tag it as the autoGeneratedUuidAttribute. If you are
6158 * using the {@link software.amazon.awssdk.enhanced.dynamodb.mapper.BeanTableSchema}, then you should use the
6259 * {@link software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbAutoGeneratedUuid} annotation. If you are using
6360 * the {@link software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableSchema}, then you should use the
64- * {@link
65- * software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedUuidExtension.AttributeTags#autoGeneratedUuidAttribute()}
66- * static attribute tag.
61+ * {@link AttributeTags#autoGeneratedUuidAttribute()} static attribute tag.
6762 * </p>
6863 * <p>
6964 * Every time a new record is successfully put into the database, the specified attribute will be automatically populated with a
70- * unique UUID generated using {@link java.util.UUID#randomUUID()}. If the UUID needs to be created only for `putItem` and should
71- * not be generated for an `updateItem`, then
65+ * unique UUID generated using {@link java.util.UUID#randomUUID()}.
66+ * <br>
67+ * If the UUID needs to be created only for `putItem` and should not be generated for an `updateItem`, then
7268 * {@link software.amazon.awssdk.enhanced.dynamodb.mapper.UpdateBehavior#WRITE_IF_NOT_EXISTS} must be along with
73- * {@link DynamoDbUpdateBehavior}
74- *
69+ * {@link DynamoDbUpdateBehavior}.
7570 * </p>
7671 */
7772@ SdkPublicApi
7873@ ThreadSafe
7974public final class AutoGeneratedUuidExtension implements DynamoDbEnhancedClientExtension {
80- private static final String CUSTOM_METADATA_KEY =
81- "software.amazon.awssdk.enhanced.dynamodb.extensions.AutoGeneratedUuidExtension:AutoGeneratedUuidAttribute" ;
82- private static final AutoGeneratedUuidAttribute AUTO_GENERATED_UUID_ATTRIBUTE = new AutoGeneratedUuidAttribute ();
75+
76+ private static final String CUSTOM_METADATA_KEY = "AutoGeneratedUuidExtension:AutoGeneratedUuidAttribute" ;
77+ private static final StaticAttributeTag AUTO_GENERATED_UUID_ATTRIBUTE = new AutoGeneratedUuidAttribute ();
78+
79+ private final Supplier <UUID > uuidSupplier ;
8380
8481 private AutoGeneratedUuidExtension () {
82+ this .uuidSupplier = UUID ::randomUUID ;
83+ }
84+
85+ private AutoGeneratedUuidExtension (Builder builder ) {
86+ this .uuidSupplier = builder .uuidSupplier == null ? UUID ::randomUUID : builder .uuidSupplier ;
87+ }
88+
89+ public static Builder builder () {
90+ return new Builder ();
91+ }
92+
93+ public Builder toBuilder () {
94+ return builder ().uuidSupplier (this .uuidSupplier );
8595 }
8696
87- /**
88- * @return an Instance of {@link AutoGeneratedUuidExtension}
89- */
9097 public static AutoGeneratedUuidExtension create () {
9198 return new AutoGeneratedUuidExtension ();
9299 }
93100
94- /**
95- * Modifies the WriteModification UUID string with the attribute updated with the extension.
96- *
97- * @param context The {@link DynamoDbExtensionContext.BeforeWrite} context containing the state of the execution.
98- * @return WriteModification String updated with attribute updated with Extension.
99- */
100101 @ Override
101102 public WriteModification beforeWrite (DynamoDbExtensionContext .BeforeWrite context ) {
102-
103-
104103 Collection <String > customMetadataObject = context .tableMetadata ()
105104 .customMetadataObject (CUSTOM_METADATA_KEY , Collection .class )
106105 .orElse (null );
@@ -116,32 +115,25 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
116115 .build ();
117116 }
118117
119- private void insertUuidInItemToTransform (Map <String , AttributeValue > itemToTransform ,
120- String key ) {
121- itemToTransform .put (key , AttributeValue .builder ().s (UUID .randomUUID ().toString ()).build ());
118+ private void insertUuidInItemToTransform (Map <String , AttributeValue > itemToTransform , String key ) {
119+ itemToTransform .put (key , AttributeValue .builder ().s (uuidSupplier .get ().toString ()).build ());
122120 }
123121
124122 public static final class AttributeTags {
125123
126124 private AttributeTags () {
127125 }
128126
129- /**
130- * Tags which indicate that the given attribute is supported wih Auto Generated UUID Record Extension.
131- *
132- * @return Tag name for AutoGenerated UUID Records
133- */
134127 public static StaticAttributeTag autoGeneratedUuidAttribute () {
135128 return AUTO_GENERATED_UUID_ATTRIBUTE ;
136129 }
130+
137131 }
138132
139133 private static class AutoGeneratedUuidAttribute implements StaticAttributeTag {
140134
141135 @ Override
142- public <R > void validateType (String attributeName , EnhancedType <R > type ,
143- AttributeValueType attributeValueType ) {
144-
136+ public <R > void validateType (String attributeName , EnhancedType <R > type , AttributeValueType attributeValueType ) {
145137 Validate .notNull (type , "type is null" );
146138 Validate .notNull (type .rawClass (), "rawClass is null" );
147139 Validate .notNull (attributeValueType , "attributeValueType is null" );
@@ -159,5 +151,34 @@ public Consumer<StaticTableMetadata.Builder> modifyMetadata(String attributeName
159151 return metadata -> metadata .addCustomMetadataObject (CUSTOM_METADATA_KEY , Collections .singleton (attributeName ))
160152 .markAttributeAsKey (attributeName , attributeValueType );
161153 }
154+
155+ }
156+
157+ @ NotThreadSafe
158+ public static final class Builder {
159+
160+ private Supplier <UUID > uuidSupplier ;
161+
162+ private Builder () {
163+ }
164+
165+ /**
166+ * Sets the UUID supplier instance , else {@link UUID#randomUUID()} is used by default. Every time a new UUID is generated
167+ * this supplier will be used to get the current UUID. If a custom supplier is not specified, the default randomUUID
168+ * supplier will be used.
169+ *
170+ * @param uuidSupplier Supplier instance to set the current UUID.
171+ * @return This builder for method chaining.
172+ */
173+ public Builder uuidSupplier (Supplier <UUID > uuidSupplier ) {
174+ this .uuidSupplier = uuidSupplier ;
175+ return this ;
176+ }
177+
178+ public AutoGeneratedUuidExtension build () {
179+ return new AutoGeneratedUuidExtension (this );
180+ }
181+
162182 }
163- }
183+
184+ }
0 commit comments