software.amazon.awssdk.enhanced.dynamodb.extensions.AtomicCounterExtension Maven / Gradle / Ivy
/*
* 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 static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.keyRef;
import static software.amazon.awssdk.enhanced.dynamodb.internal.EnhancedClientUtils.valueRef;
import static software.amazon.awssdk.enhanced.dynamodb.internal.update.UpdateExpressionUtils.ifNotExists;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.stream.Collectors;
import software.amazon.awssdk.annotations.SdkPublicApi;
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient;
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClientExtension;
import software.amazon.awssdk.enhanced.dynamodb.DynamoDbExtensionContext;
import software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbAtomicCounter;
import software.amazon.awssdk.enhanced.dynamodb.internal.extensions.AtomicCounterTag;
import software.amazon.awssdk.enhanced.dynamodb.internal.mapper.AtomicCounter;
import software.amazon.awssdk.enhanced.dynamodb.mapper.BeanTableSchema;
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags;
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableSchema;
import software.amazon.awssdk.enhanced.dynamodb.update.SetAction;
import software.amazon.awssdk.enhanced.dynamodb.update.UpdateExpression;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.utils.CollectionUtils;
/**
* This extension enables atomic counter attributes to be written to the database.
* The extension is loaded by default when you instantiate a
* {@link DynamoDbEnhancedClient} and only needs to be added to the client if you
* are adding custom extensions to the client.
*
* To utilize atomic counters, first create a field in your model that will be used to store the counter.
* This class field should of type {@link Long} and you need to tag it as an atomic counter:
*
* - If you are using the
* {@link BeanTableSchema}, you should annotate with
* {@link DynamoDbAtomicCounter}
* - If you are using the {@link StaticTableSchema},
* use the {@link StaticAttributeTags#atomicCounter()} static attribute tag.
*
*
* Every time a new update of the record is successfully written to the database, the counter will be updated automatically.
* By default, the counter starts at 0 and increments by 1 for each update. The tags provide the capability of adjusting
* the counter start and increment/decrement values such as described in
* {@link DynamoDbAtomicCounter}.
*
* Example 1: Using a bean based table schema
*
* {@code
* @DynamoDbBean
* public class CounterRecord {
* @DynamoDbAtomicCounter(delta = 5, startValue = 10)
* public Long getCustomCounter() {
* return customCounter;
* }
* }
* }
*
*
* Example 2: Using a static table schema
*
* {@code
* private static final StaticTableSchema ITEM_MAPPER =
* StaticTableSchema.builder(AtomicCounterItem.class)
* .newItemSupplier(AtomicCounterItem::new)
* .addAttribute(Long.class, a -> a.name("defaultCounter")
* .getter(AtomicCounterItem::getDefaultCounter)
* .setter(AtomicCounterItem::setDefaultCounter)
* .addTag(StaticAttributeTags.atomicCounter()))
* .build();
* }
*
*
* NOTE: When using putItem, the counter will be reset to its start value.
*/
@SdkPublicApi
public final class AtomicCounterExtension implements DynamoDbEnhancedClientExtension {
private AtomicCounterExtension() {
}
public static AtomicCounterExtension.Builder builder() {
return new AtomicCounterExtension.Builder();
}
/**
* @param context The {@link DynamoDbExtensionContext.BeforeWrite} context containing the state of the execution.
* @return WriteModification contains an update expression representing the counters.
*/
@Override
public WriteModification beforeWrite(DynamoDbExtensionContext.BeforeWrite context) {
Map counters = AtomicCounterTag.resolve(context.tableMetadata());
WriteModification.Builder modificationBuilder = WriteModification.builder();
if (CollectionUtils.isNullOrEmpty(counters)) {
return modificationBuilder.build();
}
switch (context.operationName()) {
case PUT_ITEM:
modificationBuilder.transformedItem(addToItem(counters, context.items()));
break;
case UPDATE_ITEM:
modificationBuilder.updateExpression(createUpdateExpression(counters));
break;
default: break;
}
return modificationBuilder.build();
}
private UpdateExpression createUpdateExpression(Map counters) {
return UpdateExpression.builder()
.actions(counters.entrySet().stream().map(this::counterAction).collect(Collectors.toList()))
.build();
}
private Map addToItem(Map counters, Map items) {
Map itemToTransform = new HashMap<>(items);
counters.forEach((attribute, counter) -> itemToTransform.put(attribute, attributeValue(counter.startValue().value())));
return Collections.unmodifiableMap(itemToTransform);
}
private SetAction counterAction(Map.Entry e) {
String attributeName = e.getKey();
AtomicCounter counter = e.getValue();
String startValueName = attributeName + counter.startValue().name();
String deltaValueName = attributeName + counter.delta().name();
String valueExpression = ifNotExists(attributeName, startValueName) + " + " + valueRef(deltaValueName);
AttributeValue startValue = attributeValue(counter.startValue().value() - counter.delta().value());
AttributeValue deltaValue = attributeValue(counter.delta().value());
return SetAction.builder()
.path(keyRef(attributeName))
.value(valueExpression)
.putExpressionName(keyRef(attributeName), attributeName)
.putExpressionValue(valueRef(startValueName), startValue)
.putExpressionValue(valueRef(deltaValueName), deltaValue)
.build();
}
private AttributeValue attributeValue(long value) {
return AtomicCounter.CounterAttribute.resolvedValue(value);
}
public static final class Builder {
private Builder() {
}
public AtomicCounterExtension build() {
return new AtomicCounterExtension();
}
}
}