Skip to content

Commit 6398d67

Browse files
RanVakninRanVaknin
committed
Adding annotation support
1 parent 3ce1c9d commit 6398d67

File tree

5 files changed

+311
-41
lines changed

5 files changed

+311
-41
lines changed

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -311,8 +311,17 @@ public int hashCode() {
311311
return result;
312312
}
313313

314+
@Override
315+
public String toString() {
316+
return "Expression{" +
317+
"expression='" + expression + '\'' +
318+
", expressionValues=" + expressionValues +
319+
", expressionNames=" + expressionNames +
320+
'}';
321+
}
322+
314323
/**
315-
* A builder for {@link Expression}
324+
* A builder for {@link Expression}v
316325
*/
317326
@NotThreadSafe
318327
public static final class Builder {

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

Lines changed: 166 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag;
3535
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableMetadata;
3636
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
37+
import software.amazon.awssdk.utils.Pair;
3738

3839
/**
3940
* This extension implements optimistic locking on record writes by means of a 'record version number' that is used
@@ -60,8 +61,20 @@ public final class VersionedRecordExtension implements DynamoDbEnhancedClientExt
6061
private static final Function<String, String> VERSIONED_RECORD_EXPRESSION_VALUE_KEY_MAPPER = key -> ":old_" + key + "_value";
6162
private static final String CUSTOM_METADATA_KEY = "VersionedRecordExtension:VersionAttribute";
6263
private static final VersionAttribute VERSION_ATTRIBUTE = new VersionAttribute();
64+
private static final AttributeValue DEFAULT_VALUE = AttributeValue.fromNul(Boolean.TRUE);
6365

64-
private VersionedRecordExtension() {
66+
private final int startAt;
67+
private final int incrementBy;
68+
69+
/**
70+
* Creates a new {@link VersionedRecordExtension} using the supplied starting and incrementing value.
71+
*
72+
* @param startAt the value used to compare if a record is the initial version of a record.
73+
* @param incrementBy the amount to increment the version by with each subsequent update.
74+
*/
75+
private VersionedRecordExtension(int startAt, int incrementBy) {
76+
this.startAt = startAt;
77+
this.incrementBy = incrementBy;
6578
}
6679

6780
public static Builder builder() {
@@ -75,19 +88,42 @@ private AttributeTags() {
7588
public static StaticAttributeTag versionAttribute() {
7689
return VERSION_ATTRIBUTE;
7790
}
91+
92+
public static StaticAttributeTag versionAttribute(int startAt, int incrementBy) {
93+
return new VersionAttribute(startAt, incrementBy);
94+
}
7895
}
7996

80-
private static class VersionAttribute implements StaticAttributeTag {
97+
private static final class VersionAttribute implements StaticAttributeTag {
98+
private static final String START_AT_METADATA_KEY = "VersionedRecordExtension:StartAt";
99+
private static final String INCREMENT_BY_METADATA_KEY = "VersionedRecordExtension:IncrementBy";
100+
101+
private final int startAt;
102+
private final int incrementBy;
103+
104+
private VersionAttribute() {
105+
this.startAt = 0;
106+
this.incrementBy = 1;
107+
}
108+
109+
private VersionAttribute(int startAt, int incrementBy) {
110+
this.startAt = startAt;
111+
this.incrementBy = incrementBy;
112+
}
113+
81114
@Override
82115
public Consumer<StaticTableMetadata.Builder> modifyMetadata(String attributeName,
83116
AttributeValueType attributeValueType) {
84117
if (attributeValueType != AttributeValueType.N) {
85118
throw new IllegalArgumentException(String.format(
86119
"Attribute '%s' of type %s is not a suitable type to be used as a version attribute. Only type 'N' " +
87-
"is supported.", attributeName, attributeValueType.name()));
120+
"is supported.", attributeName, attributeValueType.name()));
88121
}
89122

90-
return metadata -> metadata.addCustomMetadataObject(CUSTOM_METADATA_KEY, attributeName)
123+
return metadata -> metadata
124+
.addCustomMetadataObject(CUSTOM_METADATA_KEY, attributeName)
125+
.addCustomMetadataObject(START_AT_METADATA_KEY, startAt)
126+
.addCustomMetadataObject(INCREMENT_BY_METADATA_KEY, incrementBy)
91127
.markAttributeAsKey(attributeName, attributeValueType);
92128
}
93129
}
@@ -101,39 +137,14 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
101137
return WriteModification.builder().build();
102138
}
103139

104-
Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
105140

106-
String attributeKeyRef = keyRef(versionAttributeKey.get());
107-
AttributeValue newVersionValue;
108-
Expression condition;
109-
Optional<AttributeValue> existingVersionValue =
110-
Optional.ofNullable(itemToTransform.get(versionAttributeKey.get()));
111-
112-
if (!existingVersionValue.isPresent() || isNullAttributeValue(existingVersionValue.get())) {
113-
// First version of the record
114-
newVersionValue = AttributeValue.builder().n("1").build();
115-
condition = Expression.builder()
116-
.expression(String.format("attribute_not_exists(%s)", attributeKeyRef))
117-
.expressionNames(Collections.singletonMap(attributeKeyRef, versionAttributeKey.get()))
118-
.build();
119-
} else {
120-
// Existing record, increment version
121-
if (existingVersionValue.get().n() == null) {
122-
// In this case a non-null version attribute is present, but it's not an N
123-
throw new IllegalArgumentException("Version attribute appears to be the wrong type. N is required.");
124-
}
141+
Pair<AttributeValue, Expression> updates = getRecordUpdates(versionAttributeKey.get(), context);
125142

126-
int existingVersion = Integer.parseInt(existingVersionValue.get().n());
127-
String existingVersionValueKey = VERSIONED_RECORD_EXPRESSION_VALUE_KEY_MAPPER.apply(versionAttributeKey.get());
128-
newVersionValue = AttributeValue.builder().n(Integer.toString(existingVersion + 1)).build();
129-
condition = Expression.builder()
130-
.expression(String.format("%s = %s", attributeKeyRef, existingVersionValueKey))
131-
.expressionNames(Collections.singletonMap(attributeKeyRef, versionAttributeKey.get()))
132-
.expressionValues(Collections.singletonMap(existingVersionValueKey,
133-
existingVersionValue.get()))
134-
.build();
135-
}
143+
// Unpack values from Pair
144+
AttributeValue newVersionValue = updates.left();
145+
Expression condition = updates.right();
136146

147+
Map<String, AttributeValue> itemToTransform = new HashMap<>(context.items());
137148
itemToTransform.put(versionAttributeKey.get(), newVersionValue);
138149

139150
return WriteModification.builder()
@@ -142,13 +153,133 @@ public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite contex
142153
.build();
143154
}
144155

156+
private Pair<AttributeValue, Expression> getRecordUpdates(String versionAttributeKey,
157+
DynamoDbExtensionContext.BeforeWrite context) {
158+
Map<String, AttributeValue> itemToTransform = context.items();
159+
160+
// Default to NUL if not present to reduce additional checks further along
161+
AttributeValue existingVersionValue = itemToTransform.getOrDefault(versionAttributeKey, DEFAULT_VALUE);
162+
163+
if (isInitialVersion(existingVersionValue, context)) {
164+
// First version of the record ensure it does not exist
165+
return createInitialRecord(versionAttributeKey, context);
166+
}
167+
// Existing record, increment version
168+
return updateExistingRecord(versionAttributeKey, existingVersionValue, context);
169+
}
170+
171+
private boolean isInitialVersion(AttributeValue existingVersionValue, DynamoDbExtensionContext.BeforeWrite context) {
172+
if (isNullAttributeValue(existingVersionValue)) {
173+
return true;
174+
}
175+
176+
177+
Optional<Integer> versionStartAtFromAnnotation = context.tableMetadata()
178+
.customMetadataObject(VersionAttribute.START_AT_METADATA_KEY,
179+
Integer.class);
180+
181+
return isNullAttributeValue(existingVersionValue)
182+
|| (versionStartAtFromAnnotation.isPresent() &&
183+
getExistingVersion(existingVersionValue) == versionStartAtFromAnnotation.get())
184+
|| (!versionStartAtFromAnnotation.isPresent() &&
185+
getExistingVersion(existingVersionValue) == this.startAt);
186+
}
187+
188+
private Pair<AttributeValue, Expression> createInitialRecord(String versionAttributeKey,
189+
DynamoDbExtensionContext.BeforeWrite context) {
190+
Optional<Integer> versionStartAtFromAnnotation = context.tableMetadata()
191+
.customMetadataObject(VersionAttribute.START_AT_METADATA_KEY,
192+
Integer.class);
193+
194+
AttributeValue newVersionValue = versionStartAtFromAnnotation.isPresent() ?
195+
incrementVersion(versionStartAtFromAnnotation.get(), context) :
196+
incrementVersion(this.startAt, context);
197+
198+
199+
String attributeKeyRef = keyRef(versionAttributeKey);
200+
201+
Expression condition = Expression.builder()
202+
// Check that the version does not exist before setting the initial value.
203+
.expression(String.format("attribute_not_exists(%s)", attributeKeyRef))
204+
.expressionNames(Collections.singletonMap(attributeKeyRef, versionAttributeKey))
205+
.build();
206+
207+
return Pair.of(newVersionValue, condition);
208+
}
209+
210+
private Pair<AttributeValue, Expression> updateExistingRecord(String versionAttributeKey,
211+
AttributeValue existingVersionValue,
212+
DynamoDbExtensionContext.BeforeWrite context) {
213+
int existingVersion = getExistingVersion(existingVersionValue);
214+
AttributeValue newVersionValue = incrementVersion(existingVersion, context);
215+
216+
String attributeKeyRef = keyRef(versionAttributeKey);
217+
String existingVersionValueKey = VERSIONED_RECORD_EXPRESSION_VALUE_KEY_MAPPER.apply(versionAttributeKey);
218+
219+
Expression condition = Expression.builder()
220+
// Check that the version matches the existing value before setting the updated value.
221+
.expression(String.format("%s = %s", attributeKeyRef, existingVersionValueKey))
222+
.expressionNames(Collections.singletonMap(attributeKeyRef, versionAttributeKey))
223+
.expressionValues(Collections.singletonMap(existingVersionValueKey,
224+
existingVersionValue))
225+
.build();
226+
227+
return Pair.of(newVersionValue, condition);
228+
}
229+
230+
private int getExistingVersion(AttributeValue existingVersionValue) {
231+
if (existingVersionValue.n() == null) {
232+
// In this case a non-null version attribute is present, but it's not an N
233+
throw new IllegalArgumentException("Version attribute appears to be the wrong type. N is required.");
234+
}
235+
236+
return Integer.parseInt(existingVersionValue.n());
237+
}
238+
239+
private AttributeValue incrementVersion(int version, DynamoDbExtensionContext.BeforeWrite context) {
240+
Optional<Integer> versionIncrementByFromAnnotation = context.tableMetadata()
241+
.customMetadataObject(VersionAttribute.INCREMENT_BY_METADATA_KEY,
242+
Integer.class);
243+
if (versionIncrementByFromAnnotation.isPresent()) {
244+
return AttributeValue.fromN(Integer.toString(version + versionIncrementByFromAnnotation.get()));
245+
}
246+
return AttributeValue.fromN(Integer.toString(version + this.incrementBy));
247+
}
248+
145249
@NotThreadSafe
146250
public static final class Builder {
251+
private int startAt = 0;
252+
private int incrementBy = 1;
253+
147254
private Builder() {
148255
}
149256

257+
/**
258+
* Sets the startAt used to compare if a record is the initial version of a record.
259+
* Default value - {@code 0}.
260+
*
261+
* @param startAt
262+
* @return the builder instance
263+
*/
264+
public Builder startAt(int startAt) {
265+
this.startAt = startAt;
266+
return this;
267+
}
268+
269+
/**
270+
* Sets the amount to increment the version by with each subsequent update.
271+
* Default value - {@code 1}.
272+
*
273+
* @param incrementBy
274+
* @return the builder instance
275+
*/
276+
public Builder incrementBy(int incrementBy) {
277+
this.incrementBy = incrementBy;
278+
return this;
279+
}
280+
150281
public VersionedRecordExtension build() {
151-
return new VersionedRecordExtension();
282+
return new VersionedRecordExtension(this.startAt, this.incrementBy);
152283
}
153284
}
154285
}

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,4 +33,20 @@
3333
@Retention(RetentionPolicy.RUNTIME)
3434
@BeanTableSchemaAttributeTag(VersionRecordAttributeTags.class)
3535
public @interface DynamoDbVersionAttribute {
36+
/**
37+
* The starting value for the version attribute.
38+
* Default value - {@code 0}.
39+
*
40+
* @return the starting value
41+
*/
42+
int startAt() default 0;
43+
44+
/**
45+
* The amount to increment the version by with each update.
46+
* Default value - {@code 1}.
47+
*
48+
* @return the increment value
49+
*/
50+
int incrementBy() default 1;
51+
3652
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,6 @@ private VersionRecordAttributeTags() {
2626
}
2727

2828
public static StaticAttributeTag attributeTagFor(DynamoDbVersionAttribute annotation) {
29-
return VersionedRecordExtension.AttributeTags.versionAttribute();
29+
return VersionedRecordExtension.AttributeTags.versionAttribute(annotation.startAt(), annotation.incrementBy());
3030
}
3131
}

0 commit comments

Comments
 (0)