KafkaIO { public static <K, V> Read<K, V> read() { return new AutoValue_KafkaIO_Read.Builder<K, V>() .setTopics(new ArrayList<>()) .setTopicPartitions(new ArrayList<>()) .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN) .setConsumerConfig(KafkaIOUtils.DEFAULT_CONSUMER_PROPERTIES) .setMaxNumRecords(Long.MAX_VALUE) .setCommitOffsetsInFinalizeEnabled(false) .setTimestampPolicyFactory(TimestampPolicyFactory.withProcessingTime()) .build(); } private  KafkaIO(); static Read<byte[], byte[]> readBytes(); static Read<K, V> read(); static ReadSourceDescriptors<K, V> readSourceDescriptors(); static Write<K, V> write(); static WriteRecords<K, V> writeRecords();  }
@Test public void testReadAvroGenericRecordsWithConfluentSchemaRegistry() { int numElements = 100; String topic = "my_topic"; String schemaRegistryUrl = "mock: String keySchemaSubject = topic + "-key"; String valueSchemaSubject = topic + "-value"; List<KV<GenericRecord, GenericRecord>> inputs = new ArrayList<>(); for (int i = 0; i < numElements; i++) { inputs.add( KV.of( new AvroGeneratedUser("KeyName" + i, i, "color" + i), new AvroGeneratedUser("ValueName" + i, i, "color" + i))); } KafkaIO.Read<GenericRecord, GenericRecord> reader = KafkaIO.<GenericRecord, GenericRecord>read() .withBootstrapServers("localhost:9092") .withTopic(topic) .withKeyDeserializer( mockDeserializerProvider(schemaRegistryUrl, keySchemaSubject, null)) .withValueDeserializer( mockDeserializerProvider(schemaRegistryUrl, valueSchemaSubject, null)) .withConsumerFactoryFn( new ConsumerFactoryFn( ImmutableList.of(topic), 1, numElements, OffsetResetStrategy.EARLIEST, new KeyAvroSerializableFunction(topic, schemaRegistryUrl), new ValueAvroSerializableFunction(topic, schemaRegistryUrl))) .withMaxNumRecords(numElements); PCollection<KV<GenericRecord, GenericRecord>> input = p.apply(reader.withoutMetadata()); PAssert.that(input).containsInAnyOrder(inputs); p.run(); }
@Test public void testReadAvroSpecificRecordsWithConfluentSchemaRegistry() { int numElements = 100; String topic = "my_topic"; String schemaRegistryUrl = "mock: String valueSchemaSubject = topic + "-value"; List<KV<Integer, AvroGeneratedUser>> inputs = new ArrayList<>(); for (int i = 0; i < numElements; i++) { inputs.add(KV.of(i, new AvroGeneratedUser("ValueName" + i, i, "color" + i))); } KafkaIO.Read<Integer, AvroGeneratedUser> reader = KafkaIO.<Integer, AvroGeneratedUser>read() .withBootstrapServers("localhost:9092") .withTopic(topic) .withKeyDeserializer(IntegerDeserializer.class) .withValueDeserializer( mockDeserializerProvider(schemaRegistryUrl, valueSchemaSubject, null)) .withConsumerFactoryFn( new ConsumerFactoryFn( ImmutableList.of(topic), 1, numElements, OffsetResetStrategy.EARLIEST, i -> ByteBuffer.wrap(new byte[4]).putInt(i).array(), new ValueAvroSerializableFunction(topic, schemaRegistryUrl))) .withMaxNumRecords(numElements); PCollection<KV<Integer, AvroGeneratedUser>> input = p.apply(reader.withoutMetadata()); PAssert.that(input).containsInAnyOrder(inputs); p.run(); }
@Test public void testDeserializationWithHeaders() throws Exception { int numElements = 1000; String topic = "my_topic"; KafkaIO.Read<Integer, Long> reader = KafkaIO.<Integer, Long>read() .withBootstrapServers("none") .withTopic("my_topic") .withConsumerFactoryFn( new ConsumerFactoryFn( ImmutableList.of(topic), 10, numElements, OffsetResetStrategy.EARLIEST)) .withMaxNumRecords(numElements) .withKeyDeserializerAndCoder( KafkaIOTest.IntegerDeserializerWithHeadersAssertor.class, BigEndianIntegerCoder.of()) .withValueDeserializerAndCoder( KafkaIOTest.LongDeserializerWithHeadersAssertor.class, BigEndianLongCoder.of()); PCollection<Long> input = p.apply(reader.withoutMetadata()).apply(Values.create()); addCountingAsserts(input, numElements); p.run(); }
@Test public void testUnreachableKafkaBrokers() { thrown.expect(Exception.class); thrown.expectMessage("Reader-0: Timeout while initializing partition 'test-0'"); int numElements = 1000; PCollection<Long> input = p.apply( KafkaIO.<Integer, Long>read() .withBootstrapServers("8.8.8.8:9092") .withTopicPartitions(ImmutableList.of(new TopicPartition("test", 0))) .withKeyDeserializer(IntegerDeserializer.class) .withValueDeserializer(LongDeserializer.class) .withConsumerConfigUpdates( ImmutableMap.of( ConsumerConfig.HEARTBEAT_INTERVAL_MS_CONFIG, 5, ConsumerConfig.SESSION_TIMEOUT_MS_CONFIG, 8, ConsumerConfig.FETCH_MAX_WAIT_MS_CONFIG, 8, "default.api.timeout.ms", 10)) .withMaxNumRecords(10) .withoutMetadata()) .apply(Values.create()); addCountingAsserts(input, numElements); p.run(); }
@Test public void testResolveDefaultApiTimeout() { final String defaultApiTimeoutConfig = "default.api.timeout.ms"; assertEquals( Duration.millis(20), KafkaUnboundedReader.resolveDefaultApiTimeout( KafkaIO.<Integer, Long>read() .withConsumerConfigUpdates( ImmutableMap.of( defaultApiTimeoutConfig, 20, ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, 30)))); assertEquals( Duration.millis(2 * 30), KafkaUnboundedReader.resolveDefaultApiTimeout( KafkaIO.<Integer, Long>read() .withConsumerConfigUpdates( ImmutableMap.of(ConsumerConfig.REQUEST_TIMEOUT_MS_CONFIG, 30)))); assertEquals( Duration.millis(60 * 1000), KafkaUnboundedReader.resolveDefaultApiTimeout(KafkaIO.<Integer, Long>read())); }
@Test public void testUnboundedSourceWithSingleTopic() { int numElements = 1000; String topic = "my_topic"; KafkaIO.Read<Integer, Long> reader = KafkaIO.<Integer, Long>read() .withBootstrapServers("none") .withTopic("my_topic") .withConsumerFactoryFn( new ConsumerFactoryFn( ImmutableList.of(topic), 10, numElements, OffsetResetStrategy.EARLIEST)) .withMaxNumRecords(numElements) .withKeyDeserializer(IntegerDeserializer.class) .withValueDeserializer(LongDeserializer.class); PCollection<Long> input = p.apply(reader.withoutMetadata()).apply(Values.create()); addCountingAsserts(input, numElements); p.run(); }
@Test public void testUnboundedSourceWithExplicitPartitions() { int numElements = 1000; List<String> topics = ImmutableList.of("test"); KafkaIO.Read<byte[], Long> reader = KafkaIO.<byte[], Long>read() .withBootstrapServers("none") .withTopicPartitions(ImmutableList.of(new TopicPartition("test", 5))) .withConsumerFactoryFn( new ConsumerFactoryFn( topics, 10, numElements, OffsetResetStrategy.EARLIEST)) .withKeyDeserializer(ByteArrayDeserializer.class) .withValueDeserializer(LongDeserializer.class) .withMaxNumRecords(numElements / 10); PCollection<Long> input = p.apply(reader.withoutMetadata()).apply(Values.create()); PAssert.that(input).satisfies(new AssertMultipleOf(5)); PAssert.thatSingleton(input.apply(Count.globally())).isEqualTo(numElements / 10L); p.run(); }
@Test public void testUnboundedSourceWithExceptionInKafkaFetch() { thrown.expectCause(isA(IOException.class)); thrown.expectCause(hasMessage(containsString("Exception while reading from Kafka"))); thrown.expectCause(hasCause(isA(KafkaException.class))); thrown.expectCause(hasCause(hasMessage(containsString("Injected error in consumer.poll()")))); int numElements = 1000; String topic = "my_topic"; KafkaIO.Read<Integer, Long> reader = KafkaIO.<Integer, Long>read() .withBootstrapServers("none") .withTopic("my_topic") .withConsumerFactoryFn( new ConsumerFactoryFn( ImmutableList.of(topic), 10, numElements, OffsetResetStrategy.EARLIEST)) .withMaxNumRecords(2 * numElements) .withConsumerConfigUpdates(ImmutableMap.of("inject.error.at.eof", true)) .withKeyDeserializer(IntegerDeserializer.class) .withValueDeserializer(LongDeserializer.class); PCollection<Long> input = p.apply(reader.withoutMetadata()).apply(Values.create()); addCountingAsserts(input, numElements); p.run(); }
@Test @Ignore public void testUnboundedSourceWithoutBoundedWrapper() { final int numElements = 1000; final int numPartitions = 10; String topic = "testUnboundedSourceWithoutBoundedWrapper"; KafkaIO.Read<byte[], Long> reader = KafkaIO.<byte[], Long>read() .withBootstrapServers(topic) .withTopic(topic) .withConsumerFactoryFn( new ConsumerFactoryFn( ImmutableList.of(topic), numPartitions, numElements, OffsetResetStrategy.EARLIEST)) .withKeyDeserializer(ByteArrayDeserializer.class) .withValueDeserializer(LongDeserializer.class) .withTimestampPolicyFactory( new TimestampPolicyWithEndOfSource<>(numElements / numPartitions - 1)); p.apply("readFromKafka", reader.withoutMetadata()) .apply(Values.create()) .apply(Window.into(FixedWindows.of(Duration.standardDays(100)))); PipelineResult result = p.run(); MetricName elementsRead = SourceMetrics.elementsRead().getName(); MetricQueryResults metrics = result .metrics() .queryMetrics( MetricsFilter.builder() .addNameFilter(MetricNameFilter.inNamespace(elementsRead.getNamespace())) .build()); assertThat( metrics.getCounters(), hasItem( attemptedMetricsResult( elementsRead.getNamespace(), elementsRead.getName(), "readFromKafka", (long) numElements))); }
@Test public void testUnboundedSourceCheckpointMarkWithEmptyPartitions() throws Exception { int initialNumElements = 5; UnboundedSource<KafkaRecord<Integer, Long>, KafkaCheckpointMark> source = mkKafkaReadTransform(initialNumElements, new ValueAsTimestampFn()) .makeSource() .split(1, PipelineOptionsFactory.create()) .get(0); UnboundedReader<KafkaRecord<Integer, Long>> reader = source.createReader(null, null); for (int l = 0; l < initialNumElements; ++l) { advanceOnce(reader, l > 0); } KafkaCheckpointMark mark = CoderUtils.clone( source.getCheckpointMarkCoder(), (KafkaCheckpointMark) reader.getCheckpointMark()); int numElements = 100; List<String> topics = ImmutableList.of("topic_a", "topic_b"); source = KafkaIO.<Integer, Long>read() .withBootstrapServers("none") .withTopics(topics) .withConsumerFactoryFn( new ConsumerFactoryFn(topics, 10, numElements, OffsetResetStrategy.LATEST)) .withKeyDeserializer(IntegerDeserializer.class) .withValueDeserializer(LongDeserializer.class) .withMaxNumRecords(numElements) .withTimestampFn(new ValueAsTimestampFn()) .makeSource() .split(1, PipelineOptionsFactory.create()) .get(0); reader = source.createReader(null, mark); List<Long> expected = new ArrayList<>(); List<Long> actual = new ArrayList<>(); for (long i = initialNumElements; i < numElements; i++) { advanceOnce(reader, i > initialNumElements); expected.add(i); actual.add(reader.getCurrent().getKV().getValue()); } assertThat(actual, IsIterableContainingInAnyOrder.containsInAnyOrder(expected.toArray())); }