You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pulsar.apache.org by ba...@apache.org on 2025/07/11 08:51:33 UTC
(pulsar) 01/02: [fix][io] Make record properties configurable for kinesis source (#24495)
This is an automated email from the ASF dual-hosted git repository.
baodi pushed a commit to branch branch-4.0
in repository https://gitbox.apache.org/repos/asf/pulsar.git
commit 26a0755d3d107f30e4793ef665e47b3258f45d08
Author: Baodi Shi <ba...@apache.org>
AuthorDate: Thu Jul 10 12:12:49 2025 +0800
[fix][io] Make record properties configurable for kinesis source (#24495)
(cherry picked from commit e0efcbb379a429119f37bcbb307fd44f67e0156c)
---
.../apache/pulsar/io/kinesis/KinesisRecord.java | 36 +++++--
.../pulsar/io/kinesis/KinesisRecordProcessor.java | 9 +-
.../pulsar/io/kinesis/KinesisSourceConfig.java | 23 +++++
.../pulsar/io/kinesis/KinesisRecordTest.java | 110 +++++++++++++++++++++
.../io/kinesis/KinesisSourceConfigTests.java | 56 ++++++++++-
5 files changed, 222 insertions(+), 12 deletions(-)
diff --git a/pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecord.java b/pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecord.java
index 754b1e070f4..187e5ae9985 100644
--- a/pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecord.java
+++ b/pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecord.java
@@ -24,31 +24,49 @@ import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
+import java.util.Set;
import org.apache.pulsar.functions.api.Record;
import software.amazon.awssdk.services.kinesis.model.EncryptionType;
import software.amazon.kinesis.retrieval.KinesisClientRecord;
public class KinesisRecord implements Record<byte[]> {
- public static final String ARRIVAL_TIMESTAMP = "";
- public static final String ENCRYPTION_TYPE = "";
- public static final String PARTITION_KEY = "";
- public static final String SEQUENCE_NUMBER = "";
+ public static final String ARRIVAL_TIMESTAMP = "kinesis.arrival.timestamp";
+ public static final String ENCRYPTION_TYPE = "kinesis.encryption.type";
+ public static final String PARTITION_KEY = "kinesis.partition.key";
+ public static final String SEQUENCE_NUMBER = "kinesis.sequence.number";
+ public static final String SHARD_ID = "kinesis.shard.id";
+ public static final String MILLIS_BEHIND_LATEST = "kinesis.millis.behind.latest";
private static final CharsetDecoder decoder = StandardCharsets.UTF_8.newDecoder();
private final Optional<String> key;
private final byte[] value;
private final HashMap<String, String> userProperties = new HashMap<>();
- public KinesisRecord(KinesisClientRecord record) {
+ public KinesisRecord(KinesisClientRecord record, String shardId, long millisBehindLatest,
+ Set<String> propertiesToInclude) {
this.key = Optional.of(record.partitionKey());
// encryption type can (annoyingly) be null, so we default to NONE
EncryptionType encType = EncryptionType.NONE;
if (record.encryptionType() != null) {
encType = record.encryptionType();
}
- setProperty(ARRIVAL_TIMESTAMP, record.approximateArrivalTimestamp().toString());
- setProperty(ENCRYPTION_TYPE, encType.toString());
- setProperty(PARTITION_KEY, record.partitionKey());
- setProperty(SEQUENCE_NUMBER, record.sequenceNumber());
+ if (propertiesToInclude.contains(ARRIVAL_TIMESTAMP)) {
+ setProperty(ARRIVAL_TIMESTAMP, record.approximateArrivalTimestamp().toString());
+ }
+ if (propertiesToInclude.contains(ENCRYPTION_TYPE)) {
+ setProperty(ENCRYPTION_TYPE, encType.toString());
+ }
+ if (propertiesToInclude.contains(PARTITION_KEY)) {
+ setProperty(PARTITION_KEY, record.partitionKey());
+ }
+ if (propertiesToInclude.contains(SEQUENCE_NUMBER)) {
+ setProperty(SEQUENCE_NUMBER, record.sequenceNumber());
+ }
+ if (propertiesToInclude.contains(SHARD_ID)) {
+ setProperty(SHARD_ID, shardId);
+ }
+ if (propertiesToInclude.contains(MILLIS_BEHIND_LATEST)) {
+ setProperty(MILLIS_BEHIND_LATEST, String.valueOf(millisBehindLatest));
+ }
if (encType == EncryptionType.NONE) {
String s = null;
diff --git a/pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecordProcessor.java b/pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecordProcessor.java
index 60dd263eab4..b51ac544be2 100644
--- a/pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecordProcessor.java
+++ b/pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisRecordProcessor.java
@@ -18,6 +18,7 @@
*/
package org.apache.pulsar.io.kinesis;
+import java.util.Set;
import java.util.concurrent.LinkedBlockingQueue;
import lombok.extern.slf4j.Slf4j;
import software.amazon.kinesis.exceptions.InvalidStateException;
@@ -43,11 +44,13 @@ public class KinesisRecordProcessor implements ShardRecordProcessor {
private final LinkedBlockingQueue<KinesisRecord> queue;
private long nextCheckpointTimeInNanos;
private String kinesisShardId;
+ private final Set<String> propertiesToInclude;
public KinesisRecordProcessor(LinkedBlockingQueue<KinesisRecord> queue, KinesisSourceConfig config) {
this.queue = queue;
this.checkpointInterval = config.getCheckpointInterval();
this.numRetries = config.getNumRetries();
this.backoffTime = config.getBackoffTime();
+ this.propertiesToInclude = config.getPropertiesToInclude();
}
private void checkpoint(RecordProcessorCheckpointer checkpointer) {
@@ -82,16 +85,20 @@ public class KinesisRecordProcessor implements ShardRecordProcessor {
@Override
public void initialize(InitializationInput initializationInput) {
kinesisShardId = initializationInput.shardId();
+ log.info("Initializing KinesisRecordProcessor for shard {}. Config: checkpointInterval={}ms, numRetries={}, "
+ + "backoffTime={}ms, propertiesToInclude={}",
+ kinesisShardId, checkpointInterval, numRetries, backoffTime, propertiesToInclude);
}
@Override
public void processRecords(ProcessRecordsInput processRecordsInput) {
log.info("Processing " + processRecordsInput.records().size() + " records from " + kinesisShardId);
+ long millisBehindLatest = processRecordsInput.millisBehindLatest();
for (KinesisClientRecord record : processRecordsInput.records()) {
try {
- queue.put(new KinesisRecord(record));
+ queue.put(new KinesisRecord(record, this.kinesisShardId, millisBehindLatest, propertiesToInclude));
} catch (InterruptedException e) {
log.warn("unable to create KinesisRecord ", e);
}
diff --git a/pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisSourceConfig.java b/pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisSourceConfig.java
index 0dd9bfce9e0..ebcd06bce12 100644
--- a/pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisSourceConfig.java
+++ b/pulsar-io/kinesis/src/main/java/org/apache/pulsar/io/kinesis/KinesisSourceConfig.java
@@ -22,8 +22,12 @@ import static com.google.common.base.Preconditions.checkArgument;
import static org.apache.commons.lang3.StringUtils.isNotBlank;
import java.io.Serializable;
import java.net.URI;
+import java.util.Arrays;
+import java.util.Collections;
import java.util.Date;
import java.util.Map;
+import java.util.Set;
+import java.util.stream.Collectors;
import lombok.Data;
import lombok.EqualsAndHashCode;
import org.apache.pulsar.io.aws.AwsCredentialProviderPlugin;
@@ -130,6 +134,16 @@ public class KinesisSourceConfig extends BaseKinesisConfig implements Serializab
)
private boolean useEnhancedFanOut = true;
+ @FieldDoc(required = false,
+ defaultValue = "kinesis.arrival.timestamp,kinesis.encryption.type,kinesis.partition.key,"
+ + "kinesis.sequence.number",
+ help = "A comma-separated list of Kinesis metadata properties to include in the Pulsar message properties."
+ + " The supported properties are: kinesis.arrival.timestamp, kinesis.encryption.type, "
+ + "kinesis.partition.key, kinesis.sequence.number, kinesis.shard.id, kinesis.millis.behind.latest")
+ private String kinesisRecordProperties = "kinesis.arrival.timestamp,kinesis.encryption.type,"
+ + "kinesis.partition.key,kinesis.sequence.number";
+ private transient Set<String> propertiesToInclude;
+
public static KinesisSourceConfig load(Map<String, Object> config, SourceContext sourceContext) {
KinesisSourceConfig kinesisSourceConfig = IOConfigUtils.loadWithSecrets(config,
KinesisSourceConfig.class, sourceContext);
@@ -143,6 +157,15 @@ public class KinesisSourceConfig extends BaseKinesisConfig implements Serializab
checkArgument((kinesisSourceConfig.getStartAtTime() != null),
"When initialPositionInStream is AT_TIMESTAMP, startAtTime must be specified");
}
+ if (isNotBlank(kinesisSourceConfig.getKinesisRecordProperties())) {
+ Set<String> properties = Arrays.stream(kinesisSourceConfig.getKinesisRecordProperties().split(","))
+ .map(String::trim)
+ .filter(s -> !s.isEmpty())
+ .collect(Collectors.toSet());
+ kinesisSourceConfig.setPropertiesToInclude(properties);
+ } else {
+ kinesisSourceConfig.setPropertiesToInclude(Collections.emptySet());
+ }
return kinesisSourceConfig;
}
diff --git a/pulsar-io/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisRecordTest.java b/pulsar-io/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisRecordTest.java
new file mode 100644
index 00000000000..bf32720a795
--- /dev/null
+++ b/pulsar-io/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisRecordTest.java
@@ -0,0 +1,110 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pulsar.io.kinesis;
+
+import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.time.Instant;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import org.mockito.Mockito;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+import software.amazon.awssdk.services.kinesis.model.EncryptionType;
+import software.amazon.kinesis.retrieval.KinesisClientRecord;
+
+public class KinesisRecordTest {
+
+ private KinesisClientRecord mockRecord;
+ private final String shardId = "shard-001";
+ private final long millisBehindLatest = 12345L;
+ private final String partitionKey = "test-key";
+ private final String sequenceNumber = "seq-123";
+ private final Instant arrivalTimestamp = Instant.now();
+
+ @BeforeMethod
+ public void setup() {
+ mockRecord = Mockito.mock(KinesisClientRecord.class);
+ when(mockRecord.partitionKey()).thenReturn(partitionKey);
+ when(mockRecord.sequenceNumber()).thenReturn(sequenceNumber);
+ when(mockRecord.approximateArrivalTimestamp()).thenReturn(arrivalTimestamp);
+ when(mockRecord.encryptionType()).thenReturn(EncryptionType.NONE);
+ when(mockRecord.data()).thenReturn(ByteBuffer.wrap("test-data".getBytes(StandardCharsets.UTF_8)));
+ }
+
+ @Test
+ public void testAllPropertiesIncluded() {
+ Set<String> propertiesToInclude = new HashSet<>(Arrays.asList(
+ KinesisRecord.ARRIVAL_TIMESTAMP,
+ KinesisRecord.ENCRYPTION_TYPE,
+ KinesisRecord.PARTITION_KEY,
+ KinesisRecord.SEQUENCE_NUMBER,
+ KinesisRecord.SHARD_ID,
+ KinesisRecord.MILLIS_BEHIND_LATEST
+ ));
+
+ KinesisRecord kinesisRecord = new KinesisRecord(mockRecord, shardId, millisBehindLatest, propertiesToInclude);
+ Map<String, String> properties = kinesisRecord.getProperties();
+
+ assertEquals(properties.size(), 6);
+ assertEquals(properties.get(KinesisRecord.SHARD_ID), shardId);
+ assertEquals(properties.get(KinesisRecord.MILLIS_BEHIND_LATEST), String.valueOf(millisBehindLatest));
+ assertEquals(properties.get(KinesisRecord.PARTITION_KEY), partitionKey);
+ assertEquals(properties.get(KinesisRecord.SEQUENCE_NUMBER), sequenceNumber);
+ assertEquals(properties.get(KinesisRecord.ARRIVAL_TIMESTAMP), arrivalTimestamp.toString());
+ assertEquals(properties.get(KinesisRecord.ENCRYPTION_TYPE), EncryptionType.NONE.toString());
+ }
+
+ @Test
+ public void testSomePropertiesIncluded() {
+ Set<String> propertiesToInclude = new HashSet<>(Arrays.asList(
+ KinesisRecord.SHARD_ID,
+ KinesisRecord.SEQUENCE_NUMBER
+ ));
+
+ KinesisRecord kinesisRecord = new KinesisRecord(mockRecord, shardId, millisBehindLatest, propertiesToInclude);
+ Map<String, String> properties = kinesisRecord.getProperties();
+
+ assertEquals(properties.size(), 2);
+ assertTrue(properties.containsKey(KinesisRecord.SHARD_ID));
+ assertTrue(properties.containsKey(KinesisRecord.SEQUENCE_NUMBER));
+
+ assertFalse(properties.containsKey(KinesisRecord.PARTITION_KEY));
+ assertFalse(properties.containsKey(KinesisRecord.ARRIVAL_TIMESTAMP));
+ assertFalse(properties.containsKey(KinesisRecord.ENCRYPTION_TYPE));
+ assertFalse(properties.containsKey(KinesisRecord.MILLIS_BEHIND_LATEST));
+ }
+
+ @Test
+ public void testNoPropertiesIncluded() {
+ Set<String> propertiesToInclude = Collections.emptySet();
+
+ KinesisRecord kinesisRecord = new KinesisRecord(mockRecord, shardId, millisBehindLatest, propertiesToInclude);
+ Map<String, String> properties = kinesisRecord.getProperties();
+
+ assertTrue(properties.isEmpty());
+ }
+}
\ No newline at end of file
diff --git a/pulsar-io/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisSourceConfigTests.java b/pulsar-io/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisSourceConfigTests.java
index 4ba3593b1d9..4f64722efde 100644
--- a/pulsar-io/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisSourceConfigTests.java
+++ b/pulsar-io/kinesis/src/test/java/org/apache/pulsar/io/kinesis/KinesisSourceConfigTests.java
@@ -20,6 +20,7 @@ package org.apache.pulsar.io.kinesis;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
import java.io.IOException;
import java.time.ZoneOffset;
@@ -28,7 +29,7 @@ import java.util.Calendar;
import java.util.Date;
import java.util.HashMap;
import java.util.Map;
-
+import java.util.Set;
import org.apache.pulsar.io.core.SourceContext;
import org.mockito.Mockito;
import org.testng.annotations.Test;
@@ -160,4 +161,55 @@ public class KinesisSourceConfigTests {
"{\"accessKey\":\"myKey\",\"secretKey\":\"my-Secret\"}");
KinesisSourceConfig.load(map, Mockito.mock(SourceContext.class));
}
-}
+
+ @Test
+ public final void propertiesDefaultTest() {
+ Map<String, Object> map = new HashMap<>();
+ map.put("awsRegion", "us-west-1");
+ map.put("awsKinesisStreamName", "my-stream");
+ map.put("awsCredentialPluginParam", "{\"accessKey\":\"myKey\",\"secretKey\":\"my-Secret\"}");
+
+ SourceContext sourceContext = Mockito.mock(SourceContext.class);
+ KinesisSourceConfig config = KinesisSourceConfig.load(map, sourceContext);
+
+ Set<String> properties = config.getPropertiesToInclude();
+ assertEquals(properties.size(), 4);
+ assertTrue(properties.contains("kinesis.sequence.number"));
+ assertTrue(properties.contains("kinesis.arrival.timestamp"));
+ assertTrue(properties.contains("kinesis.encryption.type"));
+ assertTrue(properties.contains("kinesis.partition.key"));
+ }
+
+ @Test
+ public final void propertiesCustomTest() {
+ Map<String, Object> map = new HashMap<>();
+ map.put("awsRegion", "us-west-1");
+ map.put("awsKinesisStreamName", "my-stream");
+ map.put("awsCredentialPluginParam", "{\"accessKey\":\"myKey\",\"secretKey\":\"my-Secret\"}");
+ // Set custom properties, note the extra whitespace to test trim()
+ map.put("kinesisRecordProperties", "kinesis.shard.id, kinesis.partition.key ");
+
+ SourceContext sourceContext = Mockito.mock(SourceContext.class);
+ KinesisSourceConfig config = KinesisSourceConfig.load(map, sourceContext);
+
+ Set<String> properties = config.getPropertiesToInclude();
+ assertEquals(properties.size(), 2);
+ assertTrue(properties.contains("kinesis.shard.id"));
+ assertTrue(properties.contains("kinesis.partition.key"));
+ }
+
+ @Test
+ public final void propertiesEmptyTest() {
+ Map<String, Object> map = new HashMap<>();
+ map.put("awsRegion", "us-west-1");
+ map.put("awsKinesisStreamName", "my-stream");
+ map.put("awsCredentialPluginParam", "{\"accessKey\":\"myKey\",\"secretKey\":\"my-Secret\"}");
+ map.put("kinesisRecordProperties", "");
+
+ SourceContext sourceContext = Mockito.mock(SourceContext.class);
+ KinesisSourceConfig config = KinesisSourceConfig.load(map, sourceContext);
+
+ Set<String> properties = config.getPropertiesToInclude();
+ assertTrue(properties.isEmpty());
+ }
+}
\ No newline at end of file