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