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

@Test
public void testReaderOneSegment() {
  List<Map<String, AttributeValue>> expected =
      DynamoDBIOTestHelper.generateTestData(tableName, numOfItems);

  PCollection<List<Map<String, AttributeValue>>> actual =
      pipeline.apply(
          DynamoDBIO.<List<Map<String, AttributeValue>>>read()
              .withDynamoDbClientProvider(
                  DynamoDbClientProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient()))
              .withScanRequestFn(
                  (SerializableFunction<Void, ScanRequest>)
                      input ->
                          ScanRequest.builder().tableName(tableName).totalSegments(1).build())
              .items());
  PAssert.that(actual).containsInAnyOrder(expected);
  pipeline.run().waitUntilFinish();
}
 

@Test
public void generateRequest_hashAndRangeKey_exclusiveStartKey() {
    FakeItemWithSort exclusiveStartKey = createUniqueFakeItemWithSort();
    Map<String, AttributeValue> keyMap =
        FakeItemWithSort.getTableSchema().itemToMap(exclusiveStartKey,
                                                    FakeItemWithSort.getTableMetadata().primaryKeys());

    ScanOperation<FakeItemWithSort> operation =
        ScanOperation.create(ScanEnhancedRequest.builder().exclusiveStartKey(keyMap).build());

    ScanRequest scanRequest = operation.generateRequest(FakeItemWithSort.getTableSchema(),
                                                         PRIMARY_CONTEXT,
                                                         null);

    assertThat(scanRequest.exclusiveStartKey(),
               hasEntry("id", AttributeValue.builder().s(exclusiveStartKey.getId()).build()));
    assertThat(scanRequest.exclusiveStartKey(),
               hasEntry("sort", AttributeValue.builder().s(exclusiveStartKey.getSort()).build()));
}
 

@Test
public void generateRequest_projectionExpression() {
    ScanOperation<FakeItem> operation = ScanOperation.create(
            ScanEnhancedRequest.builder()
                    .attributesToProject("id")
                    .addAttributeToProject("version")
                    .build()
    );
    ScanRequest request = operation.generateRequest(FakeItem.getTableSchema(),
            PRIMARY_CONTEXT,
            null);

    Map<String, String> expectedExpressionAttributeNames = new HashMap<>();
    expectedExpressionAttributeNames.put("#AMZN_MAPPED_id", "id");
    expectedExpressionAttributeNames.put("#AMZN_MAPPED_version", "version");

    ScanRequest expectedRequest = ScanRequest.builder()
            .tableName(TABLE_NAME)
            .projectionExpression("#AMZN_MAPPED_id,#AMZN_MAPPED_version")
            .expressionAttributeNames(expectedExpressionAttributeNames)
            .build();
    assertThat(request, is(expectedRequest));
}
 

public static void scanItems( DynamoDbClient ddb,String tableName ) {

        try {
            ScanRequest scanRequest = ScanRequest.builder()
                    .tableName(tableName)
                    .build();

            ScanResponse response = ddb.scan(scanRequest);
            for (Map<String, AttributeValue> item : response.items()) {
                Set<String> keys = item.keySet();
                for (String key : keys) {

                    System.out.println ("The key name is "+key +"\n" );
                    System.out.println("The value is "+item.get(key).s());
                }
            }

        } catch (DynamoDbException e) {
            e.printStackTrace();
        }
        // snippet-end:[dynamodb.java2.dynamoDB_scan.main]
    }
 

@Override
public PCollection<T> expand(PBegin input) {
  checkArgument((getScanRequestFn() != null), "withScanRequestFn() is required");
  checkArgument(
      (getDynamoDbClientProvider() != null), "withDynamoDbClientProvider() is required");
  ScanRequest scanRequest = getScanRequestFn().apply(null);
  checkArgument(
      (scanRequest.totalSegments() != null && scanRequest.totalSegments() > 0),
      "TotalSegments is required with withScanRequestFn() and greater zero");

  PCollection<Read<T>> splits =
      (PCollection<Read<T>>)
          input.apply("Create", Create.of(this)).apply("Split", ParDo.of(new SplitFn()));
  splits.setCoder(SerializableCoder.of(new TypeDescriptor<Read<T>>() {}));

  PCollection<T> output =
      (PCollection<T>)
          splits
              .apply("Reshuffle", Reshuffle.viaRandomKey())
              .apply("Read", ParDo.of(new ReadFn()));
  output.setCoder(getCoder());
  return output;
}
 

@Test
public void generateRequest_filterCondition_expressionAndValues() {
    Map<String, AttributeValue> expressionValues = singletonMap(":test-key", stringValue("test-value"));
    Expression filterExpression =
        Expression.builder().expression("test-expression").expressionValues(expressionValues).build();
    ScanOperation<FakeItem> operation =
        ScanOperation.create(ScanEnhancedRequest.builder().filterExpression(filterExpression).build());

    ScanRequest request = operation.generateRequest(FakeItem.getTableSchema(),
                                                          PRIMARY_CONTEXT,
                                                          null);
    ScanRequest expectedRequest = ScanRequest.builder()
                                             .tableName(TABLE_NAME)
                                             .filterExpression("test-expression")
                                             .expressionAttributeValues(expressionValues)
                                             .build();
    assertThat(request, is(expectedRequest));
}
 

@Test
public void testMissingTotalSegments() {
  thrown.expectMessage("TotalSegments is required with withScanRequestFn()");
  pipeline.apply(
      DynamoDBIO.read()
          .withScanRequestFn(
              (SerializableFunction<Void, ScanRequest>)
                  input -> ScanRequest.builder().tableName(tableName).build())
          .withDynamoDbClientProvider(
              DynamoDbClientProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient())));
  try {
    pipeline.run().waitUntilFinish();
    fail("TotalSegments is required with withScanRequestFn()");
  } catch (IllegalArgumentException ex) {
    assertEquals("TotalSegments is required with withScanRequestFn()", ex.getMessage());
  }
}
 

@Test
public void testNegativeTotalSegments() {
  thrown.expectMessage("TotalSegments is required with withScanRequestFn() and greater zero");
  pipeline.apply(
      DynamoDBIO.read()
          .withScanRequestFn(
              (SerializableFunction<Void, ScanRequest>)
                  input -> ScanRequest.builder().tableName(tableName).totalSegments(-1).build())
          .withDynamoDbClientProvider(
              DynamoDbClientProviderMock.of(DynamoDBIOTestHelper.getDynamoDBClient())));
  try {
    pipeline.run().waitUntilFinish();
    fail("withTotalSegments() is expected and greater than zero");
  } catch (IllegalArgumentException ex) {
    assertEquals(
        "TotalSegments is required with withScanRequestFn() and greater zero", ex.getMessage());
  }
}
 

@Test
public void testMissingDynamoDbClientProvider() {
  thrown.expectMessage("withDynamoDbClientProvider() is required");
  pipeline.apply(
      DynamoDBIO.read()
          .withScanRequestFn(
              (SerializableFunction<Void, ScanRequest>)
                  input -> ScanRequest.builder().tableName(tableName).totalSegments(3).build()));
  try {
    pipeline.run().waitUntilFinish();
    fail("withDynamoDbClientProvider() is required");
  } catch (IllegalArgumentException ex) {
    assertEquals("withDynamoDbClientProvider() is required", ex.getMessage());
  }
}
 

@Test
public void test_MultipleIteration_On_Responses_Stream() {
    int results_per_page = 2;
    ScanRequest request = ScanRequest.builder().tableName(TABLE_NAME).limit(results_per_page).build();
    ScanIterable scanResponses = dynamo.scanPaginator(request);

    // Iterate once
    assertEquals(ITEM_COUNT, scanResponses.stream()
                                          .mapToInt(response -> response.count())
                                          .sum());

    // Iterate second time
    assertEquals(ITEM_COUNT, scanResponses.stream()
                                          .mapToInt(response -> response.count())
                                          .sum());

    // Number of pages
    assertEquals(Math.ceil((double) ITEM_COUNT / results_per_page), scanResponses.stream().count(), 0);
}
 

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

@ProcessElement
public void processElement(@Element Read<T> spec, OutputReceiver<T> out) {
  DynamoDbClient client = spec.getDynamoDbClientProvider().getDynamoDbClient();

  ScanRequest scanRequest = spec.getScanRequestFn().apply(null);
  ScanRequest scanRequestWithSegment =
      scanRequest.toBuilder().segment(spec.getSegmentId()).build();

  ScanResponse scanResponse = client.scan(scanRequestWithSegment);
  out.output(spec.getScanResponseMapperFn().apply(scanResponse));
}
 

@ProcessElement
public void processElement(@Element Read<T> spec, OutputReceiver<Read<T>> out) {
  ScanRequest scanRequest = spec.getScanRequestFn().apply(null);
  for (int i = 0; i < scanRequest.totalSegments(); i++) {
    out.output(spec.withSegmentId(i));
  }
}
 

@AfterEach
public void truncateLockTable() {
    List<Map<String, AttributeValue>> items = dynamodb.scan(ScanRequest.builder().tableName(TABLE_NAME).build()).items();
    for (Map<String, AttributeValue> item : items) {
        dynamodb.deleteItem(DeleteItemRequest.builder()
            .tableName(TABLE_NAME)
            .key(Collections.singletonMap(ID, item.get(ID)))
            .build());
    }
}
 

/**
 * Reads a page from a standard DynamoDB table.
 * @param <P> type of object
 * @param appid the app identifier (name)
 * @param p a {@link Pager}
 * @return the last row key of the page, or null.
 */
public static <P extends ParaObject> List<P> readPageFromTable(String appid, Pager p) {
	Pager pager = (p != null) ? p : new Pager();
	ScanRequest.Builder scanRequest = ScanRequest.builder().
			tableName(getTableNameForAppid(appid)).
			limit(pager.getLimit()).
			returnConsumedCapacity(ReturnConsumedCapacity.TOTAL);

	if (!StringUtils.isBlank(pager.getLastKey())) {
		scanRequest.exclusiveStartKey(Collections.
				singletonMap(Config._KEY, AttributeValue.builder().s(pager.getLastKey()).build()));
	}

	ScanResponse result = getClient().scan(scanRequest.build());
	String lastKey = null;
	LinkedList<P> results = new LinkedList<>();
	for (Map<String, AttributeValue> item : result.items()) {
		P obj = fromRow(item);
		if (obj != null) {
			lastKey = item.get(Config._KEY).s();
			results.add(obj);
		}
	}

	if (result.lastEvaluatedKey() != null && !result.lastEvaluatedKey().isEmpty()) {
		pager.setLastKey(result.lastEvaluatedKey().get(Config._KEY).s());
	} else if (!results.isEmpty()) {
		// set last key to be equal to the last result - end reached.
		pager.setLastKey(lastKey);
	}
	return results;
}
 

@Test
public void testListLeasesHandlesTimeout() throws Exception {
    TimeoutException te = setRuleForDependencyTimeout();

    when(mockScanFuture.get(anyLong(), any(TimeUnit.class))).thenThrow(te);
    when(dynamoDbClient.scan(any(ScanRequest.class))).thenReturn(mockScanFuture);

    verifyCancel(mockScanFuture, () -> leaseRefresher.listLeases());
}
 

@Test
public void testDeleteAllLeasesTimesOut() throws Exception {
    TimeoutException te = setRuleForDependencyTimeout();
    when(dynamoDbClient.scan(any(ScanRequest.class))).thenReturn(mockScanFuture);
    when(mockScanFuture.get(anyLong(), any())).thenReturn(ScanResponse.builder().items(Collections.emptyMap()).build());
    when(leaseSerializer.fromDynamoRecord(any())).thenReturn(lease);
    when(leaseSerializer.getDynamoHashKey(any(Lease.class))).thenReturn(Collections.emptyMap());

    when(dynamoDbClient.deleteItem(any(DeleteItemRequest.class))).thenReturn(mockDeleteFuture);
    when(mockDeleteFuture.get(anyLong(), any())).thenThrow(te);

    verifyCancel(mockDeleteFuture, () -> leaseRefresher.deleteAll());
}
 

@Test
public void generateRequest_hashKeyOnly_exclusiveStartKey() {
    FakeItem exclusiveStartKey = createUniqueFakeItem();
    Map<String, AttributeValue> keyMap = FakeItem.getTableSchema().itemToMap(exclusiveStartKey, singletonList("id"));
    ScanOperation<FakeItem> operation =
        ScanOperation.create(ScanEnhancedRequest.builder().exclusiveStartKey(keyMap).build());

    ScanRequest scanRequest = operation.generateRequest(FakeItem.getTableSchema(),
                                                         PRIMARY_CONTEXT,
                                                         null);

    assertThat(scanRequest.exclusiveStartKey(),
               hasEntry("id", AttributeValue.builder().s(exclusiveStartKey.getId()).build()));
}
 

@Test
public void generateRequest_consistentRead() {
    ScanOperation<FakeItem> operation = ScanOperation.create(ScanEnhancedRequest.builder().consistentRead(true).build());
    ScanRequest request = operation.generateRequest(FakeItem.getTableSchema(),
                                                          PRIMARY_CONTEXT,
                                                          null);

    ScanRequest expectedRequest = ScanRequest.builder()
                                             .tableName(TABLE_NAME)
                                             .consistentRead(true)
                                             .build();
    assertThat(request, is(expectedRequest));
}
 

@Test
public void test_MultipleIteration_On_PaginatedMember_Stream() {
    int results_per_page = 2;
    ScanRequest request = ScanRequest.builder().tableName(TABLE_NAME).limit(results_per_page).build();
    SdkIterable<Map<String, AttributeValue>> items = dynamo.scanPaginator(request).items();

    // Iterate once
    assertEquals(ITEM_COUNT, items.stream().distinct().count());

    // Iterate second time
    assertEquals(ITEM_COUNT, items.stream().distinct().count());
}
 

@Test(expected = IllegalStateException.class)
public void iteration_On_SameStream_ThrowsError() {
    int results_per_page = 2;
    ScanRequest request = ScanRequest.builder().tableName(TABLE_NAME).limit(results_per_page).build();
    Stream<Map<String, AttributeValue>> itemsStream = dynamo.scanPaginator(request).items().stream();

    // Iterate once
    assertEquals(ITEM_COUNT, itemsStream.distinct().count());

    // Iterate second time
    assertEquals(ITEM_COUNT, itemsStream.distinct().count());
}
 

@Test
public void generateRequest_filterCondition_expressionOnly() {
    Expression filterExpression = Expression.builder().expression("test-expression").build();
    ScanOperation<FakeItem> operation =
        ScanOperation.create(ScanEnhancedRequest.builder().filterExpression(filterExpression).build());

    ScanRequest request = operation.generateRequest(FakeItem.getTableSchema(),
                                                          PRIMARY_CONTEXT,
                                                          null);
    ScanRequest expectedRequest = ScanRequest.builder()
                                             .tableName(TABLE_NAME)
                                             .filterExpression("test-expression")
                                             .build();
    assertThat(request, is(expectedRequest));
}
 

@Test
public void generateRequest_limit() {
    ScanOperation<FakeItem> operation = ScanOperation.create(ScanEnhancedRequest.builder().limit(10).build());
    ScanRequest request = operation.generateRequest(FakeItem.getTableSchema(),
                                                          PRIMARY_CONTEXT,
                                                          null);

    ScanRequest expectedRequest = ScanRequest.builder()
        .tableName(TABLE_NAME)
        .limit(10)
        .build();
    assertThat(request, is(expectedRequest));
}
 

@Test
public void generateRequest_defaultScan() {
    ScanRequest request = scanOperation.generateRequest(FakeItem.getTableSchema(),
                                                        PRIMARY_CONTEXT,
                                                        null);

    ScanRequest expectedRequest = ScanRequest.builder()
        .tableName(TABLE_NAME)
        .build();
    assertThat(request, is(expectedRequest));
}
 

@Test
public void getAsyncServiceCall_makesTheRightCallAndReturnsResponse() {
    ScanRequest scanRequest = ScanRequest.builder().build();
    ScanPublisher mockScanPublisher = mock(ScanPublisher.class);
    when(mockDynamoDbAsyncClient.scanPaginator(any(ScanRequest.class))).thenReturn(mockScanPublisher);

    SdkPublisher<ScanResponse> response = scanOperation.asyncServiceCall(mockDynamoDbAsyncClient)
                                                       .apply(scanRequest);

    assertThat(response, is(mockScanPublisher));
    verify(mockDynamoDbAsyncClient).scanPaginator(scanRequest);
}
 

@Test
public void getServiceCall_makesTheRightCallAndReturnsResponse() {
    ScanRequest scanRequest = ScanRequest.builder().build();
    ScanIterable mockScanIterable = mock(ScanIterable.class);
    when(mockDynamoDbClient.scanPaginator(any(ScanRequest.class))).thenReturn(mockScanIterable);

    SdkIterable<ScanResponse> response = scanOperation.serviceCall(mockDynamoDbClient).apply(scanRequest);

    assertThat(response, is(mockScanIterable));
    verify(mockDynamoDbClient).scanPaginator(scanRequest);
}
 

protected ScanRequest scanRequest() {
    return ScanRequest.builder().tableName(getTableName())
            .attributesToGet(FRUIT_NAME_COL, FRUIT_DESC_COL).build();
}
 

@Test
public void testListLeasesSucceedsThenFails() throws Exception {
    TimeoutException te = setRuleForDependencyTimeout();

    when(dynamoDbClient.scan(any(ScanRequest.class))).thenReturn(mockScanFuture);

    Map<String, AttributeValue> lastEvaluatedKey = new HashMap<>();
    lastEvaluatedKey.put("Test", AttributeValue.builder().s("test").build());

    when(mockScanFuture.get(anyLong(), any(TimeUnit.class)))
            .thenReturn(ScanResponse.builder().lastEvaluatedKey(lastEvaluatedKey).build())
            .thenThrow(te);

    verifyCancel(mockScanFuture, () -> leaseRefresher.listLeases());

    verify(mockScanFuture, times(2)).get(anyLong(), any(TimeUnit.class));
    verify(dynamoDbClient, times(2)).scan(any(ScanRequest.class));

}
 

@Override
public ScanRequest generateRequest(TableSchema<T> tableSchema,
                                   OperationContext operationContext,
                                   DynamoDbEnhancedClientExtension extension) {
    Map<String, AttributeValue> expressionValues = null;
    Map<String, String> expressionNames = null;

    if (this.request.filterExpression() != null) {
        expressionValues = this.request.filterExpression().expressionValues();
        expressionNames = this.request.filterExpression().expressionNames();
    }

    String projectionExpression = null;
    if (this.request.attributesToProject() != null) {
        List<String> placeholders = new ArrayList<>();
        Map<String, String> projectionPlaceholders = new HashMap<>();
        this.request.attributesToProject().forEach(attr -> {
            String placeholder = PROJECTION_EXPRESSION_KEY_MAPPER.apply(attr);
            placeholders.add(placeholder);
            projectionPlaceholders.put(placeholder, attr);
        });
        projectionExpression = String.join(",", placeholders);
        expressionNames = Expression.joinNames(expressionNames, projectionPlaceholders);
    }

    ScanRequest.Builder scanRequest = ScanRequest.builder()
        .tableName(operationContext.tableName())
        .limit(this.request.limit())
        .exclusiveStartKey(this.request.exclusiveStartKey())
        .consistentRead(this.request.consistentRead())
        .expressionAttributeValues(expressionValues)
        .expressionAttributeNames(expressionNames)
        .projectionExpression(projectionExpression);

    if (!TableMetadata.primaryIndexName().equals(operationContext.indexName())) {
        scanRequest = scanRequest.indexName(operationContext.indexName());
    }

    if (this.request.filterExpression() != null) {
        scanRequest = scanRequest.filterExpression(this.request.filterExpression().expression());
    }

    return scanRequest.build();
}
 

@Override
public Function<ScanRequest, SdkIterable<ScanResponse>> serviceCall(DynamoDbClient dynamoDbClient) {
    return dynamoDbClient::scanPaginator;
}