KafkaIO { public static <K, V> Write<K, V> write() { return new AutoValue_KafkaIO_Write.Builder<K, V>() .setWriteRecordsTransform( new AutoValue_KafkaIO_WriteRecords.Builder<K, V>() .setProducerConfig(WriteRecords.DEFAULT_PRODUCER_PROPERTIES) .setEOS(false) .setNumShards(0) .setConsumerFactoryFn(KafkaIOUtils.KAFKA_CONSUMER_FACTORY_FN) .build()) .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 testSink() throws Exception { int numElements = 1000; try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { ProducerSendCompletionThread completionThread = new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); String topic = "test"; p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) .apply( KafkaIO.<Integer, Long>write() .withBootstrapServers("none") .withTopic(topic) .withKeySerializer(IntegerSerializer.class) .withValueSerializer(LongSerializer.class) .withInputTimestamp() .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); p.run(); completionThread.shutdown(); verifyProducerRecords(producerWrapper.mockProducer, topic, numElements, false, true); } }
@Test public void testValuesSink() throws Exception { int numElements = 1000; try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { ProducerSendCompletionThread completionThread = new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); String topic = "test"; p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) .apply(Values.create()) .apply( KafkaIO.<Integer, Long>write() .withBootstrapServers("none") .withTopic(topic) .withValueSerializer(LongSerializer.class) .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey)) .values()); p.run(); completionThread.shutdown(); verifyProducerRecords(producerWrapper.mockProducer, topic, numElements, true, false); } }
@Test public void testExactlyOnceSink() { if (!ProducerSpEL.supportsTransactions()) { LOG.warn( "testExactlyOnceSink() is disabled as Kafka client version does not support transactions."); return; } int numElements = 1000; try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { ProducerSendCompletionThread completionThread = new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); String topic = "test"; p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) .apply( KafkaIO.<Integer, Long>write() .withBootstrapServers("none") .withTopic(topic) .withKeySerializer(IntegerSerializer.class) .withValueSerializer(LongSerializer.class) .withEOS(1, "test") .withConsumerFactoryFn( new ConsumerFactoryFn( Lists.newArrayList(topic), 10, 10, OffsetResetStrategy.EARLIEST)) .withPublishTimestampFunction((e, ts) -> ts) .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); p.run(); completionThread.shutdown(); verifyProducerRecords(producerWrapper.mockProducer, topic, numElements, false, true); } }
@Test public void testSinkWithSendErrors() throws Throwable { thrown.expect(InjectedErrorException.class); thrown.expectMessage("Injected Error #1"); int numElements = 1000; try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { ProducerSendCompletionThread completionThreadWithErrors = new ProducerSendCompletionThread(producerWrapper.mockProducer, 10, 100).start(); String topic = "test"; p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) .apply( KafkaIO.<Integer, Long>write() .withBootstrapServers("none") .withTopic(topic) .withKeySerializer(IntegerSerializer.class) .withValueSerializer(LongSerializer.class) .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); try { p.run(); } catch (PipelineExecutionException e) { throw e.getCause().getCause(); } finally { completionThreadWithErrors.shutdown(); } } }
@Test public void testSinkDisplayData() { try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { KafkaIO.Write<Integer, Long> write = KafkaIO.<Integer, Long>write() .withBootstrapServers("myServerA:9092,myServerB:9092") .withTopic("myTopic") .withValueSerializer(LongSerializer.class) .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey)) .withProducerConfigUpdates(ImmutableMap.of("retry.backoff.ms", 100)); DisplayData displayData = DisplayData.from(write); assertThat(displayData, hasDisplayItem("topic", "myTopic")); assertThat(displayData, hasDisplayItem("bootstrap.servers", "myServerA:9092,myServerB:9092")); assertThat(displayData, hasDisplayItem("retries", 3)); assertThat(displayData, hasDisplayItem("retry.backoff.ms", 100)); } }
@Test public void testSinkMetrics() throws Exception { int numElements = 1000; try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) { ProducerSendCompletionThread completionThread = new ProducerSendCompletionThread(producerWrapper.mockProducer).start(); String topic = "test"; p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()) .apply( "writeToKafka", KafkaIO.<Integer, Long>write() .withBootstrapServers("none") .withTopic(topic) .withKeySerializer(IntegerSerializer.class) .withValueSerializer(LongSerializer.class) .withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey))); PipelineResult result = p.run(); MetricName elementsWritten = SinkMetrics.elementsWritten().getName(); MetricQueryResults metrics = result .metrics() .queryMetrics( MetricsFilter.builder() .addNameFilter(MetricNameFilter.inNamespace(elementsWritten.getNamespace())) .build()); assertThat( metrics.getCounters(), hasItem( attemptedMetricsResult( elementsWritten.getNamespace(), elementsWritten.getName(), "writeToKafka", 1000L))); completionThread.shutdown(); } }