software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest Java Examples

@Test
public void generateRequest_withExtension_modifiesKeyPortionOfItem() {
    FakeItem baseFakeItem = createUniqueFakeItem();
    FakeItem fakeItem = createUniqueFakeItem();

    Map<String, AttributeValue> baseMap = FakeItem.getTableSchema().itemToMap(baseFakeItem, false);
    Map<String, AttributeValue> fakeMap = FakeItem.getTableSchema().itemToMap(fakeItem, false);
    Map<String, AttributeValue> keyMap = FakeItem.getTableSchema().itemToMap(fakeItem, singletonList("id"));

    when(mockDynamoDbEnhancedClientExtension.beforeWrite(any(DynamoDbExtensionContext.BeforeWrite.class)))
        .thenReturn(WriteModification.builder().transformedItem(fakeMap).build());

    UpdateItemOperation<FakeItem> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class).item(baseFakeItem).build());

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    mockDynamoDbEnhancedClientExtension);

    assertThat(request.key(), is(keyMap));
    verify(mockDynamoDbEnhancedClientExtension).beforeWrite(DefaultDynamoDbExtensionContext.builder()
                                                                                           .tableMetadata(FakeItem.getTableMetadata())
                                                                                           .operationContext(PRIMARY_CONTEXT)
                                                                                           .items(baseMap).build());
}
 

@Override
public void doUnlock() {
    // Set lockUntil to now or lockAtLeastUntil whichever is later
    String unlockTimeIso = toIsoString(lockConfiguration.getUnlockTime());

    Map<String, AttributeValue> key = singletonMap(ID, attr(lockConfiguration.getName()));

    Map<String, AttributeValue> attributeUpdates = singletonMap(":lockUntil", attr(unlockTimeIso));

    UpdateItemRequest request = UpdateItemRequest.builder()
            .tableName(tableName)
            .key(key)
            .updateExpression(RELEASE_LOCK_QUERY)
            .expressionAttributeValues(attributeUpdates)
            .returnValues(ReturnValue.UPDATED_NEW)
            .build();

    dynamoDbClient.updateItem(request);
}
 

@Test
public void generateRequest_withMinimalConditionExpression() {
    FakeItemWithSort item = createUniqueFakeItemWithSort();
    item.setOtherAttribute1("value-1");

    UpdateItemOperation<FakeItemWithSort> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class)
                                                            .item(item)
                                                            .conditionExpression(MINIMAL_CONDITION_EXPRESSION)
                                                            .build());

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    null);

    Map<String, AttributeValue> expectedValues = new HashMap<>();
    expectedValues.put(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build());
    Map<String, String> expectedNames = new HashMap<>();
    expectedNames.put(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1");
    expectedNames.put(OTHER_ATTRIBUTE_2_NAME, "other_attribute_2");
    assertThat(request.conditionExpression(), is(MINIMAL_CONDITION_EXPRESSION.expression()));
    assertThat(request.expressionAttributeNames(), is(expectedNames));
    assertThat(request.expressionAttributeValues(), is(expectedValues));
}
 

@Test
public void testUpdateStatus() throws Exception {
  DyneinJobSpec jobSpec = getTestJobSpec(validToken, "test1");
  Schedule schedule = jobSpecToSchedule(jobSpec);
  UpdateItemRequest updateItemRequest =
      getUpdateItemReq(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);

  when(ddbClient.updateItem(updateItemRequest))
      .thenReturn(
          CompletableFuture.completedFuture(
              UpdateItemResponse.builder()
                  .attributes(
                      ImmutableMap.of(
                          Attribute.JOB_STATUS.columnName,
                          AttributeValue.builder().s(JobStatus.ACQUIRED.name()).build()))
                  .build()));

  CompletableFuture<Schedule> response =
      scheduleManager.updateStatus(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);
  Assert.assertEquals(response.get(1, TimeUnit.SECONDS), schedule.withStatus(JobStatus.ACQUIRED));

  verify(ddbClient, times(1)).updateItem(updateItemRequest);
  verifyNoMoreInteractions(ddbClient);
}
 

@Test
public void testUpdateStatus_emptyResponse() throws Exception {
  DyneinJobSpec jobSpec = getTestJobSpec(validToken, "test1");
  Schedule schedule = jobSpecToSchedule(jobSpec);
  UpdateItemRequest updateItemRequest =
      getUpdateItemReq(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);

  when(ddbClient.updateItem(updateItemRequest))
      .thenReturn(CompletableFuture.completedFuture(UpdateItemResponse.builder().build()));

  CompletableFuture<Schedule> response =
      scheduleManager.updateStatus(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);

  Throwable exception = getException(response);
  assertTrue(exception instanceof IllegalStateException);
  assertEquals(exception.getMessage(), "Status update successful but status isn't returned.");

  verify(ddbClient, times(1)).updateItem(updateItemRequest);
  verifyNoMoreInteractions(ddbClient);
}
 

@Test
public void generateRequest_keyOnlyItem() {
    FakeItemWithSort item = createUniqueFakeItemWithSort();
    UpdateItemOperation<FakeItemWithSort> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class)
                                                            .item(item)
                                                            .ignoreNulls(true)
                                                            .build());
    Map<String, AttributeValue> expectedKey = new HashMap<>();
    expectedKey.put("id", AttributeValue.builder().s(item.getId()).build());
    expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build());
    UpdateItemRequest expectedRequest = UpdateItemRequest.builder()
                                                         .tableName(TABLE_NAME)
                                                         .key(expectedKey)
                                                         .returnValues(ReturnValue.ALL_NEW)
                                                         .build();


    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    null);

    assertThat(request, is(expectedRequest));
}
 

@Test
public void testUpdateStatus_failure() throws Exception {
  DyneinJobSpec jobSpec = getTestJobSpec(validToken, "test1");
  Schedule schedule = jobSpecToSchedule(jobSpec);
  UpdateItemRequest updateItemRequest =
      getUpdateItemReq(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);

  Exception exception = new Exception();
  CompletableFuture<UpdateItemResponse> response = new CompletableFuture<>();
  response.completeExceptionally(exception);
  when(ddbClient.updateItem(updateItemRequest)).thenReturn(response);

  CompletableFuture<Schedule> ret =
      scheduleManager.updateStatus(schedule, JobStatus.SCHEDULED, JobStatus.ACQUIRED);

  assertSame(getException(ret), exception);

  verify(ddbClient, times(1)).updateItem(updateItemRequest);
  verifyNoMoreInteractions(ddbClient);
}
 

private UpdateItemRequest getUpdateItemReq(
    Schedule schedule, Schedule.JobStatus oldStatus, Schedule.JobStatus newStatus) {
  Map<String, AttributeValue> primaryKey = DynamoDBUtils.getPrimaryKey(schedule);
  Map<String, String> attributeNames = new HashMap<>();
  Map<String, AttributeValue> attributeValues = new HashMap<>();
  attributeNames.put("#jobStatus", DynamoDBUtils.Attribute.JOB_STATUS.columnName);
  attributeValues.put(":oldStatus", AttributeValue.builder().s(oldStatus.name()).build());
  attributeValues.put(":newStatus", AttributeValue.builder().s(newStatus.name()).build());

  String updated = "SET #jobStatus = :newStatus";

  return UpdateItemRequest.builder()
      .tableName(this.tableName)
      .key(primaryKey)
      .conditionExpression("#jobStatus = :oldStatus")
      .expressionAttributeNames(attributeNames)
      .expressionAttributeValues(attributeValues)
      .updateExpression(updated)
      .returnValues(ReturnValue.UPDATED_NEW)
      .build();
}
 

@Test
public void generateRequest_withExtension_noModifications() {
    FakeItem baseFakeItem = createUniqueFakeItem();
    when(mockDynamoDbEnhancedClientExtension.beforeWrite(any(DynamoDbExtensionContext.BeforeWrite.class)))
        .thenReturn(WriteModification.builder().build());

    UpdateItemOperation<FakeItem> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class)
                                                            .item(baseFakeItem)
                                                            .ignoreNulls(true)
                                                            .build());

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    mockDynamoDbEnhancedClientExtension);
    assertThat(request.conditionExpression(), is(nullValue()));
    assertThat(request.expressionAttributeValues().size(), is(0));
}
 

@Test
public void generateRequest_withExtension_correctlyCoalescesIdenticalExpressionNames() {
    FakeItem baseFakeItem = createUniqueFakeItem();
    baseFakeItem.setSubclassAttribute("something");
    Map<String, String> names = singletonMap(SUBCLASS_ATTRIBUTE_NAME, "subclass_attribute");
    Expression condition = Expression.builder().expression("condition").expressionNames(names).build();
    when(mockDynamoDbEnhancedClientExtension.beforeWrite(any(DynamoDbExtensionContext.BeforeWrite.class)))
        .thenReturn(WriteModification.builder().additionalConditionalExpression(condition).build());

    UpdateItemOperation<FakeItem> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class)
                                                            .item(baseFakeItem)
                                                            .ignoreNulls(true)
                                                            .build());

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    mockDynamoDbEnhancedClientExtension);

    assertThat(request.expressionAttributeNames(), is(names));
}
 

@Override
public TransactWriteItem generateTransactWriteItem(TableSchema<T> tableSchema, OperationContext operationContext,
                                                   DynamoDbEnhancedClientExtension dynamoDbEnhancedClientExtension) {
    UpdateItemRequest updateItemRequest = generateRequest(tableSchema, operationContext, dynamoDbEnhancedClientExtension);

    Update update = Update.builder()
                          .key(updateItemRequest.key())
                          .tableName(updateItemRequest.tableName())
                          .updateExpression(updateItemRequest.updateExpression())
                          .conditionExpression(updateItemRequest.conditionExpression())
                          .expressionAttributeValues(updateItemRequest.expressionAttributeValues())
                          .expressionAttributeNames(updateItemRequest.expressionAttributeNames())
                          .build();

    return TransactWriteItem.builder()
                            .update(update)
                            .build();
}
 

@Test
public void generateRequest_withExtension_correctlyCoalescesIdenticalExpressionValues() {
    FakeItem baseFakeItem = createUniqueFakeItem();
    baseFakeItem.setSubclassAttribute("something");
    Map<String, AttributeValue> values = singletonMap(SUBCLASS_ATTRIBUTE_VALUE,
                                                       AttributeValue.builder().s("something").build());
    Expression condition = Expression.builder().expression("condition").expressionValues(values).build();
    when(mockDynamoDbEnhancedClientExtension.beforeWrite(any(DynamoDbExtensionContext.BeforeWrite.class)))
        .thenReturn(WriteModification.builder().additionalConditionalExpression(condition).build());

    UpdateItemOperation<FakeItem> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class)
                                                            .item(baseFakeItem)
                                                            .ignoreNulls(true)
                                                            .build());

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    mockDynamoDbEnhancedClientExtension);

    assertThat(request.expressionAttributeValues(), is(values));
}
 

@Test
public void generateRequest_withExtensions_singleCondition() {
    FakeItem baseFakeItem = createUniqueFakeItem();
    FakeItem fakeItem = createUniqueFakeItem();
    Map<String, AttributeValue> fakeMap = FakeItem.getTableSchema().itemToMap(fakeItem, true);
    Expression condition = Expression.builder().expression("condition").expressionValues(fakeMap).build();
    when(mockDynamoDbEnhancedClientExtension.beforeWrite(any(DynamoDbExtensionContext.BeforeWrite.class)))
        .thenReturn(WriteModification.builder().additionalConditionalExpression(condition).build());
    UpdateItemOperation<FakeItem> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class)
                                                            .item(baseFakeItem)
                                                            .ignoreNulls(true)
                                                            .build());

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    mockDynamoDbEnhancedClientExtension);

    assertThat(request.conditionExpression(), is("condition"));
    assertThat(request.expressionAttributeValues(), is(fakeMap));
}
 

@Test
public void generateRequest_multipleDeletes() {
    FakeItemWithSort item = createUniqueFakeItemWithSort();
    UpdateItemOperation<FakeItemWithSort> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class)
                                                            .item(item)
                                                            .ignoreNulls(false)
                                                            .build());
    Map<String, AttributeValue> expectedKey = new HashMap<>();
    expectedKey.put("id", AttributeValue.builder().s(item.getId()).build());
    expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build());
    Map<String, String> expectedNames = new HashMap<>();
    expectedNames.put(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1");
    expectedNames.put(OTHER_ATTRIBUTE_2_NAME, "other_attribute_2");
    UpdateItemRequest.Builder baseExpectedRequest = UpdateItemRequest.builder()
                                                                     .tableName(TABLE_NAME)
                                                                     .expressionAttributeNames(expectedNames)
                                                                     .key(expectedKey)
                                                                     .returnValues(ReturnValue.ALL_NEW);
    UpdateItemRequest expectedRequest1 =
        baseExpectedRequest.updateExpression("REMOVE " + OTHER_ATTRIBUTE_1_NAME + ", " + OTHER_ATTRIBUTE_2_NAME)
                           .build();
    UpdateItemRequest expectedRequest2 =
        baseExpectedRequest.updateExpression("REMOVE " + OTHER_ATTRIBUTE_2_NAME + ", " + OTHER_ATTRIBUTE_1_NAME)
                           .build();

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    null);

    assertThat(request,either(is(expectedRequest1)).or(is(expectedRequest2)));
}
 

private void setupUpdateItemTest() throws Exception {
    TimeoutException te = setRuleForDependencyTimeout();

    when(leaseSerializer.getDynamoHashKey(any(Lease.class))).thenReturn(Collections.emptyMap());
    when(leaseSerializer.getDynamoLeaseCounterExpectation(any(Lease.class))).thenReturn(Collections.emptyMap());
    when(leaseSerializer.getDynamoLeaseCounterUpdate(any(Lease.class))).thenReturn(Collections.emptyMap());
    when(leaseSerializer.getDynamoTakeLeaseUpdate(any(), anyString())).thenReturn(Collections.emptyMap());

    when(dynamoDbClient.updateItem(any(UpdateItemRequest.class))).thenReturn(mockUpdateFuture);
    when(mockUpdateFuture.get(anyLong(), any())).thenThrow(te);
}
 

@Test
public void generateRequest_canIgnoreNullValues() {
    FakeItemWithSort item = createUniqueFakeItemWithSort();
    item.setOtherAttribute1("value-1");
    UpdateItemOperation<FakeItemWithSort> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class)
                                                            .item(item)
                                                            .ignoreNulls(true)
                                                            .build());
    Map<String, AttributeValue> expectedKey = new HashMap<>();
    expectedKey.put("id", AttributeValue.builder().s(item.getId()).build());
    expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build());
    Map<String, AttributeValue> expectedValues =
        singletonMap(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build());
    Map<String, String> expectedNames = singletonMap(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1");
    UpdateItemRequest expectedRequest = UpdateItemRequest.builder()
        .tableName(TABLE_NAME)
        .updateExpression("SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE)
        .expressionAttributeValues(expectedValues)
        .expressionAttributeNames(expectedNames)
        .key(expectedKey)
        .returnValues(ReturnValue.ALL_NEW)
        .build();


    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    null);

    assertThat(request, is(expectedRequest));
}
 

@Test
public void generateRequest_withExtension_modifiesUpdateExpression() {
    FakeItem fakeItem = createUniqueFakeItem();
    Map<String, AttributeValue> baseMap = new HashMap<>(FakeItem.getTableSchema().itemToMap(fakeItem, true));

    Map<String, AttributeValue> fakeMap = new HashMap<>(baseMap);
    fakeMap.put("subclass_attribute", AttributeValue.builder().s("1").build());

    when(mockDynamoDbEnhancedClientExtension.beforeWrite(any(DynamoDbExtensionContext.BeforeWrite.class)))
        .thenReturn(WriteModification.builder().transformedItem(fakeMap).build());


    UpdateItemOperation<FakeItem> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class)
                                                            .item(fakeItem)
                                                            .ignoreNulls(true)
                                                            .build());

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    mockDynamoDbEnhancedClientExtension);

    assertThat(request.updateExpression(), is("SET " + SUBCLASS_ATTRIBUTE_NAME + " = " + SUBCLASS_ATTRIBUTE_VALUE));
    assertThat(request.expressionAttributeValues(), hasEntry(SUBCLASS_ATTRIBUTE_VALUE,
        AttributeValue.builder().s("1").build()));
    assertThat(request.expressionAttributeNames(), hasEntry(SUBCLASS_ATTRIBUTE_NAME, "subclass_attribute"));
}
 

@Test
public void generateRequest_withExtension_conditionAndModification() {
    FakeItem baseFakeItem = createUniqueFakeItem();
    Map<String, AttributeValue> baseMap = new HashMap<>(FakeItem.getTableSchema().itemToMap(baseFakeItem, true));

    Map<String, AttributeValue> fakeMap = new HashMap<>(baseMap);
    fakeMap.put("subclass_attribute", AttributeValue.builder().s("1").build());

    Map<String, AttributeValue> conditionValues = new HashMap<>();
    conditionValues.put(":condition_value", AttributeValue.builder().s("2").build());

    when(mockDynamoDbEnhancedClientExtension.beforeWrite(any(DynamoDbExtensionContext.BeforeWrite.class)))
        .thenReturn(WriteModification.builder()
                                     .transformedItem(fakeMap)
                                     .additionalConditionalExpression(Expression.builder()
                                                                                .expression("condition")
                                                                                .expressionValues(conditionValues)
                                                                                .build())
                                     .build());

    UpdateItemOperation<FakeItem> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class)
                                                            .item(baseFakeItem)
                                                            .ignoreNulls(true)
                                                            .build());

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItem.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    mockDynamoDbEnhancedClientExtension);

    assertThat(request.updateExpression(), is("SET " + SUBCLASS_ATTRIBUTE_NAME + " = " + SUBCLASS_ATTRIBUTE_VALUE));
    assertThat(request.expressionAttributeValues(), hasEntry(SUBCLASS_ATTRIBUTE_VALUE,
                                                             AttributeValue.builder().s("1").build()));
    assertThat(request.expressionAttributeValues(), hasEntry(":condition_value",
                                                             AttributeValue.builder().s("2").build()));
    assertThat(request.expressionAttributeNames(), hasEntry(SUBCLASS_ATTRIBUTE_NAME, "subclass_attribute"));
}
 

@Test
public void generateTransactWriteItem_basicRequest() {
    FakeItem fakeItem = createUniqueFakeItem();
    Map<String, AttributeValue> fakeItemMap = FakeItem.getTableSchema().itemToMap(fakeItem, true);
    UpdateItemOperation<FakeItem> updateItemOperation =
        spy(UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class).item(fakeItem).build()));
    OperationContext context = DefaultOperationContext.create(TABLE_NAME, TableMetadata.primaryIndexName());
    String updateExpression = "update-expression";
    Map<String, AttributeValue> attributeValues = Collections.singletonMap("key", stringValue("value1"));
    Map<String, String> attributeNames = Collections.singletonMap("key", "value2");

    UpdateItemRequest updateItemRequest = UpdateItemRequest.builder()
                                                           .tableName(TABLE_NAME)
                                                           .key(fakeItemMap)
                                                           .updateExpression(updateExpression)
                                                           .expressionAttributeValues(attributeValues)
                                                           .expressionAttributeNames(attributeNames)
                                                           .build();
    doReturn(updateItemRequest).when(updateItemOperation).generateRequest(any(), any(), any());

    TransactWriteItem actualResult = updateItemOperation.generateTransactWriteItem(FakeItem.getTableSchema(),
                                                                                   context,
                                                                                   mockDynamoDbEnhancedClientExtension);

    TransactWriteItem expectedResult = TransactWriteItem.builder()
                                                        .update(Update.builder()
                                                                      .key(fakeItemMap)
                                                                      .tableName(TABLE_NAME)
                                                                      .updateExpression(updateExpression)
                                                                      .expressionAttributeNames(attributeNames)
                                                                      .expressionAttributeValues(attributeValues)
                                                                      .build())
                                                        .build();
    assertThat(actualResult, is(expectedResult));
    verify(updateItemOperation).generateRequest(FakeItem.getTableSchema(), context, mockDynamoDbEnhancedClientExtension);
}
 

@Test
public void generateTransactWriteItem_conditionalRequest() {
    FakeItem fakeItem = createUniqueFakeItem();
    Map<String, AttributeValue> fakeItemMap = FakeItem.getTableSchema().itemToMap(fakeItem, true);
    UpdateItemOperation<FakeItem> updateItemOperation =
        spy(UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItem.class).item(fakeItem).build()));
    OperationContext context = DefaultOperationContext.create(TABLE_NAME, TableMetadata.primaryIndexName());
    String updateExpression = "update-expression";
    String conditionExpression = "condition-expression";
    Map<String, AttributeValue> attributeValues = Collections.singletonMap("key", stringValue("value1"));
    Map<String, String> attributeNames = Collections.singletonMap("key", "value2");

    UpdateItemRequest updateItemRequest = UpdateItemRequest.builder()
                                                           .tableName(TABLE_NAME)
                                                           .key(fakeItemMap)
                                                           .updateExpression(updateExpression)
                                                           .conditionExpression(conditionExpression)
                                                           .expressionAttributeValues(attributeValues)
                                                           .expressionAttributeNames(attributeNames)
                                                           .build();
    doReturn(updateItemRequest).when(updateItemOperation).generateRequest(any(), any(), any());

    TransactWriteItem actualResult = updateItemOperation.generateTransactWriteItem(FakeItem.getTableSchema(),
                                                                                   context,
                                                                                   mockDynamoDbEnhancedClientExtension);

    TransactWriteItem expectedResult = TransactWriteItem.builder()
                                                        .update(Update.builder()
                                                                      .key(fakeItemMap)
                                                                      .tableName(TABLE_NAME)
                                                                      .updateExpression(updateExpression)
                                                                      .conditionExpression(conditionExpression)
                                                                      .expressionAttributeNames(attributeNames)
                                                                      .expressionAttributeValues(attributeValues)
                                                                      .build())
                                                        .build();
    assertThat(actualResult, is(expectedResult));
    verify(updateItemOperation).generateRequest(FakeItem.getTableSchema(), context, mockDynamoDbEnhancedClientExtension);
}
 

@Override
@NonNull
public Optional<SimpleLock> lock(@NonNull LockConfiguration lockConfiguration) {
    String nowIso = toIsoString(now());
    String lockUntilIso = toIsoString(lockConfiguration.getLockAtMostUntil());

    Map<String, AttributeValue> key = singletonMap(ID, attr(lockConfiguration.getName()));

    Map<String, AttributeValue> attributeUpdates = new HashMap<>(3);
    attributeUpdates.put(":lockUntil", attr(lockUntilIso));
    attributeUpdates.put(":lockedAt", attr(nowIso));
    attributeUpdates.put(":lockedBy", attr(hostname));

    UpdateItemRequest request = UpdateItemRequest.builder()
            .tableName(tableName)
            .key(key)
            .updateExpression(OBTAIN_LOCK_QUERY)
            .conditionExpression(OBTAIN_LOCK_CONDITION)
            .expressionAttributeValues(attributeUpdates)
            .returnValues(ReturnValue.UPDATED_NEW)
            .build();

    try {
        // There are three possible situations:
        // 1. The lock document does not exist yet - it is inserted - we have the lock
        // 2. The lock document exists and lockUtil <= now - it is updated - we have the lock
        // 3. The lock document exists and lockUtil > now - ConditionalCheckFailedException is thrown
        dynamoDbClient.updateItem(request);
        return Optional.of(new DynamoDBLock(dynamoDbClient, tableName, lockConfiguration));
    } catch (ConditionalCheckFailedException e) {
        // Condition failed. This means there was a lock with lockUntil > now.
        return Optional.empty();
    }
}
 

private boolean updateRow(String key, String appid, Map<String, AttributeValue> row) {
	if (StringUtils.isBlank(key) || StringUtils.isBlank(appid) || row == null || row.isEmpty()) {
		return false;
	}
	String table = getTableNameForAppid(appid);
	try {
		UpdateItemRequest.Builder updateRequest = UpdateItemRequest.builder();
		StringBuilder updateExpression = new StringBuilder("SET ");
		Map<String, String> names = new HashMap<>(row.size() + 1);
		Map<String, AttributeValue> values = new HashMap<>(row.size() + 1);
		boolean isLockingEnabledForRow = false;
		AttributeValue version = row.remove(Config._VERSION); // ignore the version field here
		if (version == null || version.n() == null) {
			version = AttributeValue.builder().n("0").build();
		}
		if (Long.parseLong(version.n()) > 0L) {
			isLockingEnabledForRow = true;
		}
		for (Entry<String, AttributeValue> attr : row.entrySet()) {
			String name = "#" + attr.getKey();
			String value = ":" + attr.getKey();
			updateExpression.append(name).append("=").append(value).append(",");
			names.put(name, attr.getKey());
			values.put(value, attr.getValue());
		}
		updateExpression.setLength(updateExpression.length() - 1); // remove comma at the end

		if (isLockingEnabledForRow) {
			names.put("#" + Config._VERSION, Config._VERSION);
			values.put(":" + Config._VERSION, version);
			values.put(":plusOne", AttributeValue.builder().n("1").build());
			updateRequest.conditionExpression("#" + Config._VERSION + " = :" + Config._VERSION);
			updateExpression.append(" ADD #").append(Config._VERSION).append(" :plusOne");
		}

		updateRequest.tableName(table);
		updateRequest.key(rowKey(key, appid));
		updateRequest.expressionAttributeNames(names);
		updateRequest.expressionAttributeValues(values);
		updateRequest.updateExpression(updateExpression.toString());
		client().updateItem(updateRequest.build());
		return true;
	} catch (ConditionalCheckFailedException ex) {
		logger.warn("Item not updated - versions don't match. table={}, appid={}, key={}.", table, appid, key);
	} catch (Exception e) {
		logger.error("Could not update row in DB - table={}, appid={}, key={}:", table, appid, key, e);
		throwIfNecessary(e);
	}
	return false;
}
 

@Test
public void generateRequest_explicitlyUnsetIgnoreNulls() {
    FakeItemWithSort item = createUniqueFakeItemWithSort();
    item.setOtherAttribute1("value-1");
    UpdateItemOperation<FakeItemWithSort> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class)
                                                            .item(item)
                                                            .ignoreNulls(false)
                                                            .build());
    Map<String, AttributeValue> expectedKey = new HashMap<>();
    expectedKey.put("id", AttributeValue.builder().s(item.getId()).build());
    expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build());
    Map<String, AttributeValue> expectedValues = new HashMap<>();
    expectedValues.put(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build());
    Map<String, String> expectedNames = new HashMap<>();
    expectedNames.put(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1");
    expectedNames.put(OTHER_ATTRIBUTE_2_NAME, "other_attribute_2");
    UpdateItemRequest.Builder baseExpectedRequest = UpdateItemRequest.builder()
                                                                     .tableName(TABLE_NAME)
                                                                     .expressionAttributeValues(expectedValues)
                                                                     .expressionAttributeNames(expectedNames)
                                                                     .key(expectedKey)
                                                                     .returnValues(ReturnValue.ALL_NEW);
    UpdateItemRequest expectedRequest =
        baseExpectedRequest.updateExpression("SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE +
                                             " REMOVE " + OTHER_ATTRIBUTE_2_NAME)
                           .build();

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    null);

    assertThat(request, is(expectedRequest));
}
 

@Test
public void generateRequest_withConditionExpression() {
    FakeItemWithSort item = createUniqueFakeItemWithSort();
    item.setOtherAttribute1("value-1");
    UpdateItemOperation<FakeItemWithSort> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class)
                                                            .item(item)
                                                            .conditionExpression(CONDITION_EXPRESSION)
                                                            .build());
    Map<String, AttributeValue> expectedKey = new HashMap<>();
    expectedKey.put("id", AttributeValue.builder().s(item.getId()).build());
    expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build());
    Map<String, AttributeValue> expectedValues = new HashMap<>(CONDITION_EXPRESSION.expressionValues());
    expectedValues.put(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build());
    Map<String, String> expectedNames = new HashMap<>(CONDITION_EXPRESSION.expressionNames());
    expectedNames.put(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1");
    expectedNames.put(OTHER_ATTRIBUTE_2_NAME, "other_attribute_2");
    UpdateItemRequest.Builder baseExpectedRequest =
        UpdateItemRequest.builder()
                         .tableName(TABLE_NAME)
                         .expressionAttributeValues(expectedValues)
                         .expressionAttributeNames(expectedNames)
                         .conditionExpression(CONDITION_EXPRESSION.expression())
                         .key(expectedKey)
                         .returnValues(ReturnValue.ALL_NEW);
    UpdateItemRequest expectedRequest =
        baseExpectedRequest.updateExpression("SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE +
                                             " REMOVE " + OTHER_ATTRIBUTE_2_NAME)
                           .build();

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    null);

    assertThat(request, is(expectedRequest));
}
 

@Test
public void generateRequest_nullValuesNotIgnoredByDefault() {
    FakeItemWithSort item = createUniqueFakeItemWithSort();
    item.setOtherAttribute1("value-1");
    UpdateItemOperation<FakeItemWithSort> updateItemOperation = UpdateItemOperation.create(
        UpdateItemEnhancedRequest.builder(FakeItemWithSort.class).item(item).build());
    Map<String, AttributeValue> expectedKey = new HashMap<>();
    expectedKey.put("id", AttributeValue.builder().s(item.getId()).build());
    expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build());
    Map<String, AttributeValue> expectedValues = new HashMap<>();
    expectedValues.put(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build());
    Map<String, String> expectedNames = new HashMap<>();
    expectedNames.put(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1");
    expectedNames.put(OTHER_ATTRIBUTE_2_NAME, "other_attribute_2");
    UpdateItemRequest.Builder baseExpectedRequest = UpdateItemRequest.builder()
                                                                     .tableName(TABLE_NAME)
                                                                     .expressionAttributeValues(expectedValues)
                                                                     .expressionAttributeNames(expectedNames)
                                                                     .key(expectedKey)
                                                                     .returnValues(ReturnValue.ALL_NEW);
    UpdateItemRequest expectedRequest =
        baseExpectedRequest.updateExpression("SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE +
                                             " REMOVE " + OTHER_ATTRIBUTE_2_NAME)
                           .build();

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    null);

    assertThat(request, is(expectedRequest));
}
 

@Test
public void getServiceCall_makesTheRightCallAndReturnsResponse() {
    FakeItem item = createUniqueFakeItem();
    UpdateItemOperation<FakeItem> updateItemOperation = UpdateItemOperation.create(
        UpdateItemEnhancedRequest.builder(FakeItem.class).item(item).build());
    UpdateItemRequest updateItemRequest = UpdateItemRequest.builder().tableName(TABLE_NAME).build();
    UpdateItemResponse expectedResponse = UpdateItemResponse.builder().build();
    when(mockDynamoDbClient.updateItem(any(UpdateItemRequest.class))).thenReturn(expectedResponse);

    UpdateItemResponse response = updateItemOperation.serviceCall(mockDynamoDbClient).apply(updateItemRequest);

    assertThat(response, sameInstance(expectedResponse));
    verify(mockDynamoDbClient).updateItem(updateItemRequest);
}
 

@Test
public void generateRequest_multipleSetters() {
    FakeItemWithSort item = createUniqueFakeItemWithSort();
    item.setOtherAttribute1("value-1");
    item.setOtherAttribute2("value-2");
    UpdateItemOperation<FakeItemWithSort> updateItemOperation =
        UpdateItemOperation.create(UpdateItemEnhancedRequest.builder(FakeItemWithSort.class)
                                                            .item(item)
                                                            .ignoreNulls(false)
                                                            .build());
    Map<String, AttributeValue> expectedKey = new HashMap<>();
    expectedKey.put("id", AttributeValue.builder().s(item.getId()).build());
    expectedKey.put("sort", AttributeValue.builder().s(item.getSort()).build());
    Map<String, AttributeValue> expectedValues = new HashMap<>();
    expectedValues.put(OTHER_ATTRIBUTE_1_VALUE, AttributeValue.builder().s("value-1").build());
    expectedValues.put(OTHER_ATTRIBUTE_2_VALUE, AttributeValue.builder().s("value-2").build());
    Map<String, String> expectedNames = new HashMap<>();
    expectedNames.put(OTHER_ATTRIBUTE_1_NAME, "other_attribute_1");
    expectedNames.put(OTHER_ATTRIBUTE_2_NAME, "other_attribute_2");
    UpdateItemRequest.Builder baseExpectedRequest = UpdateItemRequest.builder()
                                                                     .tableName(TABLE_NAME)
                                                                     .expressionAttributeValues(expectedValues)
                                                                     .expressionAttributeNames(expectedNames)
                                                                     .key(expectedKey)
                                                                     .returnValues(ReturnValue.ALL_NEW);
    UpdateItemRequest expectedRequest1 =
        baseExpectedRequest.updateExpression("SET " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE +
                                             ", " + OTHER_ATTRIBUTE_2_NAME + " = " + OTHER_ATTRIBUTE_2_VALUE)
                           .build();
    UpdateItemRequest expectedRequest2 =
        baseExpectedRequest.updateExpression("SET " + OTHER_ATTRIBUTE_2_NAME + " = " + OTHER_ATTRIBUTE_2_VALUE +
                                             ", " + OTHER_ATTRIBUTE_1_NAME + " = " + OTHER_ATTRIBUTE_1_VALUE)
                           .build();

    UpdateItemRequest request = updateItemOperation.generateRequest(FakeItemWithSort.getTableSchema(),
                                                                    PRIMARY_CONTEXT,
                                                                    null);

    assertThat(request, either(is(expectedRequest1)).or(is(expectedRequest2)));
}
 

@Test
public void updateApplication_author() {
  String userId = UUID.randomUUID().toString();
  String applicationId = UUID.randomUUID().toString();
  String author = UUID.randomUUID().toString();
  Long version = 1L;
  Map<String, AttributeValue> recordMap = keyMap(userId, applicationId);
  recordMap.put("version", AttributeValue.builder().n(version.toString()).build());
  GetItemResponse response = GetItemResponse.builder()
        .item(recordMap)
        .build();
  Map<String, AttributeValue> updateMap = new HashMap<>();
  updateMap.put(":nv", AttributeValue.builder().n("2").build());
  updateMap.put(":v", AttributeValue.builder().n("1").build());
  updateMap.put(":a", AttributeValue.builder().s(author).build());

  when(principal.getName()).thenReturn(userId);
  when(dynamodb.getItem(any(GetItemRequest.class))).thenReturn(response);

  UpdateApplicationInput input = new UpdateApplicationInput().author(author);
  UpdateItemRequest expectedUpdateItemRequest = UpdateItemRequest.builder()
        .tableName(TABLE_NAME)
        .key(keyMap(userId, applicationId))
        .updateExpression("SET author = :a,version = :nv")
        .expressionAttributeValues(updateMap)
        .conditionExpression("version = :v")
        .build();

  Application application = service.updateApplication(input, applicationId);
  ArgumentCaptor<UpdateItemRequest> updateItemRequestArgumentCaptor = ArgumentCaptor.forClass(UpdateItemRequest.class);
  verify(dynamodb).updateItem(updateItemRequestArgumentCaptor.capture());

  UpdateItemRequest updateItemRequest = updateItemRequestArgumentCaptor.getValue();
  assertThat(updateItemRequest).isEqualTo(expectedUpdateItemRequest);
  assertThat(application.getApplicationId()).isEqualTo(applicationId);
  assertThat(application.getAuthor()).isEqualTo(author);
}
 

@Test
public void updateApplication_description() {
  String userId = UUID.randomUUID().toString();
  String applicationId = UUID.randomUUID().toString();
  String description = UUID.randomUUID().toString();
  Long version = 1L;
  Map<String, AttributeValue> recordMap = keyMap(userId, applicationId);
  recordMap.put("version", AttributeValue.builder().n(version.toString()).build());
  GetItemResponse response = GetItemResponse.builder()
        .item(recordMap)
        .build();
  Map<String, AttributeValue> updateMap = new HashMap<>();
  updateMap.put(":nv", AttributeValue.builder().n("2").build());
  updateMap.put(":v", AttributeValue.builder().n("1").build());
  updateMap.put(":d", AttributeValue.builder().s(description).build());

  when(principal.getName()).thenReturn(userId);
  when(dynamodb.getItem(any(GetItemRequest.class))).thenReturn(response);

  UpdateApplicationInput input = new UpdateApplicationInput().description(description);
  UpdateItemRequest expectedUpdateItemRequest = UpdateItemRequest.builder()
        .tableName(TABLE_NAME)
        .key(keyMap(userId, applicationId))
        .updateExpression("SET description = :d,version = :nv")
        .expressionAttributeValues(updateMap)
        .conditionExpression("version = :v")
        .build();

  Application application = service.updateApplication(input, applicationId);
  ArgumentCaptor<UpdateItemRequest> updateItemRequestArgumentCaptor = ArgumentCaptor.forClass(UpdateItemRequest.class);
  verify(dynamodb).updateItem(updateItemRequestArgumentCaptor.capture());

  UpdateItemRequest updateItemRequest = updateItemRequestArgumentCaptor.getValue();
  assertThat(updateItemRequest).isEqualTo(expectedUpdateItemRequest);
  assertThat(application.getApplicationId()).isEqualTo(applicationId);
  assertThat(application.getDescription()).isEqualTo(description);
}
 

@Test
public void updateApplication_homePageUrl() {
  String userId = UUID.randomUUID().toString();
  String applicationId = UUID.randomUUID().toString();
  String homePageUrl = UUID.randomUUID().toString();
  Long version = 1L;
  Map<String, AttributeValue> recordMap = keyMap(userId, applicationId);
  recordMap.put("version", AttributeValue.builder().n(version.toString()).build());
  GetItemResponse response = GetItemResponse.builder()
        .item(recordMap)
        .build();
  Map<String, AttributeValue> updateMap = new HashMap<>();
  updateMap.put(":nv", AttributeValue.builder().n("2").build());
  updateMap.put(":v", AttributeValue.builder().n("1").build());
  updateMap.put(":h", AttributeValue.builder().s(homePageUrl).build());

  when(principal.getName()).thenReturn(userId);
  when(dynamodb.getItem(any(GetItemRequest.class))).thenReturn(response);

  UpdateApplicationInput input = new UpdateApplicationInput().homePageUrl(homePageUrl);
  UpdateItemRequest expectedUpdateItemRequest = UpdateItemRequest.builder()
        .tableName(TABLE_NAME)
        .key(keyMap(userId, applicationId))
        .updateExpression("SET homePageUrl = :h,version = :nv")
        .expressionAttributeValues(updateMap)
        .conditionExpression("version = :v")
        .build();

  Application application = service.updateApplication(input, applicationId);
  ArgumentCaptor<UpdateItemRequest> updateItemRequestArgumentCaptor = ArgumentCaptor.forClass(UpdateItemRequest.class);
  verify(dynamodb).updateItem(updateItemRequestArgumentCaptor.capture());

  UpdateItemRequest updateItemRequest = updateItemRequestArgumentCaptor.getValue();
  assertThat(updateItemRequest).isEqualTo(expectedUpdateItemRequest);
  assertThat(application.getApplicationId()).isEqualTo(applicationId);
  assertThat(application.getHomePageUrl()).isEqualTo(homePageUrl);
}