[FLINK-36637] Fix ListShards for Kinesis streams with large number of…

… shards
apache · Oct 31, 2024 · 6763f0a · 6763f0a
1 parent b922a2d
commit 6763f0a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/...ams/src/main/java/org/apache/flink/connector/kinesis/source/proxy/KinesisStreamProxy.java b/...ams/src/main/java/org/apache/flink/connector/kinesis/source/proxy/KinesisStreamProxy.java
@@ -76,7 +76,7 @@ public List<Shard> listShards(String streamArn, ListShardsStartingPosition start
             listShardsResponse =
                     kinesisClient.listShards(
                             ListShardsRequest.builder()
-                                    .streamARN(streamArn)
+                                    .streamARN(nextToken == null ? streamArn : null)
                                     .shardFilter(
                                             nextToken == null
                                                     ? startingPosition.getShardFilter()

diff --git a/...src/test/java/org/apache/flink/connector/kinesis/source/proxy/KinesisStreamProxyTest.java b/...src/test/java/org/apache/flink/connector/kinesis/source/proxy/KinesisStreamProxyTest.java
@@ -151,14 +151,14 @@ void testListShardsMultipleCalls() {
                                 ListShardItem.builder()
                                         .validation(
                                                 getListShardRequestValidation(
-                                                        STREAM_ARN, null, "next-token-1"))
+                                                        null, null, "next-token-1"))
                                         .shards(expectedShards.subList(1, 2))
                                         .nextToken("next-token-2")
                                         .build(),
                                 ListShardItem.builder()
                                         .validation(
                                                 getListShardRequestValidation(
-                                                        STREAM_ARN, null, "next-token-2"))
+                                                        null, null, "next-token-2"))
                                         .shards(expectedShards.subList(2, 4))
                                         .nextToken(null)
                                         .build())
@@ -461,6 +461,9 @@ private Consumer<ListShardsRequest> getListShardRequestValidation(
                             .nextToken(nextToken)
                             .build();
             assertThat(req).isEqualTo(expectedReq);
+            if (nextToken != null) {
+                assertThat(streamArn).isNull();
+            }
         };
     }