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

@Test
public void testBatchWriteTooManyItemsErrorHandling() throws Exception {
    int itemNumber = 26;
    HashMap<String, List<WriteRequest>> requestItems = new HashMap<String, List<WriteRequest>>();
    List<WriteRequest> writeRequests = new ArrayList<WriteRequest>();
    for (int i = 0; i < itemNumber; i++) {
        HashMap<String, AttributeValue> writeAttributes = new HashMap<String, AttributeValue>();
        writeAttributes.put(HASH_KEY_NAME, AttributeValue.builder().s("" + System.currentTimeMillis()).build());
        writeAttributes.put("bar", AttributeValue.builder().s("" + System.currentTimeMillis()).build());
        writeRequests.add(WriteRequest.builder().putRequest(PutRequest.builder().item(writeAttributes).build()).build());
    }
    requestItems.put(tableName, writeRequests);
    try {
        dynamo.batchWriteItem(BatchWriteItemRequest.builder().requestItems(requestItems).build());
    } catch (AwsServiceException exception) {
        assertEquals("ValidationException", exception.awsErrorDetails().errorCode());
        assertNotEmpty(exception.awsErrorDetails().errorMessage());
        assertNotEmpty(exception.requestId());
        assertNotEmpty(exception.awsErrorDetails().serviceName());
        assertEquals(400, exception.statusCode());
    }
}
 

@Override
public BatchWriteItemRequest generateRequest(DynamoDbEnhancedClientExtension extension) {
    Map<String, List<WriteRequest>> allRequestItems = new HashMap<>();

    request.writeBatches().forEach(writeBatch -> {
        Collection<WriteRequest> writeRequestsForTable = allRequestItems.computeIfAbsent(
            writeBatch.tableName(),
            ignored -> new ArrayList<>());
        writeRequestsForTable.addAll(writeBatch.writeRequests());
    });

    return BatchWriteItemRequest.builder()
                                .requestItems(
                                    Collections.unmodifiableMap(CollectionUtils.deepCopyMap(allRequestItems)))
                                .build();
}
 

@Test
public void getServiceCall_makesTheRightCallAndReturnsResponse() {

    WriteBatch batch = WriteBatch.builder(FakeItem.class)
                                 .mappedTableResource(fakeItemMappedTable)
                                 .addPutItem(r -> r.item(FAKE_ITEMS.get(0)))
                                 .build();

    BatchWriteItemEnhancedRequest batchWriteItemEnhancedRequest =
        BatchWriteItemEnhancedRequest.builder()
                                     .writeBatches(batch)
                                     .build();

    BatchWriteItemOperation operation = BatchWriteItemOperation.create(batchWriteItemEnhancedRequest);

    WriteRequest writeRequest =
        WriteRequest.builder()
                    .putRequest(PutRequest.builder().item(FAKE_ITEM_MAPS.get(0)).build())
                    .build();

    BatchWriteItemRequest request =
        BatchWriteItemRequest.builder()
                             .requestItems(singletonMap("table", singletonList(writeRequest)))
                             .build();

    BatchWriteItemResponse expectedResponse = BatchWriteItemResponse.builder().build();
    when(mockDynamoDbClient.batchWriteItem(any(BatchWriteItemRequest.class))).thenReturn(expectedResponse);

    BatchWriteItemResponse response = operation.serviceCall(mockDynamoDbClient).apply(request);

    assertThat(response, sameInstance(expectedResponse));
    verify(mockDynamoDbClient).batchWriteItem(request);
}
 

@Test
public void generateRequest_multipleTables_mixedCommands_usingShortcutForm() {
    BatchWriteItemEnhancedRequest batchWriteItemEnhancedRequest =
        BatchWriteItemEnhancedRequest.builder()
                                     .writeBatches(
                                         WriteBatch.builder(FakeItem.class)
                                                   .mappedTableResource(fakeItemMappedTable)
                                                   .addPutItem(FAKE_ITEMS.get(0))
                                                   .addDeleteItem(FAKE_ITEM_KEYS.get(1))
                                                   .addPutItem(FAKE_ITEMS.get(2))
                                                   .build(),
                                         WriteBatch.builder(FakeItemWithSort.class)
                                                   .mappedTableResource(fakeItemWithSortMappedTable)
                                                   .addDeleteItem(FAKESORT_ITEM_KEYS.get(0))
                                                   .addPutItem(FAKESORT_ITEMS.get(1))
                                                   .addDeleteItem(FAKESORT_ITEM_KEYS.get(2))
                                                   .build())
                                     .build();

    BatchWriteItemOperation operation = BatchWriteItemOperation.create(batchWriteItemEnhancedRequest);

    BatchWriteItemRequest request = operation.generateRequest(mockExtension);

    List<WriteRequest> writeRequests1 = request.requestItems().get(TABLE_NAME);
    List<WriteRequest> writeRequests2 = request.requestItems().get(TABLE_NAME_2);
    assertThat(writeRequests1, containsInAnyOrder(putRequest(FAKE_ITEM_MAPS.get(0)),
                                                  deleteRequest(FAKE_ITEM_MAPS.get(1)),
                                                  putRequest(FAKE_ITEM_MAPS.get(2))));
    assertThat(writeRequests2, containsInAnyOrder(deleteRequest(FAKESORT_ITEM_MAPS.get(0)),
                                                  putRequest(FAKESORT_ITEM_MAPS.get(1)),
                                                  deleteRequest(FAKESORT_ITEM_MAPS.get(2))));
}
 

@Test
public void generateRequest_multipleTables_mixedCommands_usingKeyItemForm() {
    BatchWriteItemEnhancedRequest batchWriteItemEnhancedRequest =
        BatchWriteItemEnhancedRequest.builder()
                                     .writeBatches(
                                         WriteBatch.builder(FakeItem.class)
                                                   .mappedTableResource(fakeItemMappedTable)
                                                   .addPutItem(FAKE_ITEMS.get(0))
                                                   .addDeleteItem(FAKE_ITEMS.get(1))
                                                   .addPutItem(FAKE_ITEMS.get(2))
                                                   .build(),
                                         WriteBatch.builder(FakeItemWithSort.class)
                                                   .mappedTableResource(fakeItemWithSortMappedTable)
                                                   .addDeleteItem(FAKESORT_ITEMS.get(0))
                                                   .addPutItem(FAKESORT_ITEMS.get(1))
                                                   .addDeleteItem(FAKESORT_ITEMS.get(2))
                                                   .build())
                                     .build();

    BatchWriteItemOperation operation = BatchWriteItemOperation.create(batchWriteItemEnhancedRequest);

    BatchWriteItemRequest request = operation.generateRequest(mockExtension);

    List<WriteRequest> writeRequests1 = request.requestItems().get(TABLE_NAME);
    List<WriteRequest> writeRequests2 = request.requestItems().get(TABLE_NAME_2);
    assertThat(writeRequests1, containsInAnyOrder(putRequest(FAKE_ITEM_MAPS.get(0)),
                                                  deleteRequest(FAKE_ITEM_MAPS.get(1)),
                                                  putRequest(FAKE_ITEM_MAPS.get(2))));
    assertThat(writeRequests2, containsInAnyOrder(deleteRequest(FAKESORT_ITEM_MAPS.get(0)),
                                                  putRequest(FAKESORT_ITEM_MAPS.get(1)),
                                                  deleteRequest(FAKESORT_ITEM_MAPS.get(2))));
}
 

static List<Map<String, AttributeValue>> generateTestData(String tableName, int numOfItems) {
  BatchWriteItemRequest batchWriteItemRequest =
      generateBatchWriteItemRequest(tableName, numOfItems);

  dynamoDBClient.batchWriteItem(batchWriteItemRequest);
  ScanResponse scanResult =
      dynamoDBClient.scan(ScanRequest.builder().tableName(tableName).build());

  List<Map<String, AttributeValue>> items = scanResult.items();
  Assert.assertEquals(numOfItems, items.size());
  return items;
}
 

static BatchWriteItemRequest generateBatchWriteItemRequest(String tableName, int numOfItems) {
  BatchWriteItemRequest batchWriteItemRequest =
      BatchWriteItemRequest.builder()
          .requestItems(ImmutableMap.of(tableName, generateWriteRequests(numOfItems)))
          .build();
  return batchWriteItemRequest;
}
 

void deleteEntriesPerBatch(List<WriteRequest> deleteRequests) {
    final int chunkSize = 25;
    final AtomicInteger counter = new AtomicInteger();
    final Collection<List<WriteRequest>> deleteRequestsSplittedByChunkSize = deleteRequests.stream()
            .collect(Collectors.groupingBy(it -> counter.getAndIncrement() / chunkSize))
            .values();

    deleteRequestsSplittedByChunkSize.forEach
            (currentDeleteRequests -> dynamoDbClient.batchWriteItem(
                    BatchWriteItemRequest.builder().requestItems(
                            ImmutableMap.of(tableName, currentDeleteRequests)).build()));

}
 

@Override
public Function<BatchWriteItemRequest, BatchWriteItemResponse> serviceCall(DynamoDbClient dynamoDbClient) {
    return dynamoDbClient::batchWriteItem;
}
 

@Override
public Function<BatchWriteItemRequest, CompletableFuture<BatchWriteItemResponse>> asyncServiceCall(
    DynamoDbAsyncClient dynamoDbAsyncClient) {

    return dynamoDbAsyncClient::batchWriteItem;
}
 

@Test
public void generateRequest_multipleTables_extensionOnlyTransformsPutsAndNotDeletes() {

    // Use the mock extension to transform every item based on table name
    IntStream.range(0, 3).forEach(i -> {
        lenient().doReturn(WriteModification.builder().transformedItem(FAKE_ITEM_MAPS.get(i + 3)).build())
            .when(mockExtension)
            .beforeWrite(
                     argThat(extensionContext ->
                                 extensionContext.operationContext().tableName().equals(TABLE_NAME) &&
                                 extensionContext.items().equals(FAKE_ITEM_MAPS.get(i))
                     ));
        lenient().doReturn(WriteModification.builder().transformedItem(FAKESORT_ITEM_MAPS.get(i + 3)).build())
            .when(mockExtension)
            .beforeWrite(
                argThat(extensionContext ->
                            extensionContext.operationContext().tableName().equals(TABLE_NAME_2) &&
                            extensionContext.items().equals(FAKESORT_ITEM_MAPS.get(i))
                ));
    });

    BatchWriteItemEnhancedRequest batchWriteItemEnhancedRequest =
        BatchWriteItemEnhancedRequest.builder()
                                     .writeBatches(
                                         WriteBatch.builder(FakeItem.class)
                                                   .mappedTableResource(fakeItemMappedTableWithExtension)
                                                   .addPutItem(r -> r.item(FAKE_ITEMS.get(0)))
                                                   .addDeleteItem(r -> r.key(FAKE_ITEM_KEYS.get(1)))
                                                   .addPutItem(r -> r.item(FAKE_ITEMS.get(2)))
                                                   .build(),
                                         WriteBatch.builder(FakeItemWithSort.class)
                                                   .mappedTableResource(fakeItemWithSortMappedTableWithExtension)
                                                   .addDeleteItem(r -> r.key(FAKESORT_ITEM_KEYS.get(0)))
                                                   .addPutItem(r -> r.item(FAKESORT_ITEMS.get(1)))
                                                   .addDeleteItem(r -> r.key(FAKESORT_ITEM_KEYS.get(2)))
                                                   .build())
                                     .build();

    BatchWriteItemOperation operation = BatchWriteItemOperation.create(batchWriteItemEnhancedRequest);

    BatchWriteItemRequest request = operation.generateRequest(mockExtension);

    List<WriteRequest> writeRequests1 = request.requestItems().get(TABLE_NAME);
    List<WriteRequest> writeRequests2 = request.requestItems().get(TABLE_NAME_2);

    // Only PutItem requests should have their attributes transformed
    assertThat(writeRequests1, containsInAnyOrder(putRequest(FAKE_ITEM_MAPS.get(3)),
                                                  deleteRequest(FAKE_ITEM_MAPS.get(1)),
                                                  putRequest(FAKE_ITEM_MAPS.get(5))));
    assertThat(writeRequests2, containsInAnyOrder(deleteRequest(FAKESORT_ITEM_MAPS.get(0)),
                                                  putRequest(FAKESORT_ITEM_MAPS.get(4)),
                                                  deleteRequest(FAKESORT_ITEM_MAPS.get(2))));
}
 

private void flushBatch() throws IOException, InterruptedException {
  if (batch.isEmpty()) {
    return;
  }

  try {
    // Since each element is a KV<tableName, writeRequest> in the batch, we need to group them
    // by tableName
    Map<String, List<WriteRequest>> mapTableRequest =
        batch.stream()
            .collect(
                Collectors.groupingBy(
                    KV::getKey, Collectors.mapping(KV::getValue, Collectors.toList())));

    BatchWriteItemRequest batchRequest =
        BatchWriteItemRequest.builder().requestItems(mapTableRequest).build();

    Sleeper sleeper = Sleeper.DEFAULT;
    BackOff backoff = retryBackoff.backoff();
    int attempt = 0;
    while (true) {
      attempt++;
      try {
        client.batchWriteItem(batchRequest);
        break;
      } catch (Exception ex) {
        // Fail right away if there is no retry configuration
        if (spec.getRetryConfiguration() == null
            || !spec.getRetryConfiguration().getRetryPredicate().test(ex)) {
          DYNAMO_DB_WRITE_FAILURES.inc();
          LOG.info(
              "Unable to write batch items {} due to {} ",
              batchRequest.requestItems().entrySet(),
              ex);
          throw new IOException("Error writing to DynamoDB (no attempt made to retry)", ex);
        }

        if (!BackOffUtils.next(sleeper, backoff)) {
          throw new IOException(
              String.format(
                  "Error writing to DynamoDB after %d attempt(s). No more attempts allowed",
                  attempt),
              ex);
        } else {
          // Note: this used in test cases to verify behavior
          LOG.warn(String.format(RETRY_ATTEMPT_LOG, attempt), ex);
        }
      }
    }
  } finally {
    batch.clear();
  }
}
 

@Test
public void testRetries() throws Throwable {
  thrown.expectMessage("Error writing to DynamoDB");

  List<KV<String, Integer>> items =
      ImmutableList.of(KV.of("test1", 111), KV.of("test2", 222), KV.of("test3", 333));

  DynamoDbClient amazonDynamoDBMock = Mockito.mock(DynamoDbClient.class);
  Mockito.when(amazonDynamoDBMock.batchWriteItem(Mockito.any(BatchWriteItemRequest.class)))
      .thenThrow(DynamoDbException.builder().message("Service unavailable").build());

  pipeline
      .apply(Create.of(items))
      .apply(
          DynamoDBIO.<KV<String, Integer>>write()
              .withWriteRequestMapperFn(
                  (SerializableFunction<KV<String, Integer>, KV<String, WriteRequest>>)
                      entry -> {
                        Map<String, AttributeValue> putRequest =
                            ImmutableMap.of(
                                "hashKey1", AttributeValue.builder().s(entry.getKey()).build(),
                                "rangeKey2",
                                    AttributeValue.builder()
                                        .n(entry.getValue().toString())
                                        .build());

                        WriteRequest writeRequest =
                            WriteRequest.builder()
                                .putRequest(PutRequest.builder().item(putRequest).build())
                                .build();
                        return KV.of(tableName, writeRequest);
                      })
              .withRetryConfiguration(
                  DynamoDBIO.RetryConfiguration.builder()
                      .setMaxAttempts(4)
                      .setMaxDuration(Duration.standardSeconds(10))
                      .setRetryPredicate(DEFAULT_RETRY_PREDICATE)
                      .build())
              .withDynamoDbClientProvider(DynamoDbClientProviderMock.of(amazonDynamoDBMock)));

  try {
    pipeline.run().waitUntilFinish();
  } catch (final Pipeline.PipelineExecutionException e) {
    // check 3 retries were initiated by inspecting the log before passing on the exception
    expectedLogs.verifyWarn(String.format(DynamoDBIO.Write.WriteFn.RETRY_ATTEMPT_LOG, 1));
    expectedLogs.verifyWarn(String.format(DynamoDBIO.Write.WriteFn.RETRY_ATTEMPT_LOG, 2));
    expectedLogs.verifyWarn(String.format(DynamoDBIO.Write.WriteFn.RETRY_ATTEMPT_LOG, 3));
    throw e.getCause();
  }
  fail("Pipeline is expected to fail because we were unable to write to DynamoDb.");
}