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

/**
 * Retrieve any unprocessed put action items belonging to the supplied table from the result .
 * Call this method once for each table present in the batch request.
 *
 * @param mappedTable the table to retrieve unprocessed items for
 * @param <T> the type of the table items
 * @return a list of items
 */
public <T> List<T> unprocessedPutItemsForTable(MappedTableResource<T> mappedTable) {
    List<WriteRequest> writeRequests =
        unprocessedRequests.getOrDefault(mappedTable.tableName(),
                                         Collections.emptyList());

    return writeRequests.stream()
                        .filter(writeRequest -> writeRequest.putRequest() != null)
                        .map(WriteRequest::putRequest)
                        .map(PutRequest::item)
                        .map(item -> readAndTransformSingleItem(item,
                                                                mappedTable.tableSchema(),
                                                                DefaultOperationContext.create(mappedTable.tableName()),
                                                                mappedTable.mapperExtension()))
                        .collect(Collectors.toList());
}
 

@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();
}
 

@Override
public <P extends ParaObject> void deleteAll(String appid, List<P> objects) {
	if (objects == null || objects.isEmpty() || StringUtils.isBlank(appid)) {
		return;
	}

	Map<String, WriteRequest> reqs = new LinkedHashMap<>(objects.size());
	int batchSteps = 1;
	if ((objects.size() > MAX_ITEMS_PER_WRITE)) {
		batchSteps = (objects.size() / MAX_ITEMS_PER_WRITE)
				+ ((objects.size() % MAX_ITEMS_PER_WRITE > 0) ? 1 : 0);
	}

	Iterator<P> it = objects.iterator();
	String tableName = getTableNameForAppid(appid);
	int j = 0;

	for (int i = 0; i < batchSteps; i++) {
		while (it.hasNext() && j < MAX_ITEMS_PER_WRITE) {
			ParaObject object = it.next();
			if (object != null) {
				reqs.put(object.getId(), WriteRequest.builder().
						deleteRequest(b -> b.key(rowKey(object.getId(), appid))).build());
				j++;
			}
		}
		batchWrite(Collections.singletonMap(tableName, reqs.values().stream().collect(Collectors.toList())), 1);
		reqs.clear();
		j = 0;
	}
	logger.debug("DAO.deleteAll() {}", objects.size());
}
 

@Override
public <P extends ParaObject> void createAll(String appid, List<P> objects) {
	if (objects == null || objects.isEmpty() || StringUtils.isBlank(appid)) {
		return;
	}

	Map<String, WriteRequest> reqs = new LinkedHashMap<>(objects.size());
	int batchSteps = 1;
	if ((objects.size() > MAX_ITEMS_PER_WRITE)) {
		batchSteps = (objects.size() / MAX_ITEMS_PER_WRITE) +
				((objects.size() % MAX_ITEMS_PER_WRITE > 0) ? 1 : 0);
	}

	Iterator<P> it = objects.iterator();
	String tableName = getTableNameForAppid(appid);
	int j = 0;

	for (int i = 0; i < batchSteps; i++) {
		while (it.hasNext() && j < MAX_ITEMS_PER_WRITE) {
			ParaObject object = it.next();
			if (StringUtils.isBlank(object.getId())) {
				object.setId(Utils.getNewId());
			}
			if (object.getTimestamp() == null) {
				object.setTimestamp(Utils.timestamp());
			}
			//if (updateOp) object.setUpdated(Utils.timestamp());
			object.setAppid(appid);
			Map<String, AttributeValue> row = toRow(object, null);
			setRowKey(getKeyForAppid(object.getId(), appid), row);
			reqs.put(object.getId(), WriteRequest.builder().putRequest(b -> b.item(row)).build());
			j++;
		}
		batchWrite(Collections.singletonMap(tableName, reqs.values().stream().collect(Collectors.toList())), 1);
		reqs.clear();
		j = 0;
	}
	logger.debug("DAO.createAll() {}->{}", appid, objects.size());
}
 

/**
 * Deletes all objects in a shared table, which belong to a given appid, by scanning the GSI.
 * @param appid app id
 */
public static void deleteAllFromSharedTable(String appid) {
	if (StringUtils.isBlank(appid) || !isSharedAppid(appid)) {
		return;
	}
	Pager pager = new Pager(25);
	QueryResponse pages;
	Map<String, AttributeValue> lastKey = null;
	do {
		// read all phase
		pages = queryGSI(appid, pager);
		if (pages == null) {
			break;
		}
		List<WriteRequest> deletePage = new LinkedList<>();
		for (Map<String, AttributeValue> item : pages.items()) {
			String key = item.get(Config._KEY).s();
			// only delete rows which belong to the given appid
			if (StringUtils.startsWith(key, keyPrefix(appid))) {
				logger.debug("Preparing to delete '{}' from shared table, appid: '{}'.", key, appid);
				pager.setLastKey(item.get(Config._ID).s());
				deletePage.add(WriteRequest.builder().deleteRequest(b -> b.
						key(Collections.singletonMap(Config._KEY, AttributeValue.builder().s(key).
								build()))).build());
			}
		}
		lastKey = pages.lastEvaluatedKey();
		// delete all phase
		logger.info("Deleting {} items belonging to app '{}', from shared table...", deletePage.size(), appid);
		if (!deletePage.isEmpty()) {
			batchWrite(Collections.singletonMap(getTableNameForAppid(appid), deletePage), 1);
		}
	} while (lastKey != null && !lastKey.isEmpty());
}
 

/**
 * Writes multiple items in batch.
 * @param items a map of tables->write requests
 * @param backoff backoff seconds
 */
protected static void batchWrite(Map<String, List<WriteRequest>> items, int backoff) {
	if (items == null || items.isEmpty()) {
		return;
	}
	try {
		logger.debug("batchWrite(): requests {}, backoff {}", items.values().iterator().next().size(), backoff);
		BatchWriteItemResponse result = getClient().batchWriteItem(b -> b.
				returnConsumedCapacity(ReturnConsumedCapacity.TOTAL).requestItems(items));
		if (result == null) {
			return;
		}
		logger.debug("batchWrite(): success - consumed capacity {}", result.consumedCapacity());

		if (result.unprocessedItems() != null && !result.unprocessedItems().isEmpty()) {
			Thread.sleep((long) backoff * 1000L);
			logger.warn("{} UNPROCESSED write requests!", result.unprocessedItems().size());
			batchWrite(result.unprocessedItems(), backoff * 2);
		}
	} catch (ProvisionedThroughputExceededException ex) {
		logger.warn("Write capacity exceeded for table '{}'. Retrying request in {} seconds.",
				items.keySet().iterator().next(), backoff);
		try {
			Thread.sleep((long) backoff * 1000L);
			// retry forever
			batchWrite(items, backoff * 2);
		} catch (InterruptedException ie) {
			logger.error(null, ie);
			Thread.currentThread().interrupt();
		}
	} catch (Exception e) {
		logger.error("Failed to execute batch write operation on table '{}'", items.keySet().iterator().next(), e);
		throwIfNecessary(e);
	}
}
 

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()));

}
 

static List<WriteRequest> generateWriteRequests(int numOfItem) {
  List<WriteRequest> writeRequests = new ArrayList<>();
  for (int i = 1; i <= numOfItem; i++) {
    WriteRequest writeRequest =
        WriteRequest.builder()
            .putRequest(generatePutRequest("hashKeyDataStr_" + i, "1000" + i))
            .build();
    writeRequests.add(writeRequest);
  }
  return writeRequests;
}
 

@ProcessElement
public void processElement(ProcessContext context) throws Exception {
  final KV<String, WriteRequest> writeRequest =
      (KV<String, WriteRequest>) spec.getWriteItemMapperFn().apply(context.element());
  batch.add(writeRequest);
  if (batch.size() >= BATCH_SIZE) {
    flushBatch();
  }
}
 

@Test
public void transformResults_multipleUnprocessedOperations() {
    BatchWriteItemOperation operation = BatchWriteItemOperation.create(emptyRequest());

    List<WriteRequest> writeRequests1 = Arrays.asList(putRequest(FAKE_ITEM_MAPS.get(0)),
                                                      deleteRequest(FAKE_ITEM_MAPS.get(1)),
                                                      deleteRequest(FAKE_ITEM_MAPS.get(2)));
    List<WriteRequest> writeRequests2 = Arrays.asList(deleteRequest(FAKESORT_ITEM_MAPS.get(0)),
                                                      putRequest(FAKESORT_ITEM_MAPS.get(1)),
                                                      putRequest(FAKESORT_ITEM_MAPS.get(2)));
    Map<String, List<WriteRequest>> writeRequests = new HashMap<>();
    writeRequests.put(TABLE_NAME, writeRequests1);
    writeRequests.put(TABLE_NAME_2, writeRequests2);
    BatchWriteItemResponse response =
        BatchWriteItemResponse.builder()
                              .unprocessedItems(writeRequests)
                              .build();

    BatchWriteResult results = operation.transformResponse(response, mockExtension);

    assertThat(results.unprocessedDeleteItemsForTable(fakeItemMappedTableWithExtension),
               containsInAnyOrder(FAKE_ITEM_KEYS.get(1), FAKE_ITEM_KEYS.get(2)));
    assertThat(results.unprocessedPutItemsForTable(fakeItemMappedTableWithExtension),
               containsInAnyOrder(FAKE_ITEMS.get(0)));
    assertThat(results.unprocessedDeleteItemsForTable(fakeItemWithSortMappedTableWithExtension),
               containsInAnyOrder(FAKESORT_ITEM_KEYS.get(0)));
    assertThat(results.unprocessedPutItemsForTable(fakeItemWithSortMappedTableWithExtension),
               containsInAnyOrder(FAKESORT_ITEMS.get(1), FAKESORT_ITEMS.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))));
}
 

@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);
}
 

/**
 * Retrieve any unprocessed delete action keys belonging to the supplied table from the result.
 * Call this method once for each table present in the batch request.
 *
 * @param mappedTable the table to retrieve unprocessed items for.
 * @return a list of keys that were not processed as part of the batch request.
 */
public List<Key> unprocessedDeleteItemsForTable(MappedTableResource<?> mappedTable) {
    List<WriteRequest> writeRequests =
        unprocessedRequests.getOrDefault(mappedTable.tableName(),
                                         Collections.emptyList());

    return writeRequests.stream()
                        .filter(writeRequest -> writeRequest.deleteRequest() != null)
                        .map(WriteRequest::deleteRequest)
                        .map(DeleteRequest::key)
                        .map(itemMap -> createKeyFromMap(itemMap,
                                                         mappedTable.tableSchema(),
                                                         TableMetadata.primaryIndexName()))
                        .collect(Collectors.toList());
}
 

@Override
public WriteRequest generateWriteRequest(TableSchema<T> tableSchema,
                                         OperationContext operationContext,
                                         DynamoDbEnhancedClientExtension extension) {
    DeleteItemRequest deleteItemRequest = generateRequest(tableSchema, operationContext, extension);

    return WriteRequest.builder()
                       .deleteRequest(DeleteRequest.builder().key(deleteItemRequest.key()).build())
                       .build();
}
 

@Override
public WriteRequest generateWriteRequest(TableSchema<T> tableSchema,
                                         OperationContext operationContext,
                                         DynamoDbEnhancedClientExtension extension) {

    PutItemRequest putItemRequest = generateRequest(tableSchema, operationContext, extension);

    if (putItemRequest.conditionExpression() != null) {
        throw new IllegalArgumentException("A mapper extension inserted a conditionExpression in a PutItem "
                                           + "request as part of a BatchWriteItemRequest. This is not supported by "
                                           + "DynamoDb. An extension known to do this is the "
                                           + "VersionedRecordExtension which is loaded by default unless overridden. "
                                           + "To fix this use a table schema that does not "
                                           + "have a versioned attribute in it or do not load the offending extension.");
    }

    return WriteRequest.builder().putRequest(PutRequest.builder().item(putItemRequest.item()).build()).build();
}
 

@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 transformResults_multipleUnprocessedOperations_extensionTransformsPutsNotDeletes() {
    BatchWriteItemOperation operation = BatchWriteItemOperation.create(emptyRequest());

    List<WriteRequest> writeRequests1 = Arrays.asList(putRequest(FAKE_ITEM_MAPS.get(0)),
                                                      deleteRequest(FAKE_ITEM_MAPS.get(1)),
                                                      deleteRequest(FAKE_ITEM_MAPS.get(2)));
    List<WriteRequest> writeRequests2 = Arrays.asList(deleteRequest(FAKESORT_ITEM_MAPS.get(0)),
                                                      putRequest(FAKESORT_ITEM_MAPS.get(1)),
                                                      putRequest(FAKESORT_ITEM_MAPS.get(2)));
    Map<String, List<WriteRequest>> writeRequests = new HashMap<>();
    writeRequests.put(TABLE_NAME, writeRequests1);
    writeRequests.put(TABLE_NAME_2, writeRequests2);
    BatchWriteItemResponse response =
        BatchWriteItemResponse.builder()
                              .unprocessedItems(writeRequests)
                              .build();

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

    BatchWriteResult results = operation.transformResponse(response, mockExtension);

    assertThat(results.unprocessedDeleteItemsForTable(fakeItemMappedTableWithExtension),
               containsInAnyOrder(FAKE_ITEM_KEYS.get(1), FAKE_ITEM_KEYS.get(2)));
    assertThat(results.unprocessedPutItemsForTable(fakeItemMappedTableWithExtension),
               containsInAnyOrder(FAKE_ITEMS.get(3)));
    assertThat(results.unprocessedDeleteItemsForTable(fakeItemWithSortMappedTableWithExtension),
               containsInAnyOrder(FAKESORT_ITEM_KEYS.get(0)));
    assertThat(results.unprocessedPutItemsForTable(fakeItemWithSortMappedTableWithExtension),
               containsInAnyOrder(FAKESORT_ITEMS.get(4), FAKESORT_ITEMS.get(5)));
}
 

private static WriteRequest putRequest(Map<String, AttributeValue> itemMap) {
    return WriteRequest.builder().putRequest(PutRequest.builder().item(itemMap).build()).build();
}
 

private static WriteRequest deleteRequest(Map<String, AttributeValue> itemMap) {
    return WriteRequest.builder().deleteRequest(DeleteRequest.builder().key(itemMap).build()).build();
}
 

@Nullable
abstract SerializableFunction<T, KV<String, WriteRequest>> getWriteItemMapperFn();
 

abstract Builder<T> setWriteItemMapperFn(
SerializableFunction<T, KV<String, WriteRequest>> writeItemMapperFn);
 

public Write<T> withWriteRequestMapperFn(
    SerializableFunction<T, KV<String, WriteRequest>> writeItemMapperFn) {
  return toBuilder().setWriteItemMapperFn(writeItemMapperFn).build();
}
 

private static WriteRequest putRequest(Map<String, AttributeValue> itemMap) {
    return WriteRequest.builder().putRequest(PutRequest.builder().item(itemMap).build()).build();
}
 

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();
  }
}
 

private static WriteRequest deleteRequest(Map<String, AttributeValue> itemMap) {
    return WriteRequest.builder().deleteRequest(DeleteRequest.builder().key(itemMap).build()).build();
}
 

@Test
public void testWriteDataToDynamo() {
  List<KV<String, Integer>> items =
      ImmutableList.of(KV.of("test1", 111), KV.of("test2", 222), KV.of("test3", 333));

  final PCollection<Void> output =
      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(5)
                          .setMaxDuration(Duration.standardMinutes(1))
                          .setRetryPredicate(DEFAULT_RETRY_PREDICATE)
                          .build())
                  .withDynamoDbClientProvider(
                      DynamoDbClientProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient())));

  final PCollection<Long> publishedResultsSize = output.apply(Count.globally());
  PAssert.that(publishedResultsSize).containsInAnyOrder(0L);

  pipeline.run().waitUntilFinish();

  // Make sure data written to the table are in the table.
  int actualItemCount = DynamoDBIOTestHelper.readDataFromTable(tableName).size();
  assertEquals(3, actualItemCount);
}
 

@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.");
}
 

@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))));
}
 

WriteRequest generateWriteRequest(TableSchema<T> tableSchema,
OperationContext context,
DynamoDbEnhancedClientExtension extension);