software.amazon.awssdk.services.kinesis.model.GetRecordsRequest Java Examples

private GetRecordsResponse tryNext() {
    GetRecordsResponse response = kinesisClient.getRecords(GetRecordsRequest.builder()
            .shardIterator(id)
            .limit(fetchRecordLimit)
            .build())
            .join();
    this.id = response.nextShardIterator();
    LOG.debug("next() with id " + this.id + " returned " + response.records().size() + " records");
    if (!response.records().isEmpty()) {
        this.shardPosition = fromPosition(
                shardPosition.shardName(),
                response.records().get(response.records().size()-1).sequenceNumber()
        );
    }
    return response;
}
 

@Test
public void testGetRecordsThrowsSdkException() throws Exception {
    expectedExceptionRule.expect(SdkException.class);
    expectedExceptionRule.expectMessage("Test Exception");

    CompletableFuture<GetShardIteratorResponse> getShardIteratorFuture = CompletableFuture
            .completedFuture(GetShardIteratorResponse.builder().shardIterator("test").build());

    // Set up proxy mock methods
    when(kinesisClient.getShardIterator(any(GetShardIteratorRequest.class))).thenReturn(getShardIteratorFuture);
    when(kinesisClient.getRecords(any(GetRecordsRequest.class))).thenReturn(getRecordsResponseFuture);
    when(getRecordsResponseFuture.get(anyLong(), any(TimeUnit.class)))
            .thenThrow(new ExecutionException(SdkException.builder().message("Test Exception").build()));

    // Create data fectcher and initialize it with latest type checkpoint
    kinesisDataFetcher.initialize(SentinelCheckpoint.LATEST.toString(), INITIAL_POSITION_LATEST);
    final GetRecordsRetrievalStrategy getRecordsRetrievalStrategy = new SynchronousGetRecordsRetrievalStrategy(
            kinesisDataFetcher);

    // Call records of dataFetcher which will throw an exception
    getRecordsRetrievalStrategy.getRecords(MAX_RECORDS);

}
 

@Test
public void testNonNullGetRecords() throws Exception {
    final String nextIterator = "TestIterator";
    final ArgumentCaptor<GetShardIteratorRequest> iteratorCaptor = ArgumentCaptor
            .forClass(GetShardIteratorRequest.class);
    final ArgumentCaptor<GetRecordsRequest> recordsCaptor = ArgumentCaptor.forClass(GetRecordsRequest.class);
    final GetShardIteratorRequest expectedIteratorRequest = makeGetShardIteratorRequest(
            ShardIteratorType.LATEST.name());
    final GetRecordsRequest expectedRecordsRequest = makeGetRecordsRequest(nextIterator);

    final CompletableFuture<GetRecordsResponse> future = mock(CompletableFuture.class);

    when(kinesisClient.getShardIterator(iteratorCaptor.capture()))
            .thenReturn(makeGetShardIteratorResonse(nextIterator));
    when(kinesisClient.getRecords(recordsCaptor.capture())).thenReturn(future);
    when(future.get(anyLong(), any(TimeUnit.class))).thenThrow(
            new ExecutionException(ResourceNotFoundException.builder().message("Test Exception").build()));

    kinesisDataFetcher.initialize(SentinelCheckpoint.LATEST.toString(), INITIAL_POSITION_LATEST);
    DataFetcherResult dataFetcherResult = kinesisDataFetcher.getRecords();

    assertNotNull(dataFetcherResult);
    assertEquals(expectedIteratorRequest.startingSequenceNumber(), iteratorCaptor.getValue().startingSequenceNumber());
    assertEquals(expectedRecordsRequest.shardIterator(), recordsCaptor.getValue().shardIterator());
}
 

@Test
public void testTimeoutExceptionIsRetryableForGetRecords() throws Exception {
    expectedExceptionRule.expect(RetryableRetrievalException.class);
    expectedExceptionRule.expectCause(isA(TimeoutException.class));
    expectedExceptionRule.expectMessage("Timeout");

    CompletableFuture<GetShardIteratorResponse> getShardIteratorFuture = CompletableFuture
            .completedFuture(GetShardIteratorResponse.builder().shardIterator("test").build());

    // Set up proxy mock methods
    when(kinesisClient.getShardIterator(any(GetShardIteratorRequest.class))).thenReturn(getShardIteratorFuture);
    when(kinesisClient.getRecords(any(GetRecordsRequest.class))).thenReturn(getRecordsResponseFuture);
    when(getRecordsResponseFuture.get(anyLong(), any(TimeUnit.class))).thenThrow(new TimeoutException("Timeout"));

    // Create data fectcher and initialize it with latest type checkpoint
    kinesisDataFetcher.initialize(SentinelCheckpoint.LATEST.toString(), INITIAL_POSITION_LATEST);
    final GetRecordsRetrievalStrategy getRecordsRetrievalStrategy = new SynchronousGetRecordsRetrievalStrategy(
            kinesisDataFetcher);

    // Call records of dataFetcher which will throw an exception
    getRecordsRetrievalStrategy.getRecords(MAX_RECORDS);
}
 

@Test
public void testGetFromNullIterator() {
    try {
        client.getRecords(GetRecordsRequest.builder().build());
        Assert.fail("Expected InvalidArgumentException");
    } catch (SdkServiceException exception) {
        // Ignored or expected.
    }
}
 

@Test
public void testGetFromBogusIterator() {
    try {
        client.getRecords(GetRecordsRequest.builder().shardIterator("bogusmonkeys").build());
        Assert.fail("Expected InvalidArgumentException");
    } catch (InvalidArgumentException exception) {
        // Ignored or expected.
    }

}
 

@Test
public void testGetRecordsWithResourceNotFoundException() throws Exception {
    final ArgumentCaptor<GetShardIteratorRequest> iteratorCaptor = ArgumentCaptor
            .forClass(GetShardIteratorRequest.class);
    final ArgumentCaptor<GetRecordsRequest> recordsCaptor = ArgumentCaptor.forClass(GetRecordsRequest.class);
    // Set up arguments used by proxy
    final String nextIterator = "TestShardIterator";

    final GetShardIteratorRequest expectedIteratorRequest = makeGetShardIteratorRequest(
            ShardIteratorType.LATEST.name());
    final GetRecordsRequest expectedRecordsRequest = makeGetRecordsRequest(nextIterator);

    final CompletableFuture<GetRecordsResponse> future = mock(CompletableFuture.class);

    // Set up proxy mock methods
    when(kinesisClient.getShardIterator(iteratorCaptor.capture()))
            .thenReturn(makeGetShardIteratorResonse(nextIterator));
    when(kinesisClient.getRecords(recordsCaptor.capture())).thenReturn(future);
    when(future.get(anyLong(), any(TimeUnit.class))).thenThrow(
            new ExecutionException(ResourceNotFoundException.builder().message("Test Exception").build()));

    // Create data fectcher and initialize it with latest type checkpoint
    kinesisDataFetcher.initialize(SentinelCheckpoint.LATEST.toString(), INITIAL_POSITION_LATEST);
    final GetRecordsRetrievalStrategy getRecordsRetrievalStrategy = new SynchronousGetRecordsRetrievalStrategy(
            kinesisDataFetcher);
    try {
        // Call records of dataFetcher which will throw an exception
        getRecordsRetrievalStrategy.getRecords(MAX_RECORDS);
    } finally {
        // Test shard has reached the end
        assertTrue("Shard should reach the end", kinesisDataFetcher.isShardEndReached());
        assertEquals(expectedIteratorRequest.startingSequenceNumber(), iteratorCaptor.getValue().startingSequenceNumber());
        assertEquals(expectedRecordsRequest.shardIterator(), recordsCaptor.getValue().shardIterator());
    }
}
 

private void testInitializeAndFetch(final String iteratorType, final String seqNo,
        final InitialPositionInStreamExtended initialPositionInStream) throws Exception {
    final ArgumentCaptor<GetShardIteratorRequest> iteratorCaptor = ArgumentCaptor
            .forClass(GetShardIteratorRequest.class);
    final ArgumentCaptor<GetRecordsRequest> recordsCaptor = ArgumentCaptor.forClass(GetRecordsRequest.class);
    final String iterator = "foo";
    final List<Record> expectedRecords = Collections.emptyList();
    GetShardIteratorRequest expectedIteratorRequest = makeGetShardIteratorRequest(iteratorType);
    if (iteratorType.equals(ShardIteratorType.AT_TIMESTAMP.toString())) {
        expectedIteratorRequest = expectedIteratorRequest.toBuilder()
                .timestamp(initialPositionInStream.getTimestamp().toInstant()).build();
    } else if (iteratorType.equals(ShardIteratorType.AT_SEQUENCE_NUMBER.toString())) {
        expectedIteratorRequest = expectedIteratorRequest.toBuilder().startingSequenceNumber(seqNo).build();
    }
    final GetRecordsRequest expectedRecordsRequest = makeGetRecordsRequest(iterator);

    when(kinesisClient.getShardIterator(iteratorCaptor.capture()))
            .thenReturn(makeGetShardIteratorResonse(iterator));

    when(kinesisClient.getRecords(recordsCaptor.capture()))
            .thenReturn(makeGetRecordsResponse(null, expectedRecords));

    Checkpointer checkpoint = mock(Checkpointer.class);
    when(checkpoint.getCheckpoint(SHARD_ID)).thenReturn(new ExtendedSequenceNumber(seqNo));

    final GetRecordsRetrievalStrategy getRecordsRetrievalStrategy = new SynchronousGetRecordsRetrievalStrategy(
            kinesisDataFetcher);
    kinesisDataFetcher.initialize(seqNo, initialPositionInStream);

    assertEquals(expectedRecords, getRecordsRetrievalStrategy.getRecords(MAX_RECORDS).records());
    verify(kinesisClient, times(1)).getShardIterator(any(GetShardIteratorRequest.class));
    verify(kinesisClient, times(1)).getRecords(any(GetRecordsRequest.class));
}
 

private void testGets(final String streamName, final Shard shard) {
    // Get an iterator for the first shard.
    GetShardIteratorResponse iteratorResult = client.getShardIterator(
            GetShardIteratorRequest.builder()
                                   .streamName(streamName)
                                   .shardId(shard.shardId())
                                   .shardIteratorType(ShardIteratorType.AT_SEQUENCE_NUMBER)
                                   .startingSequenceNumber(shard.sequenceNumberRange().startingSequenceNumber())
                                   .build());
    Assert.assertNotNull(iteratorResult);

    String iterator = iteratorResult.shardIterator();
    Assert.assertNotNull(iterator);

    int tries = 0;
    GetRecordsResponse result;
    List<Record> records;

    // Read the first record from the first shard (looping until it's
    // available).
    while (true) {
        tries += 1;
        if (tries > 100) {
            Assert.fail("Failed to read any records after 100 seconds");
        }

        result = client.getRecords(GetRecordsRequest.builder()
                                                    .shardIterator(iterator)
                                                    .limit(1)
                                                    .build());
        Assert.assertNotNull(result);
        Assert.assertNotNull(result.records());
        Assert.assertNotNull(result.nextShardIterator());

        records = result.records();
        if (records.size() > 0) {
            long arrivalTime = records.get(0).approximateArrivalTimestamp().toEpochMilli();
            Long delta = Math.abs(Instant.now().minusMillis(arrivalTime).toEpochMilli());
            // Assert that the arrival date is within 5 minutes of the current date to make sure it unmarshalled correctly.
            assertThat(delta, Matchers.lessThan(60 * 5000L));
            break;
        }

        try {
            Thread.sleep(1000);
        } catch (InterruptedException exception) {
            throw new RuntimeException(exception);
        }

        iterator = result.nextShardIterator();
    }

    System.out.println("  [Succeeded after " + tries + " tries]");
    Assert.assertEquals(1, records.size());
    validateRecord(records.get(0), "See No Evil");

    // Read the second record from the first shard.
    result = client.getRecords(GetRecordsRequest.builder()
                                                .shardIterator(result.nextShardIterator())
                                                .build());
    Assert.assertNotNull(result);
    Assert.assertNotNull(result.records());
    Assert.assertNotNull(result.nextShardIterator());

    records = result.records();
    Assert.assertEquals(1, records.size());
    validateRecord(records.get(0), "See No Evil");

    // Try to read some more, get EOF.
    result = client.getRecords(GetRecordsRequest.builder()
                                                .shardIterator(result.nextShardIterator())
                                                .build());
    Assert.assertNotNull(result);
    Assert.assertNotNull(result.records());
    Assert.assertTrue(result.records().isEmpty());
    Assert.assertNull(result.nextShardIterator());
}
 

public static void getStockTrades(KinesisClient kinesisClient, String streamName) {

            String shardIterator;
            String lastShardId = null;

            // Retrieve the shards from a data stream
            DescribeStreamRequest describeStreamRequest = DescribeStreamRequest.builder()
                    .streamName(streamName)
                    .build();
            List<Shard> shards = new ArrayList<>();

            DescribeStreamResponse streamRes;
            do {
                streamRes = kinesisClient.describeStream(describeStreamRequest);
                shards.addAll(streamRes.streamDescription().shards());

                if (shards.size() > 0) {
                    lastShardId = shards.get(shards.size() - 1).shardId();
                }
            } while (streamRes.streamDescription().hasMoreShards());

            GetShardIteratorRequest itReq = GetShardIteratorRequest.builder()
                    .streamName(streamName)
                    .shardIteratorType("TRIM_HORIZON")
                    .shardId(shards.get(0).shardId())
                    .build();

            GetShardIteratorResponse shardIteratorResult = kinesisClient.getShardIterator(itReq);
            shardIterator = shardIteratorResult.shardIterator();

            // Continuously read data records from a shard
            List<Record> records;

            // Create a GetRecordsRequest with the existing shardIterator,
            // and set maximum records to return to 1000
            GetRecordsRequest recordsRequest = GetRecordsRequest.builder()
                     .shardIterator(shardIterator)
                     .limit(1000)
                     .build();

           GetRecordsResponse result = kinesisClient.getRecords(recordsRequest);

           // Put result into a record list, result might be empty
           records = result.records();

            // Print records
            for (Record record : records) {
                SdkBytes byteBuffer = record.data();
                System.out.println(String.format("Seq No: %s - %s", record.sequenceNumber(),
                 new String(byteBuffer.asByteArray())));
             }

             }
 

public static GetRecordsRequest.Builder getRecordsRequestBuilder() {
    return appendUserAgent(GetRecordsRequest.builder());
}
 

private GetRecordsRequest makeGetRecordsRequest(String shardIterator) {
    return GetRecordsRequest.builder().shardIterator(shardIterator).limit(MAX_RECORDS).build();
}
 

@Test
public void testFetcherDoesNotAdvanceWithoutAccept() throws InterruptedException, ExecutionException {
    final ArgumentCaptor<GetShardIteratorRequest> iteratorCaptor = ArgumentCaptor
            .forClass(GetShardIteratorRequest.class);
    final ArgumentCaptor<GetRecordsRequest> recordsCaptor = ArgumentCaptor.forClass(GetRecordsRequest.class);
    final String initialIterator = "InitialIterator";
    final String nextIterator1 = "NextIteratorOne";
    final String nextIterator2 = "NextIteratorTwo";
    final CompletableFuture<GetRecordsResponse> nonAdvancingResult1 = makeGetRecordsResponse(initialIterator, null);
    final CompletableFuture<GetRecordsResponse> nonAdvancingResult2 = makeGetRecordsResponse(nextIterator1, null);
    final CompletableFuture<GetRecordsResponse> finalNonAdvancingResult = makeGetRecordsResponse(nextIterator2,
            null);
    final CompletableFuture<GetRecordsResponse> advancingResult1 = makeGetRecordsResponse(nextIterator1, null);
    final CompletableFuture<GetRecordsResponse> advancingResult2 = makeGetRecordsResponse(nextIterator2, null);
    final CompletableFuture<GetRecordsResponse> finalAdvancingResult = makeGetRecordsResponse(null, null);

    when(kinesisClient.getShardIterator(iteratorCaptor.capture()))
            .thenReturn(makeGetShardIteratorResonse(initialIterator));
    when(kinesisClient.getRecords(recordsCaptor.capture())).thenReturn(nonAdvancingResult1, advancingResult1,
            nonAdvancingResult2, advancingResult2, finalNonAdvancingResult, finalAdvancingResult);

    kinesisDataFetcher.initialize("TRIM_HORIZON",
            InitialPositionInStreamExtended.newInitialPosition(InitialPositionInStream.TRIM_HORIZON));

    assertNoAdvance(nonAdvancingResult1.get(), initialIterator);
    assertAdvanced(advancingResult1.get(), initialIterator, nextIterator1);
    verify(kinesisClient, times(2)).getRecords(any(GetRecordsRequest.class));

    assertNoAdvance(nonAdvancingResult2.get(), nextIterator1);
    assertAdvanced(advancingResult2.get(), nextIterator1, nextIterator2);
    verify(kinesisClient, times(4)).getRecords(any(GetRecordsRequest.class));

    assertNoAdvance(finalNonAdvancingResult.get(), nextIterator2);
    assertAdvanced(finalAdvancingResult.get(), nextIterator2, null);
    verify(kinesisClient, times(6)).getRecords(any(GetRecordsRequest.class));



    reset(kinesisClient);

    DataFetcherResult terminal = kinesisDataFetcher.getRecords();
    assertTrue(terminal.isShardEnd());
    assertNotNull(terminal.getResult());

    final GetRecordsResponse terminalResult = terminal.getResult();
    assertNotNull(terminalResult.records());
    assertEquals(0, terminalResult.records().size());
    assertNull(terminalResult.nextShardIterator());
    assertEquals(kinesisDataFetcher.TERMINAL_RESULT, terminal);

    verify(kinesisClient, never()).getRecords(any(GetRecordsRequest.class));
}