You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by ch...@apache.org on 2023/06/30 18:29:38 UTC
[beam] branch master updated: Add topicPattern property to KafkaIO.Read to match topics using a regex (#26948)
This is an automated email from the ASF dual-hosted git repository.
chamikara pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new c36f0f13d35 Add topicPattern property to KafkaIO.Read to match topics using a regex (#26948)
c36f0f13d35 is described below
commit c36f0f13d35b00d43dc99392eeaf0e1b12f3c9b5
Author: Steven van Rossum <sj...@google.com>
AuthorDate: Fri Jun 30 20:29:30 2023 +0200
Add topicPattern property to KafkaIO.Read to match topics using a regex (#26948)
* Add topicPattern property to KafkaIO.Read to match topics using a regex
* Add partially matched and unmatched topic pattern test
* Change method signature of withTopicPattern to use String
---
.../java/org/apache/beam/sdk/io/kafka/KafkaIO.java | 71 ++++++++++++---
.../KafkaIOReadImplementationCompatibility.java | 1 +
.../beam/sdk/io/kafka/KafkaUnboundedSource.java | 30 +++++--
.../sdk/io/kafka/WatchForKafkaTopicPartitions.java | 26 ++++--
.../org/apache/beam/sdk/io/kafka/KafkaIOTest.java | 100 +++++++++++++++++++++
.../io/kafka/WatchForKafkaTopicPartitionsTest.java | 48 +++++++++-
6 files changed, 246 insertions(+), 30 deletions(-)
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
index daa3ff23928..3b1ee1163e4 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java
@@ -37,6 +37,7 @@ import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
+import java.util.regex.Pattern;
import java.util.stream.Collectors;
import org.apache.beam.runners.core.construction.PTransformMatchers;
import org.apache.beam.runners.core.construction.ReplacementOutputs;
@@ -350,10 +351,11 @@ import org.slf4j.LoggerFactory;
* href="https://beam.apache.org/blog/splittable-do-fn/">blog post</a> and <a
* href="https://s.apache.org/beam-fn-api">design doc</a>. The major difference from {@link
* KafkaIO.Read} is, {@link ReadSourceDescriptors} doesn't require source descriptions(e.g., {@link
- * KafkaIO.Read#getTopicPartitions()}, {@link KafkaIO.Read#getTopics()}, {@link
- * KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline construction time. Instead, the
- * pipeline can populate these source descriptions during runtime. For example, the pipeline can
- * query Kafka topics from a BigQuery table and read these topics via {@link ReadSourceDescriptors}.
+ * KafkaIO.Read#getTopicPattern()}, {@link KafkaIO.Read#getTopicPartitions()}, {@link
+ * KafkaIO.Read#getTopics()}, {@link KafkaIO.Read#getStartReadTime()}, etc.) during the pipeline
+ * construction time. Instead, the pipeline can populate these source descriptions during runtime.
+ * For example, the pipeline can query Kafka topics from a BigQuery table and read these topics via
+ * {@link ReadSourceDescriptors}.
*
* <h3>Common Kafka Consumer Configurations</h3>
*
@@ -633,6 +635,9 @@ public class KafkaIO {
@Pure
abstract @Nullable List<TopicPartition> getTopicPartitions();
+ @Pure
+ abstract @Nullable Pattern getTopicPattern();
+
@Pure
abstract @Nullable Coder<K> getKeyCoder();
@@ -692,6 +697,8 @@ public class KafkaIO {
abstract Builder<K, V> setTopicPartitions(List<TopicPartition> topicPartitions);
+ abstract Builder<K, V> setTopicPattern(Pattern topicPattern);
+
abstract Builder<K, V> setKeyCoder(Coder<K> keyCoder);
abstract Builder<K, V> setValueCoder(Coder<V> valueCoder);
@@ -922,8 +929,9 @@ public class KafkaIO {
*/
public Read<K, V> withTopics(List<String> topics) {
checkState(
- getTopicPartitions() == null || getTopicPartitions().isEmpty(),
- "Only topics or topicPartitions can be set, not both");
+ (getTopicPartitions() == null || getTopicPartitions().isEmpty())
+ && getTopicPattern() == null,
+ "Only one of topics, topicPartitions or topicPattern can be set");
return toBuilder().setTopics(ImmutableList.copyOf(topics)).build();
}
@@ -936,11 +944,26 @@ public class KafkaIO {
*/
public Read<K, V> withTopicPartitions(List<TopicPartition> topicPartitions) {
checkState(
- getTopics() == null || getTopics().isEmpty(),
- "Only topics or topicPartitions can be set, not both");
+ (getTopics() == null || getTopics().isEmpty()) && getTopicPattern() == null,
+ "Only one of topics, topicPartitions or topicPattern can be set");
return toBuilder().setTopicPartitions(ImmutableList.copyOf(topicPartitions)).build();
}
+ /**
+ * Internally sets a {@link java.util.regex.Pattern} of topics to read from. All the partitions
+ * from each of the matching topics are read.
+ *
+ * <p>See {@link KafkaUnboundedSource#split(int, PipelineOptions)} for description of how the
+ * partitions are distributed among the splits.
+ */
+ public Read<K, V> withTopicPattern(String topicPattern) {
+ checkState(
+ (getTopics() == null || getTopics().isEmpty())
+ && (getTopicPartitions() == null || getTopicPartitions().isEmpty()),
+ "Only one of topics, topicPartitions or topicPattern can be set");
+ return toBuilder().setTopicPattern(Pattern.compile(topicPattern)).build();
+ }
+
/**
* Sets a Kafka {@link Deserializer} to interpret key bytes read from Kafka.
*
@@ -1274,8 +1297,9 @@ public class KafkaIO {
if (!isDynamicRead()) {
checkArgument(
(getTopics() != null && getTopics().size() > 0)
- || (getTopicPartitions() != null && getTopicPartitions().size() > 0),
- "Either withTopic(), withTopics() or withTopicPartitions() is required");
+ || (getTopicPartitions() != null && getTopicPartitions().size() > 0)
+ || getTopicPattern() != null,
+ "Either withTopic(), withTopics(), withTopicPartitions() or withTopicPattern() is required");
} else {
checkArgument(
ExperimentalOptions.hasExperiment(input.getPipeline().getOptions(), "beam_fn_api"),
@@ -1537,6 +1561,7 @@ public class KafkaIO {
kafkaRead.getConsumerConfig(),
kafkaRead.getCheckStopReadingFn(),
topics,
+ kafkaRead.getTopicPattern(),
kafkaRead.getStartReadTime(),
kafkaRead.getStopReadTime()));
} else {
@@ -1561,6 +1586,7 @@ public class KafkaIO {
this.consumerFactoryFn = read.getConsumerFactoryFn();
this.topics = read.getTopics();
this.topicPartitions = read.getTopicPartitions();
+ this.topicPattern = read.getTopicPattern();
this.startReadTime = read.getStartReadTime();
this.stopReadTime = read.getStopReadTime();
}
@@ -1578,15 +1604,30 @@ public class KafkaIO {
@VisibleForTesting final @Nullable List<String> topics;
+ private final @Nullable Pattern topicPattern;
+
@ProcessElement
public void processElement(OutputReceiver<KafkaSourceDescriptor> receiver) {
List<TopicPartition> partitions =
new ArrayList<>(Preconditions.checkStateNotNull(topicPartitions));
if (partitions.isEmpty()) {
try (Consumer<?, ?> consumer = consumerFactoryFn.apply(consumerConfig)) {
- for (String topic : Preconditions.checkStateNotNull(topics)) {
- for (PartitionInfo p : consumer.partitionsFor(topic)) {
- partitions.add(new TopicPartition(p.topic(), p.partition()));
+ List<String> topics = Preconditions.checkStateNotNull(this.topics);
+ if (topics.isEmpty()) {
+ Pattern pattern = Preconditions.checkStateNotNull(topicPattern);
+ for (Map.Entry<String, List<PartitionInfo>> entry :
+ consumer.listTopics().entrySet()) {
+ if (pattern.matcher(entry.getKey()).matches()) {
+ for (PartitionInfo p : entry.getValue()) {
+ partitions.add(new TopicPartition(p.topic(), p.partition()));
+ }
+ }
+ }
+ } else {
+ for (String topic : topics) {
+ for (PartitionInfo p : consumer.partitionsFor(topic)) {
+ partitions.add(new TopicPartition(p.topic(), p.partition()));
+ }
}
}
}
@@ -1634,12 +1675,16 @@ public class KafkaIO {
super.populateDisplayData(builder);
List<String> topics = Preconditions.checkStateNotNull(getTopics());
List<TopicPartition> topicPartitions = Preconditions.checkStateNotNull(getTopicPartitions());
+ Pattern topicPattern = getTopicPattern();
if (topics.size() > 0) {
builder.add(DisplayData.item("topics", Joiner.on(",").join(topics)).withLabel("Topic/s"));
} else if (topicPartitions.size() > 0) {
builder.add(
DisplayData.item("topicPartitions", Joiner.on(",").join(topicPartitions))
.withLabel("Topic Partition/s"));
+ } else if (topicPattern != null) {
+ builder.add(
+ DisplayData.item("topicPattern", topicPattern.pattern()).withLabel("Topic Pattern"));
}
Set<String> disallowedConsumerPropertiesKeys =
KafkaIOUtils.DISALLOWED_CONSUMER_PROPERTIES.keySet();
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibility.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibility.java
index 1db28d2a3e6..1c974f3bcc1 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibility.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIOReadImplementationCompatibility.java
@@ -79,6 +79,7 @@ class KafkaIOReadImplementationCompatibility {
CONSUMER_CONFIG,
TOPICS,
TOPIC_PARTITIONS,
+ TOPIC_PATTERN,
KEY_CODER,
VALUE_CODER,
CONSUMER_FACTORY_FN,
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedSource.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedSource.java
index 4af13bbf474..bfb375b5eec 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedSource.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedSource.java
@@ -24,6 +24,8 @@ import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
+import java.util.Map;
+import java.util.regex.Pattern;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.extensions.avro.coders.AvroCoder;
import org.apache.beam.sdk.io.UnboundedSource;
@@ -65,14 +67,26 @@ class KafkaUnboundedSource<K, V> extends UnboundedSource<KafkaRecord<K, V>, Kafk
if (partitions.isEmpty()) {
try (Consumer<?, ?> consumer = spec.getConsumerFactoryFn().apply(spec.getConsumerConfig())) {
- for (String topic : Preconditions.checkStateNotNull(spec.getTopics())) {
- List<PartitionInfo> partitionInfoList = consumer.partitionsFor(topic);
- checkState(
- partitionInfoList != null,
- "Could not find any partitions info. Please check Kafka configuration and make sure "
- + "that provided topics exist.");
- for (PartitionInfo p : partitionInfoList) {
- partitions.add(new TopicPartition(p.topic(), p.partition()));
+ List<String> topics = Preconditions.checkStateNotNull(spec.getTopics());
+ if (topics.isEmpty()) {
+ Pattern pattern = Preconditions.checkStateNotNull(spec.getTopicPattern());
+ for (Map.Entry<String, List<PartitionInfo>> entry : consumer.listTopics().entrySet()) {
+ if (pattern.matcher(entry.getKey()).matches()) {
+ for (PartitionInfo p : entry.getValue()) {
+ partitions.add(new TopicPartition(p.topic(), p.partition()));
+ }
+ }
+ }
+ } else {
+ for (String topic : topics) {
+ List<PartitionInfo> partitionInfoList = consumer.partitionsFor(topic);
+ checkState(
+ partitionInfoList != null,
+ "Could not find any partitions info. Please check Kafka configuration and make sure "
+ + "that provided topics exist.");
+ for (PartitionInfo p : partitionInfoList) {
+ partitions.add(new TopicPartition(p.topic(), p.partition()));
+ }
}
}
}
diff --git a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/WatchForKafkaTopicPartitions.java b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/WatchForKafkaTopicPartitions.java
index 21fec1bd27b..6a25257bdac 100644
--- a/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/WatchForKafkaTopicPartitions.java
+++ b/sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/WatchForKafkaTopicPartitions.java
@@ -24,6 +24,7 @@ import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
+import java.util.regex.Pattern;
import org.apache.beam.sdk.metrics.Counter;
import org.apache.beam.sdk.metrics.Metrics;
import org.apache.beam.sdk.transforms.DoFn;
@@ -64,6 +65,7 @@ class WatchForKafkaTopicPartitions extends PTransform<PBegin, PCollection<KafkaS
private final Map<String, Object> kafkaConsumerConfig;
private final @Nullable SerializableFunction<TopicPartition, Boolean> checkStopReadingFn;
private final Set<String> topics;
+ private final @Nullable Pattern topicPattern;
private final @Nullable Instant startReadTime;
private final @Nullable Instant stopReadTime;
@@ -73,6 +75,7 @@ class WatchForKafkaTopicPartitions extends PTransform<PBegin, PCollection<KafkaS
Map<String, Object> kafkaConsumerConfig,
@Nullable SerializableFunction<TopicPartition, Boolean> checkStopReadingFn,
Set<String> topics,
+ @Nullable Pattern topicPattern,
@Nullable Instant startReadTime,
@Nullable Instant stopReadTime) {
this.checkDuration = firstNonNull(checkDuration, DEFAULT_CHECK_DURATION);
@@ -80,6 +83,7 @@ class WatchForKafkaTopicPartitions extends PTransform<PBegin, PCollection<KafkaS
this.kafkaConsumerConfig = kafkaConsumerConfig;
this.checkStopReadingFn = checkStopReadingFn;
this.topics = topics;
+ this.topicPattern = topicPattern;
this.startReadTime = startReadTime;
this.stopReadTime = stopReadTime;
}
@@ -91,7 +95,8 @@ class WatchForKafkaTopicPartitions extends PTransform<PBegin, PCollection<KafkaS
.apply(
"Match new TopicPartitions",
Watch.growthOf(
- new WatchPartitionFn(kafkaConsumerFactoryFn, kafkaConsumerConfig, topics))
+ new WatchPartitionFn(
+ kafkaConsumerFactoryFn, kafkaConsumerConfig, topics, topicPattern))
.withPollInterval(checkDuration))
.apply(ParDo.of(new ConvertToDescriptor(checkStopReadingFn, startReadTime, stopReadTime)));
}
@@ -134,14 +139,17 @@ class WatchForKafkaTopicPartitions extends PTransform<PBegin, PCollection<KafkaS
kafkaConsumerFactoryFn;
private final Map<String, Object> kafkaConsumerConfig;
private final Set<String> topics;
+ private final @Nullable Pattern topicPattern;
private WatchPartitionFn(
SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> kafkaConsumerFactoryFn,
Map<String, Object> kafkaConsumerConfig,
- Set<String> topics) {
+ Set<String> topics,
+ @Nullable Pattern topicPattern) {
this.kafkaConsumerFactoryFn = kafkaConsumerFactoryFn;
this.kafkaConsumerConfig = kafkaConsumerConfig;
this.topics = topics;
+ this.topicPattern = topicPattern;
}
@Override
@@ -149,7 +157,9 @@ class WatchForKafkaTopicPartitions extends PTransform<PBegin, PCollection<KafkaS
throws Exception {
Instant now = Instant.now();
return Watch.Growth.PollResult.incomplete(
- now, getAllTopicPartitions(kafkaConsumerFactoryFn, kafkaConsumerConfig, topics))
+ now,
+ getAllTopicPartitions(
+ kafkaConsumerFactoryFn, kafkaConsumerConfig, topics, topicPattern))
.withWatermark(now);
}
}
@@ -158,7 +168,8 @@ class WatchForKafkaTopicPartitions extends PTransform<PBegin, PCollection<KafkaS
static List<TopicPartition> getAllTopicPartitions(
SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>> kafkaConsumerFactoryFn,
Map<String, Object> kafkaConsumerConfig,
- Set<String> topics) {
+ Set<String> topics,
+ @Nullable Pattern topicPattern) {
List<TopicPartition> current = new ArrayList<>();
try (Consumer<byte[], byte[]> kafkaConsumer =
kafkaConsumerFactoryFn.apply(kafkaConsumerConfig)) {
@@ -168,12 +179,13 @@ class WatchForKafkaTopicPartitions extends PTransform<PBegin, PCollection<KafkaS
current.add(new TopicPartition(topic, partition.partition()));
}
}
-
} else {
for (Map.Entry<String, List<PartitionInfo>> topicInfo :
kafkaConsumer.listTopics().entrySet()) {
- for (PartitionInfo partition : topicInfo.getValue()) {
- current.add(new TopicPartition(topicInfo.getKey(), partition.partition()));
+ if (topicPattern == null || topicPattern.matcher(topicInfo.getKey()).matches()) {
+ for (PartitionInfo partition : topicInfo.getValue()) {
+ current.add(new TopicPartition(partition.topic(), partition.partition()));
+ }
}
}
}
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
index d170b2a237b..577dbaf9827 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaIOTest.java
@@ -723,6 +723,87 @@ public class KafkaIOTest {
p.run();
}
+ @Test
+ public void testUnboundedSourceWithPattern() {
+ int numElements = 1000;
+
+ List<String> topics =
+ ImmutableList.of(
+ "best", "gest", "hest", "jest", "lest", "nest", "pest", "rest", "test", "vest", "west",
+ "zest");
+
+ KafkaIO.Read<byte[], Long> reader =
+ KafkaIO.<byte[], Long>read()
+ .withBootstrapServers("none")
+ .withTopicPattern("[a-z]est")
+ .withConsumerFactoryFn(
+ new ConsumerFactoryFn(topics, 10, numElements, OffsetResetStrategy.EARLIEST))
+ .withKeyDeserializer(ByteArrayDeserializer.class)
+ .withValueDeserializer(LongDeserializer.class)
+ .withMaxNumRecords(numElements);
+
+ PCollection<Long> input = p.apply(reader.withoutMetadata()).apply(Values.create());
+
+ addCountingAsserts(input, numElements);
+ p.run();
+ }
+
+ @Test
+ public void testUnboundedSourceWithPartiallyMatchedPattern() {
+ int numElements = 1000;
+ long numMatchedElements = numElements / 2; // Expected elements if split across 2 topics
+
+ List<String> topics = ImmutableList.of("test", "Test");
+
+ KafkaIO.Read<byte[], Long> reader =
+ KafkaIO.<byte[], Long>read()
+ .withBootstrapServers("none")
+ .withTopicPattern("[a-z]est")
+ .withConsumerFactoryFn(
+ new ConsumerFactoryFn(topics, 1, numElements, OffsetResetStrategy.EARLIEST))
+ .withKeyDeserializer(ByteArrayDeserializer.class)
+ .withValueDeserializer(LongDeserializer.class)
+ .withMaxNumRecords(numMatchedElements);
+
+ PCollection<Long> input = p.apply(reader.withoutMetadata()).apply(Values.create());
+
+ // With 1 partition per topic element to partition allocation alternates between test and Test,
+ // producing even elements for test and odd elements for Test.
+ // The pattern only matches test, so we expect even elements.
+ PAssert.that(input).satisfies(new AssertMultipleOf(2));
+
+ PAssert.thatSingleton(input.apply("Count", Count.globally())).isEqualTo(numMatchedElements);
+
+ p.run();
+ }
+
+ @Test
+ public void testUnboundedSourceWithUnmatchedPattern() {
+ // Expect an exception when provided pattern doesn't match any Kafka topics.
+ thrown.expect(PipelineExecutionException.class);
+ thrown.expectCause(instanceOf(IllegalStateException.class));
+ thrown.expectMessage(
+ "Could not find any partitions. Please check Kafka configuration and topic names");
+
+ int numElements = 1000;
+
+ List<String> topics = ImmutableList.of("chest", "crest", "egest", "guest", "quest", "wrest");
+
+ KafkaIO.Read<byte[], Long> reader =
+ KafkaIO.<byte[], Long>read()
+ .withBootstrapServers("none")
+ .withTopicPattern("[a-z]est")
+ .withConsumerFactoryFn(
+ new ConsumerFactoryFn(topics, 10, numElements, OffsetResetStrategy.EARLIEST))
+ .withKeyDeserializer(ByteArrayDeserializer.class)
+ .withValueDeserializer(LongDeserializer.class)
+ .withMaxNumRecords(numElements);
+
+ p.apply(reader.withoutMetadata()).apply(Values.create());
+
+ p.run();
+ }
+
@Test
public void testUnboundedSourceWithWrongTopic() {
// Expect an exception when provided Kafka topic doesn't exist.
@@ -1829,6 +1910,25 @@ public class KafkaIOTest {
assertThat(displayData, hasDisplayItem("receive.buffer.bytes", 524288));
}
+ @Test
+ public void testSourceWithPatternDisplayData() {
+ KafkaIO.Read<byte[], byte[]> read =
+ KafkaIO.readBytes()
+ .withBootstrapServers("myServer1:9092,myServer2:9092")
+ .withTopicPattern("[a-z]est")
+ .withConsumerFactoryFn(
+ new ConsumerFactoryFn(
+ Lists.newArrayList("test"), 10, 10, OffsetResetStrategy.EARLIEST));
+
+ DisplayData displayData = DisplayData.from(read);
+
+ assertThat(displayData, hasDisplayItem("topicPattern", "[a-z]est"));
+ assertThat(displayData, hasDisplayItem("enable.auto.commit", false));
+ assertThat(displayData, hasDisplayItem("bootstrap.servers", "myServer1:9092,myServer2:9092"));
+ assertThat(displayData, hasDisplayItem("auto.offset.reset", "latest"));
+ assertThat(displayData, hasDisplayItem("receive.buffer.bytes", 524288));
+ }
+
@Test
public void testSinkDisplayData() {
try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) {
diff --git a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/WatchForKafkaTopicPartitionsTest.java b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/WatchForKafkaTopicPartitionsTest.java
index 78ec5772819..25ca17ba91a 100644
--- a/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/WatchForKafkaTopicPartitionsTest.java
+++ b/sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/WatchForKafkaTopicPartitionsTest.java
@@ -23,6 +23,7 @@ import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import java.util.Set;
+import java.util.regex.Pattern;
import org.apache.beam.sdk.io.kafka.KafkaMocks.PartitionGrowthMockConsumer;
import org.apache.beam.sdk.testing.PAssert;
import org.apache.beam.sdk.testing.SerializableMatcher;
@@ -77,7 +78,8 @@ public class WatchForKafkaTopicPartitionsTest {
new TopicPartition("topic1", 1),
new TopicPartition("topic2", 0),
new TopicPartition("topic2", 1)),
- WatchForKafkaTopicPartitions.getAllTopicPartitions((input) -> mockConsumer, null, null));
+ WatchForKafkaTopicPartitions.getAllTopicPartitions(
+ (input) -> mockConsumer, null, null, null));
}
@Test
@@ -103,7 +105,47 @@ public class WatchForKafkaTopicPartitionsTest {
new TopicPartition("topic2", 0),
new TopicPartition("topic2", 1)),
WatchForKafkaTopicPartitions.getAllTopicPartitions(
- (input) -> mockConsumer, null, givenTopics));
+ (input) -> mockConsumer, null, givenTopics, null));
+ }
+
+ @Test
+ public void testGetAllTopicPartitionsWithGivenPattern() throws Exception {
+ Consumer<byte[], byte[]> mockConsumer = Mockito.mock(Consumer.class);
+ when(mockConsumer.listTopics())
+ .thenReturn(
+ ImmutableMap.of(
+ "topic1",
+ ImmutableList.of(
+ new PartitionInfo("topic1", 0, null, null, null),
+ new PartitionInfo("topic1", 1, null, null, null)),
+ "topic2",
+ ImmutableList.of(
+ new PartitionInfo("topic2", 0, null, null, null),
+ new PartitionInfo("topic2", 1, null, null, null)),
+ "topicA",
+ ImmutableList.of(
+ new PartitionInfo("topicA", 0, null, null, null),
+ new PartitionInfo("topicA", 1, null, null, null)),
+ "topicB",
+ ImmutableList.of(
+ new PartitionInfo("topicB", 0, null, null, null),
+ new PartitionInfo("topicB", 1, null, null, null))));
+ assertEquals(
+ ImmutableList.of(
+ new TopicPartition("topic1", 0),
+ new TopicPartition("topic1", 1),
+ new TopicPartition("topic2", 0),
+ new TopicPartition("topic2", 1)),
+ WatchForKafkaTopicPartitions.getAllTopicPartitions(
+ (input) -> mockConsumer, null, null, Pattern.compile("topic[0-9]")));
+ assertEquals(
+ ImmutableList.of(
+ new TopicPartition("topicA", 0),
+ new TopicPartition("topicA", 1),
+ new TopicPartition("topicB", 0),
+ new TopicPartition("topicB", 1)),
+ WatchForKafkaTopicPartitions.getAllTopicPartitions(
+ (input) -> mockConsumer, null, null, Pattern.compile("topic[A-Z]")));
}
@Test
@@ -120,6 +162,7 @@ public class WatchForKafkaTopicPartitionsTest {
null,
givenTopics,
null,
+ null,
null);
PCollection<KafkaSourceDescriptor> descriptors = p.apply(watchForKafkaTopicPartitions);
@@ -145,6 +188,7 @@ public class WatchForKafkaTopicPartitionsTest {
null,
givenTopics,
null,
+ null,
null);
PCollection<KafkaSourceDescriptor> descriptors = p.apply(watchForKafkaTopicPartitions);