software.amazon.awssdk.services.kinesis.KinesisAsyncClient Java Examples

@Test
public void shouldIterateToNextIdWithRetryAfterException() {
    // given
    GetRecordsResponse response = GetRecordsResponse.builder()
            .records(emptyList())
            .nextShardIterator("nextIteratorId")
            .millisBehindLatest(42L)
            .build();
    KinesisAsyncClient kinesisClient = someKinesisClient();
    when(kinesisClient.getRecords(any(GetRecordsRequest.class)))
            .thenThrow(new CompletionException(ExpiredIteratorException.builder().message("forced test exception").build()))
            .thenReturn(completedFuture(response));
    final KinesisShardIterator shardIterator = new KinesisShardIterator(kinesisClient, "", fromHorizon("someShard"));

    // when
    shardIterator.next();

    // then
    assertThat(shardIterator.getId(), is("nextIteratorId"));
}
 

@Test
public void shouldIterateToNextId() {
    // given
    GetRecordsResponse response = GetRecordsResponse.builder()
            .records(emptyList())
            .nextShardIterator("nextIteratorId")
            .millisBehindLatest(42L)
            .build();
    KinesisAsyncClient kinesisClient = someKinesisClient();
    when(kinesisClient.getRecords(any(GetRecordsRequest.class))).thenReturn(completedFuture(response));
    final KinesisShardIterator shardIterator = new KinesisShardIterator(kinesisClient, "", fromHorizon("someShard"));

    // when
    shardIterator.next();

    // then
    assertThat(shardIterator.getId(), is("nextIteratorId"));
}
 

@Test
public void shouldFetchRecords() {
    // given
    GetRecordsResponse response = GetRecordsResponse.builder()
            .millisBehindLatest(42L)
            .records(builder()
                    .sequenceNumber("someSeqNumber")
                    .partitionKey("foo")
                    .approximateArrivalTimestamp(now())
                    .build())
            .build();
    final KinesisAsyncClient kinesisClient = someKinesisClient();
    when(kinesisClient.getRecords(any(GetRecordsRequest.class))).thenReturn(completedFuture(response));
    final KinesisShardIterator shardIterator = new KinesisShardIterator(kinesisClient, "", fromHorizon("someShard"));

    // when
    final ShardResponse fetchedResponse = shardIterator.next();

    // then
    assertThat(fetchedResponse.getShardPosition(), is(fromPosition("someShard", "someSeqNumber")));
    GetRecordsRequest expectedRequest = GetRecordsRequest.builder()
            .shardIterator("someShardIterator")
            .limit(KinesisShardIterator.FETCH_RECORDS_LIMIT)
            .build();
    verify(kinesisClient).getRecords(expectedRequest);
}
 

@Test
public void shouldReturnAtTimestampIterator() {
    // when
    final Instant now = now();
    final KinesisAsyncClient kinesisClient = someKinesisClient();
    final KinesisShardIterator shardIterator = new KinesisShardIterator(kinesisClient, "someChannel", fromTimestamp("someShard", now));

    // then
    assertThat(shardIterator.getId(), is("someShardIterator"));

    GetShardIteratorRequest expectedRequest = GetShardIteratorRequest.builder()
            .streamName("someChannel")
            .shardId("someShard")
            .shardIteratorType(ShardIteratorType.AT_TIMESTAMP)
            .timestamp(now)
            .build();
    verify(kinesisClient).getShardIterator(expectedRequest);
}
 

@Test
public void shouldReturnAfterSequenceNumberIterator() {
    // when
    final KinesisAsyncClient kinesisClient = someKinesisClient();
    final KinesisShardIterator shardIterator = new KinesisShardIterator(kinesisClient, "someChannel", fromPosition("someShard", "1"));

    // then
    assertThat(shardIterator.getId(), is("someShardIterator"));

    GetShardIteratorRequest expectedRequest = GetShardIteratorRequest.builder()
            .streamName("someChannel")
            .shardId("someShard")
            .shardIteratorType(ShardIteratorType.AFTER_SEQUENCE_NUMBER)
            .startingSequenceNumber("1")
            .build();
    verify(kinesisClient).getShardIterator(expectedRequest);
}
 

@Test
public void shouldFetchSingleMessage() {
    final Record record = Record.builder()
            .sequenceNumber("43")
            .approximateArrivalTimestamp(now())
            .partitionKey("someKey")
            .build();
    final GetRecordsResponse response = GetRecordsResponse.builder()
            .records(record)
            .nextShardIterator("nextShardIterator")
            .millisBehindLatest(1234L)
            .build();
    final KinesisAsyncClient kinesisClient = someKinesisClient();
    when(kinesisClient.getRecords(any(GetRecordsRequest.class))).thenReturn(completedFuture(response));

    final KinesisShardIterator iterator = new KinesisShardIterator(kinesisClient, "", fromPosition("someShard", "42"), 1);
    final ShardResponse shardResponse = iterator.next();

    assertThat(shardResponse.getMessages(), hasSize(1));
    final Message<String> message = shardResponse.getMessages().get(0);
    assertThat(message.getKey(), is(Key.of("someKey")));
    assertThat(message.getPayload(), is(nullValue()));
    assertThat(message.getHeader().getShardPosition().get(), is(fromPosition("someShard", "43")));
}
 

public static void main(String[] args) {

        // snippet-start:[kinesis.java2.stream_example.setup]
        Region region = Region.US_EAST_1;
        KinesisAsyncClient client = KinesisAsyncClient.builder()
        .region(region)
        .build();

        SubscribeToShardRequest request = SubscribeToShardRequest.builder()
                .consumerARN(CONSUMER_ARN)
                .shardId("arn:aws:kinesis:us-east-1:814548047983:stream/StockTradeStream")
                .startingPosition(s -> s.type(ShardIteratorType.LATEST)).build();

        // snippet-end:[kinesis.java2.stream_example.setup]
        SubscribeToShardResponseHandler responseHandler = SubscribeToShardResponseHandler
                .builder()
                .onError(t -> System.err.println("Error during stream - " + t.getMessage()))
                .subscriber(MySubscriber::new)
                .build();

        client.subscribeToShard(request, responseHandler);
        client.close();
    }
 

private KinesisAsyncClient createClient() {

        ClientAsyncConfiguration clientConfiguration = ClientAsyncConfiguration.builder().build();

        // Reading credentials from ENV-variables
        AwsCredentialsProvider awsCredentialsProvider = DefaultCredentialsProvider.builder().build();

        // Configuring Kinesis-client with configuration
        String tmp = System.getenv("REGION");

        Region myRegion;
        if (tmp == null || tmp.trim().length() == 0) {
            myRegion = Region.US_EAST_1;
            LOGGER.info("Using default region");
        } else {
            myRegion = Region.of(tmp);
        }

        LOGGER.info("Deploying in Region " + myRegion.toString());

        return KinesisAsyncClient.builder()
                .asyncConfiguration(clientConfiguration)
                .credentialsProvider(awsCredentialsProvider)
                .region(myRegion)
                .build();
    }
 

static boolean doesChannelExist(final KinesisAsyncClient kinesisClient,
                                final String channelName,
                                final String from) {

    final ListStreamsResponse response = kinesisClient
            .listStreams(ListStreamsRequest.builder().exclusiveStartStreamName(from).build())
            .join();
    final List<String> streamNames = response.streamNames();
    if (streamNames.stream().anyMatch(channelName::equals)) {
        return true;
    } else if (response.hasMoreStreams()) {
        return doesChannelExist(kinesisClient, channelName, streamNames.get(streamNames.size() - 1));
    }
    return false;
}
 

/**
 * Creates a SubscribeToShardResponseHandler.Visitor using the builder, which lets you register an event handler for
 * all events you're interested in instead of implementing the interface
 */
// snippet-start:[kinesis.java2.stream_example.visitor]
private static CompletableFuture<Void> responseHandlerBuilderVisitorBuilder(KinesisAsyncClient client, SubscribeToShardRequest request) {
    SubscribeToShardResponseHandler.Visitor visitor = SubscribeToShardResponseHandler.Visitor
            .builder()
            .onSubscribeToShardEvent(e -> System.out.println("Received subscribe to shard event " + e))
            .build();
    SubscribeToShardResponseHandler responseHandler = SubscribeToShardResponseHandler
            .builder()
            .onError(t -> System.err.println("Error during stream - " + t.getMessage()))
            .subscriber(visitor)
            .build();
    return client.subscribeToShard(request, responseHandler);
}
 

public static void createChannelIfNotExists(final KinesisAsyncClient kinesisClient,
                                            final String channelName,
                                            final int numberOfShards) {
    if (!doesChannelExist(kinesisClient, channelName, null)) {
        kinesisClient.createStream(CreateStreamRequest.builder()
                .streamName(channelName)
                .shardCount(numberOfShards)
                .build())
                .join();
        DescribeStreamResponse describeStreamResponse;
        int retries = 0;
        do {
            describeStreamResponse = kinesisClient.describeStream(DescribeStreamRequest.builder()
                    .streamName(channelName)
                    .build())
                    .join();
            
            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
            retries++;
        } while ((StreamStatus.ACTIVE != describeStreamResponse.streamDescription().streamStatus()) && retries < MAX_RETRIES);

        if (describeStreamResponse.streamDescription().streamStatus() != StreamStatus.ACTIVE) {
            throw new RuntimeException(String.format("Timeout while waiting to create stream %s", channelName));
        }
    }
}
 

@Bean
@Primary
public KinesisAsyncClient kinesisAsyncClient(final @Value("${test.environment:local}") String testEnvironment,
                                             final AwsCredentialsProvider credentialsProvider) {
    // kinesalite does not support cbor at the moment (v1.11.6)
    System.setProperty("aws.cborEnabled", "false");
    LOG.info("kinesis client for local tests");
    final KinesisAsyncClient kinesisClient;
    if (testEnvironment.equals("local")) {
        kinesisClient = KinesisAsyncClient.builder()
                .httpClient(
                        // Disables HTTP2 because of problems with LocalStack
                        NettyNioAsyncHttpClient.builder()
                                .protocol(Protocol.HTTP1_1)
                                .build())
                .endpointOverride(URI.create("http://localhost:4568"))
                .region(Region.EU_CENTRAL_1)
                .credentialsProvider(credentialsProvider)
                .build();
    } else {
        kinesisClient = KinesisAsyncClient.builder()
                .credentialsProvider(credentialsProvider)
                .build();
    }
    createChannelIfNotExists(kinesisClient, INTEGRATION_TEST_STREAM, 2);
    return kinesisClient;
}
 

/**
 * Because a Flux is also a publisher, the publisherTransformer method integrates nicely with Reactor. Notice that
 * you must adapt to an SdkPublisher.
 */
private static CompletableFuture<Void> responseHandlerBuilder_OnEventStream_Reactor(KinesisAsyncClient client, SubscribeToShardRequest request) {
    SubscribeToShardResponseHandler responseHandler = SubscribeToShardResponseHandler
        .builder()
        .onError(t -> System.err.println("Error during stream - " + t.getMessage()))
        .publisherTransformer(p -> Flux.from(p).limitRate(100).as(SdkPublisher::adapt))
        .build();
    return client.subscribeToShard(request, responseHandler);
}
 

/**
 * Subscribes to the publisher using the onEventStream lifecycle callback method, which allows greater control
 * over the publisher and allows transformation methods on the publisher, like map and buffer
 */
private static CompletableFuture<Void> responseHandlerBuilderOnEventStream(KinesisAsyncClient client, SubscribeToShardRequest request) {
    SubscribeToShardResponseHandler responseHandler = SubscribeToShardResponseHandler
            .builder()
            .onError(t -> System.err.println("Error during stream - " + t.getMessage()))
            .onEventStream(p -> p.filter(SubscribeToShardEvent.class).subscribe(new MySubscriber()))
            .build();
    return client.subscribeToShard(request, responseHandler);
}
 

/**
 * Subscribes to the stream of events by implementing the SubscribeToShardResponseHandler.Visitor interface
 */
private static CompletableFuture<Void> responseHandlerBuilderVisitor(KinesisAsyncClient client, SubscribeToShardRequest request) {
    SubscribeToShardResponseHandler.Visitor visitor = new SubscribeToShardResponseHandler.Visitor() {
        @Override
        public void visit(SubscribeToShardEvent event) {
            System.out.println("Received subscribe to shard event " + event);
        }
    };
    SubscribeToShardResponseHandler responseHandler = SubscribeToShardResponseHandler
            .builder()
            .onError(t -> System.err.println("Error during stream - " + t.getMessage()))
            .subscriber(visitor)
            .build();
    return client.subscribeToShard(request, responseHandler);
}
 

/**
 * Creates a SubscribeToShardResponseHandler the classic way by implementing the interface
 */
// snippet-start:[kinesis.java2.stream_example.custom_handler]
private static CompletableFuture<Void> responseHandlerBuilderClassic(KinesisAsyncClient client, SubscribeToShardRequest request) {
    SubscribeToShardResponseHandler responseHandler = new SubscribeToShardResponseHandler() {

        @Override
        public void responseReceived(SubscribeToShardResponse response) {
            System.out.println("Receieved initial response");
        }

        @Override
        public void onEventStream(SdkPublisher<SubscribeToShardEventStream> publisher) {
            publisher
                    // Filter to only SubscribeToShardEvents
                    .filter(SubscribeToShardEvent.class)
                    // Flat map into a publisher of just records
                    .flatMapIterable(SubscribeToShardEvent::records)
                    // Limit to 1000 total records
                    .limit(1000)
                    // Batch records into lists of 25
                    .buffer(25)
                    // Print out each record batch
                    .subscribe(batch -> System.out.println("Record Batch - " + batch));
        }

        @Override
        public void complete() {
            System.out.println("All records stream successfully");
        }

        @Override
        public void exceptionOccurred(Throwable throwable) {
            System.err.println("Error during stream - " + throwable.getMessage());
        }
    };
    return client.subscribeToShard(request, responseHandler);
}
 

/**
 * Use the SubscribeToShardResponseHandler.Builder and a traditional subscriber
 */
// snippet-start:[kinesis.java2.stream_example.subscribe]
private static CompletableFuture<Void> responseHandlerBuilderSubscriber(KinesisAsyncClient client, SubscribeToShardRequest request) {
    SubscribeToShardResponseHandler responseHandler = SubscribeToShardResponseHandler
            .builder()
            .onError(t -> System.err.println("Error during stream - " + t.getMessage()))
            .subscriber(MySubscriber::new)
            .build();
    return client.subscribeToShard(request, responseHandler);
}
 

@Override
protected KinesisAsyncClient createClient(AwsCredentialsProvider credentials, ClientOverrideConfiguration configuration,
    ThreadPoolExecutor executor, Region region, Optional<URI> endpointOverride) {
  KinesisAsyncClientBuilder builder = KinesisAsyncClient.builder()
    .credentialsProvider(credentials)
    .region(region)
    .overrideConfiguration(configuration)
    .asyncConfiguration(b -> b.advancedOption(SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR, executor));
  
  endpointOverride.ifPresent(endpoint -> builder.endpointOverride(endpoint));
  
  return builder.build();
}
 

@Override
protected void validateStreamName(KinesisAsyncClient client, String streamName) {
  DescribeStreamResponse describeResult;
  try {
    describeResult = getClient().describeStream(b -> b.streamName(streamName).build()).get();
    StreamStatus streamStatus = describeResult.streamDescription().streamStatus();
    if(!StreamStatus.ACTIVE.equals(streamStatus) && !StreamStatus.UPDATING.equals(streamStatus)) {
      setInitializationFailed(true);
      addError("Stream " + streamName + " is not ready (in active/updating status) for appender: " + name);
    }
  }
  catch(InterruptedException ie) {
    setInitializationFailed(true);
    addError("Interrupted while attempting to describe " + streamName, ie);
  }
  catch(ExecutionException ee) {
    setInitializationFailed(true);
    addError("Error executing the operation", ee);
  }
  catch(ResourceNotFoundException rnfe) {
    setInitializationFailed(true);
    addError("Stream " + streamName + " doesn't exist for appender: " + name, rnfe);
  }
  catch(AwsServiceException ase) {
    setInitializationFailed(true);
    addError("Error connecting to AWS to verify stream " + streamName + " for appender: " + name, ase);
  }
}
 

@Test
public void shouldNotMatchMessageQueueSelectors() {
    final MessageInterceptorRegistry interceptorRegistry = mock(MessageInterceptorRegistry.class);
    final KinesisAsyncClient kinesisClient = mock(KinesisAsyncClient.class);

    final KinesisMessageSenderEndpointFactory factory = new KinesisMessageSenderEndpointFactory(interceptorRegistry, kinesisClient);

    assertThat(factory.matches(MessageQueue.class), is(false));
}
 

@Test
public void shouldMatchMessageLogSelectors() {
    final MessageInterceptorRegistry interceptorRegistry = mock(MessageInterceptorRegistry.class);
    final KinesisAsyncClient kinesisClient = mock(KinesisAsyncClient.class);

    final KinesisMessageSenderEndpointFactory factory = new KinesisMessageSenderEndpointFactory(interceptorRegistry, kinesisClient);

    assertThat(factory.matches(MessageLog.class), is(true));
    assertThat(factory.matches(Kinesis.class), is(true));
}
 

public SynchronousPrefetchingRetrievalFactory(String streamName, KinesisAsyncClient kinesisClient,
        RecordsFetcherFactory recordsFetcherFactory, int maxRecords, ExecutorService executorService,
        long idleMillisBetweenCalls, Duration maxFutureWait) {
    this.streamName = streamName;
    this.kinesisClient = kinesisClient;
    this.recordsFetcherFactory = recordsFetcherFactory;
    this.maxRecords = maxRecords;
    this.executorService = executorService;
    this.idleMillisBetweenCalls = idleMillisBetweenCalls;
    this.maxFutureWait = maxFutureWait;
}
 

@Test
public void shouldBuildKinesisMessageSender() {
    final KinesisAsyncClient kinesisClient = mock(KinesisAsyncClient.class);
    when(kinesisClient.listStreams()).thenReturn(completedFuture(ListStreamsResponse.builder().build()));

    final KinesisMessageSenderEndpointFactory factory = new KinesisMessageSenderEndpointFactory(new MessageInterceptorRegistry(), kinesisClient);

    final KinesisMessageSender sender = (KinesisMessageSender) factory.create("foo-stream", MessageFormat.V2);
    assertThat(sender.getChannelName(), is("foo-stream"));
    assertThat(sender.getMessageFormat(), is(MessageFormat.V2));
    assertThat(sender, is(instanceOf(KinesisMessageSender.class)));
}
 

private void describeStreamResponse(final KinesisAsyncClient kinesisClient) {
    DescribeStreamResponse response = createResponseForShards(
            Shard.builder()
                    .shardId("foo")
                    .sequenceNumberRange(SequenceNumberRange.builder()
                            .startingSequenceNumber("42")
                            .endingSequenceNumber("4711")
                            .build())
                    .build()
    );

    when(kinesisClient.describeStream(any(DescribeStreamRequest.class))).thenReturn(completedFuture(response));
}
 

private static KinesisAsyncClient someKinesisClient() {
    final KinesisAsyncClient kinesisClient = mock(KinesisAsyncClient.class);
    when(kinesisClient.getShardIterator(any(GetShardIteratorRequest.class))).thenReturn(completedFuture(GetShardIteratorResponse
            .builder()
            .shardIterator("someShardIterator")
            .build()));
    return kinesisClient;
}
 

public SynchronousBlockingRetrievalFactory(String streamName, KinesisAsyncClient kinesisClient, RecordsFetcherFactory recordsFetcherFactory, int maxRecords, Duration kinesisRequestTimeout) {
    this.streamName = streamName;
    this.kinesisClient = kinesisClient;
    this.recordsFetcherFactory = recordsFetcherFactory;
    this.maxRecords = maxRecords;
    this.kinesisRequestTimeout = kinesisRequestTimeout;
}
 

public KinesisDataFetcher(KinesisAsyncClient kinesisClient, String streamName, String shardId, int maxRecords, MetricsFactory metricsFactory, Duration maxFutureWait) {
    this.kinesisClient = kinesisClient;
    this.streamName = streamName;
    this.shardId = shardId;
    this.maxRecords = maxRecords;
    this.metricsFactory = metricsFactory;
    this.maxFutureWait = maxFutureWait;
}
 

@Override
protected FanOutConsumerRegistration createConsumerRegistration(KinesisAsyncClient client, String stream,
                                                                String consumerToCreate) {
    this.stream = stream;
    this.consumerToCreate = consumerToCreate;
    return consumerRegistration;
}
 

public KinesisShardDetector(KinesisAsyncClient kinesisClient, String streamName, long listShardsBackoffTimeInMillis,
        int maxListShardsRetryAttempts, long listShardsCacheAllowedAgeInSeconds, int maxCacheMissesBeforeReload,
        int cacheMissWarningModulus, Duration kinesisRequestTimeout) {
    this.kinesisClient = kinesisClient;
    this.streamName = streamName;
    this.listShardsBackoffTimeInMillis = listShardsBackoffTimeInMillis;
    this.maxListShardsRetryAttempts = maxListShardsRetryAttempts;
    this.listShardsCacheAllowedAgeInSeconds = listShardsCacheAllowedAgeInSeconds;
    this.maxCacheMissesBeforeReload = maxCacheMissesBeforeReload;
    this.cacheMissWarningModulus = cacheMissWarningModulus;
    this.kinesisRequestTimeout = kinesisRequestTimeout;
}
 

public static void main(String[] args) {

        KinesisAsyncClient client = KinesisAsyncClient.create();

        SubscribeToShardRequest request = SubscribeToShardRequest.builder()
                .consumerARN(CONSUMER_ARN)
                .shardId("shardId-000000000000")
                .startingPosition(StartingPosition.builder().type(ShardIteratorType.LATEST).build())
                .build();

        responseHandlerBuilder_Reactor(client, request).join();

        client.close();
    }