From 17c61021ee368cbbea41bc56bee5b89ad8fdf1c7 Mon Sep 17 00:00:00 2001 From: Bohan Zhang Date: Wed, 17 Jul 2024 14:12:19 +0800 Subject: [PATCH] fix: Kinesis: NextToken and StreamName cannot be provided together (#17687) Signed-off-by: tabVersion Signed-off-by: tabVersion Co-authored-by: tabVersion --- .../src/source/kinesis/enumerator/client.rs | 30 +++++++++++++------ 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/src/connector/src/source/kinesis/enumerator/client.rs b/src/connector/src/source/kinesis/enumerator/client.rs index 840def08f685..423516fa5bd4 100644 --- a/src/connector/src/source/kinesis/enumerator/client.rs +++ b/src/connector/src/source/kinesis/enumerator/client.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use anyhow::Context as _; +use anyhow::anyhow; use async_trait::async_trait; use aws_sdk_kinesis::types::Shard; use aws_sdk_kinesis::Client as kinesis_client; @@ -52,14 +52,26 @@ impl SplitEnumerator for KinesisSplitEnumerator { let mut shard_collect: Vec = Vec::new(); loop { - let list_shard_output = self - .client - .list_shards() - .set_next_token(next_token) - .stream_name(&self.stream_name) - .send() - .await - .context("failed to list kinesis shards")?; + let mut req = self.client.list_shards(); + if let Some(token) = next_token.take() { + req = req.next_token(token); + } else { + req = req.stream_name(&self.stream_name); + } + + let list_shard_output = match req.send().await { + Ok(output) => output, + Err(e) => { + if let Some(e_inner) = e.as_service_error() + && e_inner.is_expired_next_token_exception() + { + tracing::info!("Kinesis ListShard token expired, retrying..."); + next_token = None; + continue; + } + return Err(anyhow!(e).context("failed to list kinesis shards").into()); + } + }; match list_shard_output.shards { Some(shard) => shard_collect.extend(shard), None => bail!("no shards in stream {}", &self.stream_name),