Search in sources :

Example 1 with Shard

use of software.amazon.awssdk.services.kinesis.model.Shard in project beam by apache.

the class ShardRecordsIterator method findSuccessiveShardRecordIterators.

List<ShardRecordsIterator> findSuccessiveShardRecordIterators() throws TransientKinesisException {
    List<Shard> shards = kinesis.listShardsFollowingClosedShard(streamName, shardId);
    List<ShardRecordsIterator> successiveShardRecordIterators = new ArrayList<>();
    for (Shard shard : shards) {
        if (shardId.equals(shard.parentShardId())) {
            ShardCheckpoint shardCheckpoint = new ShardCheckpoint(streamName, shard.shardId(), new StartingPoint(InitialPositionInStream.TRIM_HORIZON));
            successiveShardRecordIterators.add(new ShardRecordsIterator(shardCheckpoint, kinesis, watermarkPolicyFactory));
        }
    }
    return successiveShardRecordIterators;
}
Also used : ArrayList(java.util.ArrayList) Shard(software.amazon.awssdk.services.kinesis.model.Shard)

Example 2 with Shard

use of software.amazon.awssdk.services.kinesis.model.Shard in project beam by apache.

the class SimplifiedKinesisClient method listShards.

private List<Shard> listShards(final String streamName, final ShardFilter shardFilter) throws TransientKinesisException {
    return wrapExceptions(() -> {
        ImmutableList.Builder<Shard> shardsBuilder = ImmutableList.builder();
        String currentNextToken = null;
        do {
            ListShardsRequest.Builder reqBuilder = ListShardsRequest.builder().maxResults(LIST_SHARDS_MAX_RESULTS).shardFilter(shardFilter);
            if (currentNextToken != null) {
                reqBuilder.nextToken(currentNextToken);
            } else {
                reqBuilder.streamName(streamName);
            }
            ListShardsResponse response = kinesis.listShards(reqBuilder.build());
            shardsBuilder.addAll(response.shards());
            currentNextToken = response.nextToken();
        } while (currentNextToken != null);
        return shardsBuilder.build();
    });
}
Also used : ListShardsRequest(software.amazon.awssdk.services.kinesis.model.ListShardsRequest) ListShardsResponse(software.amazon.awssdk.services.kinesis.model.ListShardsResponse) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) Shard(software.amazon.awssdk.services.kinesis.model.Shard)

Example 3 with Shard

use of software.amazon.awssdk.services.kinesis.model.Shard in project beam by apache.

the class KinesisIOReadTest method testRecords.

private List<List<Record>> testRecords(int shards, int events) {
    final Instant now = DateTime.now().toInstant();
    Function<Integer, List<Record>> dataStream = shard -> range(0, events).mapToObj(off -> record(now, shard, off)).collect(toList());
    return range(0, shards).boxed().map(dataStream).collect(toList());
}
Also used : PipelineExecutionException(org.apache.beam.sdk.Pipeline.PipelineExecutionException) IntStream.range(java.util.stream.IntStream.range) ArgumentMatchers.argThat(org.mockito.ArgumentMatchers.argThat) AwsCredentialsProvider(software.amazon.awssdk.auth.credentials.AwsCredentialsProvider) Shard(software.amazon.awssdk.services.kinesis.model.Shard) Assertions.assertThat(org.assertj.core.api.Assertions.assertThat) Duration.standardSeconds(org.joda.time.Duration.standardSeconds) BiFunction(java.util.function.BiFunction) CloudWatchClient(software.amazon.awssdk.services.cloudwatch.CloudWatchClient) SdkBytes(software.amazon.awssdk.core.SdkBytes) ArgumentMatcher(org.mockito.ArgumentMatcher) GetRecordsResponse(software.amazon.awssdk.services.kinesis.model.GetRecordsResponse) URI(java.net.URI) LimitExceededException(software.amazon.awssdk.services.kinesis.model.LimitExceededException) DefaultCredentialsProvider(software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider) List(java.util.List) ParDo(org.apache.beam.sdk.transforms.ParDo) Function.identity(java.util.function.Function.identity) Read(org.apache.beam.sdk.io.aws2.kinesis.KinesisIO.Read) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) MockitoJUnitRunner(org.mockito.junit.MockitoJUnitRunner) Mockito.mock(org.mockito.Mockito.mock) ArgumentMatchers.any(org.mockito.ArgumentMatchers.any) GetShardIteratorRequest(software.amazon.awssdk.services.kinesis.model.GetShardIteratorRequest) Mock(org.mockito.Mock) TRIM_HORIZON(software.amazon.kinesis.common.InitialPositionInStream.TRIM_HORIZON) RunWith(org.junit.runner.RunWith) StaticCredentialsProvider(software.amazon.awssdk.auth.credentials.StaticCredentialsProvider) Function(java.util.function.Function) TestPipeline(org.apache.beam.sdk.testing.TestPipeline) GetShardIteratorResponse(software.amazon.awssdk.services.kinesis.model.GetShardIteratorResponse) Region(software.amazon.awssdk.regions.Region) MockClientBuilderFactory(org.apache.beam.sdk.io.aws2.MockClientBuilderFactory) IntFunction(java.util.function.IntFunction) Before(org.junit.Before) DoFn(org.apache.beam.sdk.transforms.DoFn) KinesisClient(software.amazon.awssdk.services.kinesis.KinesisClient) PAssert(org.apache.beam.sdk.testing.PAssert) UTF_8(java.nio.charset.StandardCharsets.UTF_8) DateTime(org.joda.time.DateTime) KinesisClientBuilder(software.amazon.awssdk.services.kinesis.KinesisClientBuilder) Test(org.junit.Test) Mockito.when(org.mockito.Mockito.when) PCollection(org.apache.beam.sdk.values.PCollection) Iterables.concat(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.Iterables.concat) GetRecordsRequest(software.amazon.awssdk.services.kinesis.model.GetRecordsRequest) Collectors.toList(java.util.stream.Collectors.toList) Rule(org.junit.Rule) StaticSupplier(org.apache.beam.sdk.io.aws2.StaticSupplier) CloudWatchClientBuilder(software.amazon.awssdk.services.cloudwatch.CloudWatchClientBuilder) Instant(org.joda.time.Instant) ListShardsRequest(software.amazon.awssdk.services.kinesis.model.ListShardsRequest) ListShardsResponse(software.amazon.awssdk.services.kinesis.model.ListShardsResponse) ClientConfiguration(org.apache.beam.sdk.io.aws2.common.ClientConfiguration) Record(software.amazon.awssdk.services.kinesis.model.Record) AwsBasicCredentials(software.amazon.awssdk.auth.credentials.AwsBasicCredentials) Instant(org.joda.time.Instant) List(java.util.List) ImmutableList(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList) Collectors.toList(java.util.stream.Collectors.toList)

Example 4 with Shard

use of software.amazon.awssdk.services.kinesis.model.Shard in project beam by apache.

the class SimplifiedKinesisClientTest method shouldListAllShardsForTimestampOutsideStreamRetentionAfterStreamCreationTimestamp.

@Test
public void shouldListAllShardsForTimestampOutsideStreamRetentionAfterStreamCreationTimestamp() throws Exception {
    Shard shard1 = Shard.builder().shardId(SHARD_1).build();
    Shard shard2 = Shard.builder().shardId(SHARD_2).build();
    Shard shard3 = Shard.builder().shardId(SHARD_3).build();
    int retentionPeriodHours = 3;
    int startingPointHours = 5;
    int hoursSinceStreamCreation = 6;
    Instant streamCreationTimestamp = CURRENT_TIMESTAMP.minus(Duration.standardHours(hoursSinceStreamCreation));
    Instant startingPointTimestampAfterStreamRetentionTimestamp = CURRENT_TIMESTAMP.minus(Duration.standardHours(startingPointHours));
    when(currentInstantSupplier.get()).thenReturn(CURRENT_TIMESTAMP);
    DescribeStreamSummaryRequest describeStreamRequest = DescribeStreamSummaryRequest.builder().streamName(STREAM).build();
    when(kinesis.describeStreamSummary(describeStreamRequest)).thenReturn(DescribeStreamSummaryResponse.builder().streamDescriptionSummary(StreamDescriptionSummary.builder().retentionPeriodHours(retentionPeriodHours).streamCreationTimestamp(TimeUtil.toJava(streamCreationTimestamp)).build()).build());
    ShardFilter shardFilter = ShardFilter.builder().type(ShardFilterType.AT_TRIM_HORIZON).build();
    when(kinesis.listShards(ListShardsRequest.builder().streamName(STREAM).shardFilter(shardFilter).maxResults(1_000).build())).thenReturn(ListShardsResponse.builder().shards(shard1, shard2, shard3).nextToken(null).build());
    List<Shard> shards = underTest.listShardsAtPoint(STREAM, new StartingPoint(startingPointTimestampAfterStreamRetentionTimestamp));
    assertThat(shards).containsOnly(shard1, shard2, shard3);
}
Also used : Instant(org.joda.time.Instant) DescribeStreamSummaryRequest(software.amazon.awssdk.services.kinesis.model.DescribeStreamSummaryRequest) ShardFilter(software.amazon.awssdk.services.kinesis.model.ShardFilter) Shard(software.amazon.awssdk.services.kinesis.model.Shard) Datapoint(software.amazon.awssdk.services.cloudwatch.model.Datapoint) Test(org.junit.Test)

Example 5 with Shard

use of software.amazon.awssdk.services.kinesis.model.Shard in project beam by apache.

the class SimplifiedKinesisClientTest method shouldListAllShardsForLatest.

@Test
public void shouldListAllShardsForLatest() throws Exception {
    Shard shard1 = Shard.builder().shardId(SHARD_1).build();
    Shard shard2 = Shard.builder().shardId(SHARD_2).build();
    Shard shard3 = Shard.builder().shardId(SHARD_3).build();
    when(kinesis.listShards(ListShardsRequest.builder().streamName(STREAM).shardFilter(ShardFilter.builder().type(ShardFilterType.AT_LATEST).build()).maxResults(1_000).build())).thenReturn(ListShardsResponse.builder().shards(shard1, shard2, shard3).nextToken(null).build());
    List<Shard> shards = underTest.listShardsAtPoint(STREAM, new StartingPoint(InitialPositionInStream.LATEST));
    assertThat(shards).containsOnly(shard1, shard2, shard3);
}
Also used : Shard(software.amazon.awssdk.services.kinesis.model.Shard) Test(org.junit.Test)

Aggregations

Shard (software.amazon.awssdk.services.kinesis.model.Shard)12 Test (org.junit.Test)9 ShardFilter (software.amazon.awssdk.services.kinesis.model.ShardFilter)6 Instant (org.joda.time.Instant)5 ArrayList (java.util.ArrayList)2 ImmutableList (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList)2 Datapoint (software.amazon.awssdk.services.cloudwatch.model.Datapoint)2 DescribeStreamSummaryRequest (software.amazon.awssdk.services.kinesis.model.DescribeStreamSummaryRequest)2 ListShardsRequest (software.amazon.awssdk.services.kinesis.model.ListShardsRequest)2 ListShardsResponse (software.amazon.awssdk.services.kinesis.model.ListShardsResponse)2 URI (java.net.URI)1 UTF_8 (java.nio.charset.StandardCharsets.UTF_8)1 List (java.util.List)1 BiFunction (java.util.function.BiFunction)1 Function (java.util.function.Function)1 Function.identity (java.util.function.Function.identity)1 IntFunction (java.util.function.IntFunction)1 Collectors.toList (java.util.stream.Collectors.toList)1 IntStream.range (java.util.stream.IntStream.range)1 PipelineExecutionException (org.apache.beam.sdk.Pipeline.PipelineExecutionException)1