Skip to content

Update additional attributes on lock release #108

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 4 commits into
base: master
Choose a base branch
from
Draft

Update additional attributes on lock release #108

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
cac6ea6
update additional attributes on lock release
nathanallen-toast Jan 24, 2025
ccdd0e1
prevent attribute collision
nathanallen-toast Jan 24, 2025
56b8f95
put attribute updates in expression
nathanallen-toast Jan 24, 2025
6287451
remove debug log
nathanallen-toast Jan 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 43 additions & 5 deletions src/main/java/com/amazonaws/services/dynamodbv2/AmazonDynamoDBLockClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Objects;
Expand Down Expand Up @@ -52,6 +53,7 @@
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeDefinition;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.AttributeValueUpdate;
import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
import software.amazon.awssdk.services.dynamodb.model.CreateTableRequest;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
Expand Down Expand Up @@ -841,7 +843,7 @@ public boolean releaseLock(final LockItem lockItem) {
return this.releaseLock(ReleaseLockOptions.builder(lockItem).withDeleteLock(lockItem.getDeleteLockItemOnClose()).build());
}

public boolean releaseLock(final ReleaseLockOptions options) {
public boolean releaseLock(final ReleaseLockOptions options) throws IllegalArgumentException {
Objects.requireNonNull(options, "ReleaseLockOptions cannot be null");

final LockItem lockItem = options.getLockItem();
Expand Down Expand Up @@ -893,20 +895,38 @@ public boolean releaseLock(final ReleaseLockOptions options) {

this.dynamoDB.deleteItem(deleteItemRequest);
} else {
final String updateExpression;
final Map<String, AttributeValueUpdate> additionalAttributeUpdates =
checkAndRetrieveAdditionalAttributeUpdates(options);

StringBuilder updateExpression;
expressionAttributeNames.put(IS_RELEASED_PATH_EXPRESSION_VARIABLE, IS_RELEASED);
expressionAttributeValues.put(IS_RELEASED_VALUE_EXPRESSION_VARIABLE, IS_RELEASED_ATTRIBUTE_VALUE);
if (data.isPresent()) {
updateExpression = UPDATE_IS_RELEASED_AND_DATA;
updateExpression = new StringBuilder(UPDATE_IS_RELEASED_AND_DATA);
expressionAttributeNames.put(DATA_PATH_EXPRESSION_VARIABLE, DATA);
expressionAttributeValues.put(DATA_VALUE_EXPRESSION_VARIABLE, AttributeValue.builder().b(SdkBytes.fromByteBuffer(data.get())).build());
} else {
updateExpression = UPDATE_IS_RELEASED;
updateExpression = new StringBuilder(UPDATE_IS_RELEASED);
}
for (Entry<String, AttributeValueUpdate> entry : additionalAttributeUpdates.entrySet()) {
String k = entry.getKey();
String attributePathExpressionVariable = "#" + k.toLowerCase(Locale.getDefault());
String attributeValueExpressionVariable = ":" + k.toLowerCase(Locale.getDefault());
expressionAttributeNames.put(attributePathExpressionVariable, k);
AttributeValueUpdate v = entry.getValue();
expressionAttributeValues.put(attributeValueExpressionVariable, v.value());
updateExpression.append(
String.format(
", %s = %s",
attributePathExpressionVariable,
attributeValueExpressionVariable
)
);
}
final UpdateItemRequest updateItemRequest = UpdateItemRequest.builder()
.tableName(this.tableName)
.key(key)
.updateExpression(updateExpression)
.updateExpression(updateExpression.toString())
.conditionExpression(conditionalExpression)
.expressionAttributeNames(expressionAttributeNames)
.expressionAttributeValues(expressionAttributeValues).build();
Expand Down Expand Up @@ -934,6 +954,24 @@ public boolean releaseLock(final ReleaseLockOptions options) {
return true;
}

private Map<String, AttributeValueUpdate> checkAndRetrieveAdditionalAttributeUpdates(ReleaseLockOptions options) {
final Map<String, AttributeValueUpdate> additionalAttributeUpdates = options.getAdditionalAttributeUpdates();
if (
additionalAttributeUpdates.containsKey(this.partitionKeyName) ||
additionalAttributeUpdates.containsKey(OWNER_NAME) ||
additionalAttributeUpdates.containsKey(LEASE_DURATION) ||
additionalAttributeUpdates.containsKey(RECORD_VERSION_NUMBER) ||
additionalAttributeUpdates.containsKey(DATA) ||
this.sortKeyName.isPresent() &&
additionalAttributeUpdates.containsKey(this.sortKeyName.get())
) {
throw new IllegalArgumentException(String
.format("Additional attribute cannot be one of the following types: " + "%s, %s, %s, %s, %s", this.partitionKeyName, OWNER_NAME, LEASE_DURATION,
RECORD_VERSION_NUMBER, DATA));
}
return additionalAttributeUpdates;
}

private Map<String, AttributeValue> getItemKeys(LockItem lockItem) {
return getKeys(lockItem.getPartitionKey(), lockItem.getSortKey());
}
Expand Down
37 changes: 30 additions & 7 deletions src/main/java/com/amazonaws/services/dynamodbv2/ReleaseLockOptions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,11 @@
*/
package com.amazonaws.services.dynamodbv2;

import software.amazon.awssdk.services.dynamodb.model.AttributeValueUpdate;

import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

/**
Expand All @@ -28,25 +32,29 @@ public class ReleaseLockOptions {
private final boolean deleteLock;
private final boolean bestEffort;
private final Optional<ByteBuffer> data;
private final Map<String, AttributeValueUpdate> additionalAttributeUpdates;

ReleaseLockOptions(final LockItem lockItem, final boolean deleteLock, final boolean bestEffort, final Optional<ByteBuffer> data) {
ReleaseLockOptions(final LockItem lockItem, final boolean deleteLock, final boolean bestEffort, final Optional<ByteBuffer> data, Map<String, AttributeValueUpdate> additionalAttributeUpdates) {
this.lockItem = lockItem;
this.deleteLock = deleteLock;
this.bestEffort = bestEffort;
this.data = data;
this.additionalAttributeUpdates = additionalAttributeUpdates;
}

public static class ReleaseLockOptionsBuilder {
private LockItem lockItem;
private final LockItem lockItem;
private boolean deleteLock;
private boolean bestEffort;
private Optional<ByteBuffer> data;
private Map<String, AttributeValueUpdate> additionalAttributeUpdates;

ReleaseLockOptionsBuilder(final LockItem lockItem) {
this.lockItem = lockItem;
this.deleteLock = true;
this.bestEffort = false;
this.data = Optional.empty();
this.additionalAttributeUpdates = new HashMap<>();
}

/**
Expand Down Expand Up @@ -90,14 +98,25 @@ public ReleaseLockOptionsBuilder withData(final ByteBuffer data) {
return this;
}

/**
* Stores some additional attributes with the lock. This can be used to add/update any arbitrary parameters to
* the lock row.
*
* @param additionalAttributeUpdates an arbitrary map of attribute updates to store with the lock row to be acquired
* @return a reference to this builder for fluent method chaining
*/
public ReleaseLockOptionsBuilder withAdditionalAttributeUpdates(Map<String, AttributeValueUpdate> additionalAttributeUpdates) {
this.additionalAttributeUpdates = additionalAttributeUpdates;
return this;
}

public ReleaseLockOptions build() {
return new ReleaseLockOptions(this.lockItem, this.deleteLock, this.bestEffort, this.data);
return new ReleaseLockOptions(this.lockItem, this.deleteLock, this.bestEffort, this.data, this.additionalAttributeUpdates);
}

@Override
public java.lang.String toString() {
return "ReleaseLockOptions.ReleaseLockOptionsBuilder(lockItem=" + this.lockItem + ", deleteLock=" + this.deleteLock + ", bestEffort=" + this.bestEffort + ", data="
+ this.data + ")";
public String toString() {
return String.format("ReleaseLockOptions.ReleaseLockOptionsBuilder(lockItem=%s, deleteLock=%s, bestEffort=%s, data=%s, additionalAttributeUpdates=%s)", lockItem, deleteLock, bestEffort, data, additionalAttributeUpdates);
}
}

Expand Down Expand Up @@ -129,4 +148,8 @@ boolean isBestEffort() {
Optional<ByteBuffer> getData() {
return this.data;
}
}

Map<String, AttributeValueUpdate> getAdditionalAttributeUpdates() {
return additionalAttributeUpdates;
}
}
98 changes: 97 additions & 1 deletion src/test/java/com/amazonaws/services/dynamodbv2/BasicLockClientTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import software.amazon.awssdk.core.exception.SdkClientException;
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
import software.amazon.awssdk.services.dynamodb.model.AttributeValueUpdate;
import software.amazon.awssdk.services.dynamodb.model.DeleteItemRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
Expand All @@ -45,6 +46,7 @@
import static org.junit.Assert.assertNotEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNotSame;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.Mockito.doReturn;
Expand Down Expand Up @@ -1002,7 +1004,7 @@ public void testAcquireLockMultipleTimesReentrant() throws InterruptedException
}

@Test
public void testSendHeatbeatWithRangeKey() throws IOException, LockNotGrantedException, InterruptedException {
public void testSendHeartbeatWithRangeKey() throws IOException, LockNotGrantedException, InterruptedException {

final String data = new String("testSendHeartbeatLeaveData" + SECURE_RANDOM.nextDouble());

Expand Down Expand Up @@ -1124,6 +1126,62 @@ public void testReleaseLockLeaveItemAndChangeData() throws LockNotGrantedExcepti
item.close();
}

@Test
public void testReleaseLockLeaveItemAndChangeAttributes() throws LockNotGrantedException, InterruptedException {
final String lockTypeAttributeName = "LockType";
final String boreDiameterAttributeName = "BoreDiameter";

final AttributeValue initialLockType = AttributeValue.fromS("Mortise");

final Map<String, AttributeValue> additionalAttributes = new HashMap<>();
additionalAttributes.put(lockTypeAttributeName, initialLockType);

LockItem item = this.lockClient.acquireLock(
AcquireLockOptions.builder("testKey1")
.withAdditionalAttributes(additionalAttributes)
.withDeleteLockOnRelease(false)
.withReplaceData(true)
.build()
);
assertEquals(initialLockType, item.getAdditionalAttributes().get(lockTypeAttributeName));
assertNull(item.getAdditionalAttributes().get(boreDiameterAttributeName));

final AttributeValue updatedLockType = AttributeValue.fromS("Deadbolt");

final AttributeValue boreDiameter = AttributeValue.fromN("55");

final Map<String, AttributeValueUpdate> additionalAttributeUpdates = new HashMap<>();
additionalAttributeUpdates.put(
lockTypeAttributeName,
AttributeValueUpdate.builder().value(updatedLockType).build()
);
additionalAttributeUpdates.put(
boreDiameterAttributeName,
AttributeValueUpdate.builder().value(boreDiameter).build()
);

this.lockClient.releaseLock(
ReleaseLockOptions.builder(item)
.withDeleteLock(false)
.withAdditionalAttributeUpdates(additionalAttributeUpdates)
.build()
);

assertEquals(Optional.empty(), this.lockClient.getLock("testKey1", Optional.empty()));

item = this.lockClient.getLockFromDynamoDB(GET_LOCK_OPTIONS_DO_NOT_DELETE_ON_RELEASE).get();
assertTrue(item.isReleased());
assertEquals(updatedLockType, item.getAdditionalAttributes().get(lockTypeAttributeName));
assertEquals(boreDiameter, item.getAdditionalAttributes().get(boreDiameterAttributeName));

item = this.lockClient.acquireLock(AcquireLockOptions.builder("testKey1")
.withDeleteLockOnRelease(false).withReplaceData(false).build());
assertEquals(updatedLockType, item.getAdditionalAttributes().get(lockTypeAttributeName));
assertEquals(boreDiameter, item.getAdditionalAttributes().get(boreDiameterAttributeName));

item.close();
}

@Test
public void testReleaseLockRemoveItem() throws LockNotGrantedException, InterruptedException {
final String data = "testReleaseLockRemoveItem";
Expand Down Expand Up @@ -1482,6 +1540,44 @@ private void testInvalidAttribute(final String invalidAttribute) throws LockNotG
.build());
}

@Test(expected = IllegalArgumentException.class)
public void testInvalidAttributeUpdatesData() throws LockNotGrantedException, InterruptedException {
this.testInvalidAttributeUpdate("data");
}

@Test(expected = IllegalArgumentException.class)
public void testInvalidAttributeUpdatesKey() throws LockNotGrantedException, InterruptedException {
this.testInvalidAttributeUpdate("key");
}

@Test(expected = IllegalArgumentException.class)
public void testInvalidAttributeUpdatesLeaseDuration() throws LockNotGrantedException, InterruptedException {
this.testInvalidAttributeUpdate("leaseDuration");
}

@Test(expected = IllegalArgumentException.class)
public void testInvalidAttributeUpdatesRecordVersionNumber() throws LockNotGrantedException, InterruptedException {
this.testInvalidAttributeUpdate("recordVersionNumber");
}

@Test(expected = IllegalArgumentException.class)
public void testInvalidAttributeUpdatesOwnerName() throws LockNotGrantedException, InterruptedException {
this.testInvalidAttributeUpdate("ownerName");
}

private void testInvalidAttributeUpdate(final String invalidAttribute) throws InterruptedException {
final LockItem item = this.lockClient.acquireLock(AcquireLockOptions.builder("testKey1").build());

final Map<String, AttributeValueUpdate> additionalAttributeUpdates = new HashMap<>();
additionalAttributeUpdates.put(
invalidAttribute,
AttributeValueUpdate.builder().value(AttributeValue.fromS("ok")).build()
);
this.lockClient.releaseLock(
ReleaseLockOptions.builder(item).withAdditionalAttributeUpdates(additionalAttributeUpdates).build()
);
}

@Test
public void testLockItemToString() throws LockNotGrantedException, InterruptedException {
final LockItem lockItem = this.lockClient.acquireLock(ACQUIRE_LOCK_OPTIONS_TEST_KEY_1);
Expand Down