DDB Enhanced: Allow custom versioning

.changes/next-release/feature-DynamoDBEnhancedClient-2a501d8.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "DynamoDB Enhanced Client",
+    "contributor": "akiesler",
+    "description": "Support for Version Starting at 0 with Configurable Increment"
+}

pom.xml

@@ -682,6 +682,9 @@
                             <includeModule>polly</includeModule>
                         </includeModules>
                         <excludes>
+                            <!-- TODO remove after release -->
+                            <exclude>software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbVersionAttribute#incrementBy()</exclude>
+                            <exclude>software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbVersionAttribute#startAt()</exclude>
                             <exclude>*.internal.*</exclude>
                             <exclude>software.amazon.awssdk.thirdparty.*</exclude>
                             <exclude>software.amazon.awssdk.regions.*</exclude>

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/Expression.java

@@ -26,6 +26,7 @@
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.annotations.ThreadSafe;
 import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
+import software.amazon.awssdk.utils.ToString;
 
 /**
  * High-level representation of a DynamoDB 'expression' that can be used in various situations where the API requires
@@ -311,6 +312,15 @@ public int hashCode() {
         return result;
     }
 
+    @Override
+    public String toString() {
+        return ToString.builder("Expression")
+                       .add("expression", expression)
+                       .add("expressionValues", expressionValues)
+                       .add("expressionNames", expressionNames)
+                       .build();
+    }
+
     /**
      * A builder for {@link Expression}
      */

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

@@ -34,6 +34,7 @@
 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;
 
 /**
  * This extension implements optimistic locking on record writes by means of a 'record version number' that is used
@@ -61,7 +62,18 @@ public final class VersionedRecordExtension implements DynamoDbEnhancedClientExt
     private static final String CUSTOM_METADATA_KEY = "VersionedRecordExtension:VersionAttribute";
     private static final VersionAttribute VERSION_ATTRIBUTE = new VersionAttribute();
 
-    private VersionedRecordExtension() {
+    private final long startAt;
+    private final long incrementBy;
+
+    private VersionedRecordExtension(Long startAt, Long incrementBy) {
+        Validate.isNotNegativeOrNull(startAt, "startAt");
+
+        if (incrementBy != null && incrementBy < 1) {
+            throw new IllegalArgumentException("incrementBy must be greater than 0.");
+        }
+
+        this.startAt = startAt != null ? startAt : 0L;
+        this.incrementBy = incrementBy != null ? incrementBy : 1L;
     }
 
     public static Builder builder() {
@@ -75,19 +87,47 @@ private AttributeTags() {
         public static StaticAttributeTag versionAttribute() {
             return VERSION_ATTRIBUTE;
         }
+
+        public static StaticAttributeTag versionAttribute(Long startAt, Long incrementBy) {
+            return new VersionAttribute(startAt, incrementBy);
+        }
     }
 
-    private static class VersionAttribute implements StaticAttributeTag {
+    private static final class VersionAttribute implements StaticAttributeTag {
+        private static final String START_AT_METADATA_KEY = "VersionedRecordExtension:StartAt";
+        private static final String INCREMENT_BY_METADATA_KEY = "VersionedRecordExtension:IncrementBy";
+
+        private final Long startAt;
+        private final Long incrementBy;
+
+        private VersionAttribute() {
+            this.startAt = null;
+            this.incrementBy = null;
+        }
+
+        private VersionAttribute(Long startAt, Long incrementBy) {
+            this.startAt = startAt;
+            this.incrementBy = incrementBy;
+        }
+
         @Override
         public Consumer<StaticTableMetadata.Builder> modifyMetadata(String attributeName,
                                                                     AttributeValueType attributeValueType) {
             if (attributeValueType != AttributeValueType.N) {
                 throw new IllegalArgumentException(String.format(
                     "Attribute '%s' of type %s is not a suitable type to be used as a version attribute. Only type 'N' " +
-                        "is supported.", attributeName, attributeValueType.name()));
+                    "is supported.", attributeName, attributeValueType.name()));
+            }
+
+            Validate.isNotNegativeOrNull(startAt, "startAt");
+
+            if (incrementBy != null && incrementBy < 1) {
+                throw new IllegalArgumentException("incrementBy must be greater than 0.");
             }
 
             return metadata -> metadata.addCustomMetadataObject(CUSTOM_METADATA_KEY, attributeName)
+                                       .addCustomMetadataObject(START_AT_METADATA_KEY, startAt)
+                                       .addCustomMetadataObject(INCREMENT_BY_METADATA_KEY, incrementBy)
                                        .markAttributeAsKey(attributeName, attributeValueType);
         }
     }
@@ -106,31 +146,53 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
         String attributeKeyRef = keyRef(versionAttributeKey.get());
         AttributeValue newVersionValue;
         Expression condition;
-        Optional<AttributeValue> existingVersionValue =
-            Optional.ofNullable(itemToTransform.get(versionAttributeKey.get()));
 
-        if (!existingVersionValue.isPresent() || isNullAttributeValue(existingVersionValue.get())) {
-            // First version of the record
-            newVersionValue = AttributeValue.builder().n("1").build();
+        AttributeValue existingVersionValue = itemToTransform.get(versionAttributeKey.get());
+        Long versionStartAtFromAnnotation = context.tableMetadata()
+                                                   .customMetadataObject(VersionAttribute.START_AT_METADATA_KEY, Long.class)
+                                                   .orElse(this.startAt);
+        Long versionIncrementByFromAnnotation = context.tableMetadata()
+                                                   .customMetadataObject(VersionAttribute.INCREMENT_BY_METADATA_KEY, Long.class)
+                                                   .orElse(this.incrementBy);
+
+
+        if (isInitialVersion(existingVersionValue, versionStartAtFromAnnotation)) {
+            newVersionValue = AttributeValue.builder()
+                                            .n(Long.toString(versionStartAtFromAnnotation + versionIncrementByFromAnnotation))
+                                            .build();
             condition = Expression.builder()
                                   .expression(String.format("attribute_not_exists(%s)", attributeKeyRef))
                                   .expressionNames(Collections.singletonMap(attributeKeyRef, versionAttributeKey.get()))
                                   .build();
         } else {
             // Existing record, increment version
-            if (existingVersionValue.get().n() == null) {
+            if (existingVersionValue.n() == null) {
                 // In this case a non-null version attribute is present, but it's not an N
                 throw new IllegalArgumentException("Version attribute appears to be the wrong type. N is required.");
             }
 
-            int existingVersion = Integer.parseInt(existingVersionValue.get().n());
+            long existingVersion = Long.parseLong(existingVersionValue.n());
             String existingVersionValueKey = VERSIONED_RECORD_EXPRESSION_VALUE_KEY_MAPPER.apply(versionAttributeKey.get());
-            newVersionValue = AttributeValue.builder().n(Integer.toString(existingVersion + 1)).build();
+
+            long increment = versionIncrementByFromAnnotation;
+
+            /*
+            Since the new incrementBy and StartAt functionality can now accept any positive number, though unlikely
+            to happen in a real life scenario, we should add overflow protection.
+            */
+            if (existingVersion > Long.MAX_VALUE - increment) {
+                throw new IllegalStateException(
+                    String.format("Version overflow detected. Current version %d + increment %d would exceed Long.MAX_VALUE",
+                                  existingVersion, increment));
+            }
+
+            newVersionValue = AttributeValue.builder().n(Long.toString(existingVersion + increment)).build();
+
             condition = Expression.builder()
                                   .expression(String.format("%s = %s", attributeKeyRef, existingVersionValueKey))
                                   .expressionNames(Collections.singletonMap(attributeKeyRef, versionAttributeKey.get()))
                                   .expressionValues(Collections.singletonMap(existingVersionValueKey,
-                                                                             existingVersionValue.get()))
+                                                                             existingVersionValue))
                                   .build();
         }
 
@@ -142,13 +204,55 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
                                 .build();
     }
 
+    private boolean isInitialVersion(AttributeValue existingVersionValue, Long versionStartAtFromAnnotation) {
+        if (existingVersionValue == null || isNullAttributeValue(existingVersionValue)) {
+            return true;
+        }
+
+        if (existingVersionValue.n() != null) {
+            long currentVersion = Long.parseLong(existingVersionValue.n());
+            // If annotation value is present, use it, otherwise fall back to the extension's value
+            Long effectiveStartAt = versionStartAtFromAnnotation != null ? versionStartAtFromAnnotation : this.startAt;
+            return currentVersion == effectiveStartAt;
+        }
+
+        return false;
+    }
+
     @NotThreadSafe
     public static final class Builder {
+        private Long startAt;
+        private Long incrementBy;
+
         private Builder() {
         }
 
+        /**
+         * Sets the startAt used to compare if a record is the initial version of a record.
+         * Default value - {@code 0}.
+         *
+         * @param startAt the starting value for version comparison, must not be negative
+         * @return the builder instance
+         */
+        public Builder startAt(Long startAt) {
+            this.startAt = startAt;
+            return this;
+        }
+
+        /**
+         * Sets the amount to increment the version by with each subsequent update.
+         * Default value - {@code 1}.
+         *
+         * @param incrementBy the amount to increment the version by, must be greater than 0
+         * @return the builder instance
+         */
+        public Builder incrementBy(Long incrementBy) {
+            this.incrementBy = incrementBy;
+            return this;
+        }
+
         public VersionedRecordExtension build() {
-            return new VersionedRecordExtension();
+            return new VersionedRecordExtension(this.startAt, this.incrementBy);
         }
     }
 }

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/extensions/annotations/DynamoDbVersionAttribute.java

@@ -33,4 +33,20 @@
 @Retention(RetentionPolicy.RUNTIME)
 @BeanTableSchemaAttributeTag(VersionRecordAttributeTags.class)
 public @interface DynamoDbVersionAttribute {
+    /**
+     * The starting value for the version attribute.
+     * Default value - {@code 0}.
+     *
+     * @return the starting value
+     */
+    long startAt() default 0;
+
+    /**
+     * The amount to increment the version by with each update.
+     * Default value - {@code 1}.
+     *
+     * @return the increment value
+     */
+    long incrementBy() default 1;
+
 }

services-custom/dynamodb-enhanced/src/main/java/software/amazon/awssdk/enhanced/dynamodb/internal/extensions/VersionRecordAttributeTags.java

@@ -26,6 +26,6 @@ private VersionRecordAttributeTags() {
     }
 
     public static StaticAttributeTag attributeTagFor(DynamoDbVersionAttribute annotation) {
-        return VersionedRecordExtension.AttributeTags.versionAttribute();
+        return VersionedRecordExtension.AttributeTags.versionAttribute(annotation.startAt(), annotation.incrementBy());
     }
 }