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

The following examples show how to use software.amazon.awssdk.services.dynamodb.model.WriteRequest. You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example #1
Source File: BatchWriteResult.java    From aws-sdk-java-v2 with Apache License 2.0 6 votes vote down vote up
/**
 * 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());
}
 
Example #2
Source File: DynamoServiceIntegrationTest.java    From aws-sdk-java-v2 with Apache License 2.0 6 votes vote down vote up
@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());
    }
}
 
Example #3
Source File: BatchWriteItemOperation.java    From aws-sdk-java-v2 with Apache License 2.0 6 votes vote down vote up
@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();
}
 
Example #4
Source File: AWSDynamoDAO.java    From para with Apache License 2.0 5 votes vote down vote up
@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());
}
 
Example #5
Source File: AWSDynamoDAO.java    From para with Apache License 2.0 5 votes vote down vote up
@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());
}
 
Example #6
Source File: AWSDynamoUtils.java    From para with Apache License 2.0 5 votes vote down vote up
/**
 * 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());
}
 
Example #7
Source File: AWSDynamoUtils.java    From para with Apache License 2.0 5 votes vote down vote up
/**
 * 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);
	}
}
 
Example #8
Source File: AbstractDynamoRepository.java    From edison-microservice with Apache License 2.0 5 votes vote down vote up
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()));

}
 
Example #9
Source File: DynamoDBIOTestHelper.java    From beam with Apache License 2.0 5 votes vote down vote up
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;
}
 
Example #10
Source File: DynamoDBIO.java    From beam with Apache License 2.0 5 votes vote down vote up
@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();
  }
}
 
Example #11
Source File: BatchWriteItemOperationTest.java    From aws-sdk-java-v2 with Apache License 2.0 5 votes vote down vote up
@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)));
}
 
Example #12
Source File: BatchWriteItemOperationTest.java    From aws-sdk-java-v2 with Apache License 2.0 5 votes vote down vote up
@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))));
}
 
Example #13
Source File: BatchWriteItemOperationTest.java    From aws-sdk-java-v2 with Apache License 2.0 5 votes vote down vote up
@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);
}
 
Example #14
Source File: BatchWriteResult.java    From aws-sdk-java-v2 with Apache License 2.0 5 votes vote down vote up
/**
 * 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());
}
 
Example #15
Source File: DeleteItemOperation.java    From aws-sdk-java-v2 with Apache License 2.0 5 votes vote down vote up
@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();
}
 
Example #16
Source File: PutItemOperation.java    From aws-sdk-java-v2 with Apache License 2.0 5 votes vote down vote up
@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();
}
 
Example #17
Source File: BatchWriteItemOperationTest.java    From aws-sdk-java-v2 with Apache License 2.0 5 votes vote down vote up
@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))));
}
 
Example #18
Source File: BatchWriteItemOperationTest.java    From aws-sdk-java-v2 with Apache License 2.0 4 votes vote down vote up
@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)));
}
 
Example #19
Source File: BatchWriteItemOperationTest.java    From aws-sdk-java-v2 with Apache License 2.0 4 votes vote down vote up
private static WriteRequest putRequest(Map<String, AttributeValue> itemMap) {
    return WriteRequest.builder().putRequest(PutRequest.builder().item(itemMap).build()).build();
}
 
Example #20
Source File: BatchWriteItemOperationTest.java    From aws-sdk-java-v2 with Apache License 2.0 4 votes vote down vote up
private static WriteRequest deleteRequest(Map<String, AttributeValue> itemMap) {
    return WriteRequest.builder().deleteRequest(DeleteRequest.builder().key(itemMap).build()).build();
}
 
Example #21
Source File: DynamoDBIO.java    From beam with Apache License 2.0 4 votes vote down vote up
@Nullable
abstract SerializableFunction<T, KV<String, WriteRequest>> getWriteItemMapperFn();
 
Example #22
Source File: DynamoDBIO.java    From beam with Apache License 2.0 4 votes vote down vote up
abstract Builder<T> setWriteItemMapperFn(
SerializableFunction<T, KV<String, WriteRequest>> writeItemMapperFn);
 
Example #23
Source File: DynamoDBIO.java    From beam with Apache License 2.0 4 votes vote down vote up
public Write<T> withWriteRequestMapperFn(
    SerializableFunction<T, KV<String, WriteRequest>> writeItemMapperFn) {
  return toBuilder().setWriteItemMapperFn(writeItemMapperFn).build();
}
 
Example #24
Source File: WriteBatchTest.java    From aws-sdk-java-v2 with Apache License 2.0 4 votes vote down vote up
private static WriteRequest putRequest(Map<String, AttributeValue> itemMap) {
    return WriteRequest.builder().putRequest(PutRequest.builder().item(itemMap).build()).build();
}
 
Example #25
Source File: DynamoDBIO.java    From beam with Apache License 2.0 4 votes vote down vote up
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();
  }
}
 
Example #26
Source File: WriteBatchTest.java    From aws-sdk-java-v2 with Apache License 2.0 4 votes vote down vote up
private static WriteRequest deleteRequest(Map<String, AttributeValue> itemMap) {
    return WriteRequest.builder().deleteRequest(DeleteRequest.builder().key(itemMap).build()).build();
}
 
Example #27
Source File: DynamoDBIOTest.java    From beam with Apache License 2.0 4 votes vote down vote up
@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);
}
 
Example #28
Source File: DynamoDBIOTest.java    From beam with Apache License 2.0 4 votes vote down vote up
@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.");
}
 
Example #29
Source File: BatchWriteItemOperationTest.java    From aws-sdk-java-v2 with Apache License 2.0 4 votes vote down vote up
@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))));
}
 
Example #30
Source File: BatchableWriteOperation.java    From aws-sdk-java-v2 with Apache License 2.0 4 votes vote down vote up
WriteRequest generateWriteRequest(TableSchema<T> tableSchema,
OperationContext context,
DynamoDbEnhancedClientExtension extension);