embulk-output-kafka 0.1.7 → 0.1.8

data/src/main/java/org/embulk/output/kafka/KafkaTransactionalPageOutput.java

@@ -0,0 +1,104 @@
+package org.embulk.output.kafka;
+
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.embulk.config.TaskReport;
+import org.embulk.spi.Page;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.TransactionalPageOutput;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.PrimitiveIterator;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicLong;
+
+public abstract class KafkaTransactionalPageOutput<P, T extends P> implements TransactionalPageOutput
+{
+    private static final Logger logger = LoggerFactory.getLogger(KafkaTransactionalPageOutput.class);
+
+    private final KafkaProducer<Object, P> producer;
+    private final PageReader pageReader;
+    private final KafkaOutputColumnVisitor<T> columnVisitor;
+    private final String topic;
+    private final int taskIndex;
+
+    private final PrimitiveIterator.OfLong randomLong = new Random().longs(1, Long.MAX_VALUE).iterator();
+    private final AtomicLong counter = new AtomicLong(0);
+    private final AtomicLong recordLoggingCount = new AtomicLong(1);
+
+    public KafkaTransactionalPageOutput(
+            KafkaProducer<Object, P> producer,
+            PageReader pageReader,
+            KafkaOutputColumnVisitor<T> columnVisitor,
+            String topic, int taskIndex)
+    {
+        this.producer = producer;
+        this.pageReader = pageReader;
+        this.columnVisitor = columnVisitor;
+        this.topic = topic;
+        this.taskIndex = taskIndex;
+    }
+
+    @Override
+    public void add(Page page)
+    {
+        pageReader.setPage(page);
+        while (pageReader.nextRecord()) {
+            columnVisitor.reset();
+
+            pageReader.getSchema().visitColumns(columnVisitor);
+
+            Object recordKey = columnVisitor.getRecordKey();
+            if (recordKey == null) {
+                recordKey = randomLong.next();
+            }
+
+            String targetTopic = columnVisitor.getTopicName() != null ? columnVisitor.getTopicName() : topic;
+
+            ProducerRecord<Object, P> producerRecord = new ProducerRecord<>(targetTopic, columnVisitor.getPartition(), recordKey, columnVisitor.getRecord());
+            producer.send(producerRecord, (metadata, exception) -> {
+                if (exception != null) {
+                    logger.error("produce error", exception);
+                }
+
+                logger.debug("sent record: {topic: {}, key: {}, value: {}, partition: {}}",
+                        producerRecord.topic(),
+                        producerRecord.key(),
+                        producerRecord.value(),
+                        producerRecord.partition());
+
+                long current = counter.incrementAndGet();
+                if (current >= recordLoggingCount.get()) {
+                    logger.info("[task-{}] Producer sent {} records", String.format("%04d", taskIndex), current);
+                    recordLoggingCount.set(recordLoggingCount.get() * 2);
+                }
+            });
+        }
+    }
+
+    @Override
+    public void finish()
+    {
+        producer.flush();
+    }
+
+    @Override
+    public void close()
+    {
+        producer.close();
+    }
+
+    @Override
+    public void abort()
+    {
+        producer.flush();
+        producer.close();
+    }
+
+    @Override
+    public TaskReport commit()
+    {
+        return null;
+    }
+};

data/src/main/java/org/embulk/output/kafka/RecordProducerFactory.java

@@ -2,7 +2,7 @@ package org.embulk.output.kafka;
 
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import com.google.common.collect.ImmutableMap;
-import io.confluent.kafka.serializers.AbstractKafkaAvroSerDeConfig;
+import io.confluent.kafka.serializers.AbstractKafkaSchemaSerDeConfig;
 import io.confluent.kafka.serializers.KafkaAvroSerializer;
 import org.apache.kafka.clients.producer.KafkaProducer;
 import org.apache.kafka.clients.producer.ProducerConfig;
@@ -96,10 +96,10 @@ class RecordProducerFactory
         String schemaRegistryUrl = task.getSchemaRegistryUrl().orElseThrow(() -> new ConfigException("avro_with_schema_registry format needs schema_registry_url"));
 
         ImmutableMap.Builder<String, String> builder = ImmutableMap.<String, String>builder()
-                .put(AbstractKafkaAvroSerDeConfig.SCHEMA_REGISTRY_URL_CONFIG, schemaRegistryUrl);
+                .put(AbstractKafkaSchemaSerDeConfig.SCHEMA_REGISTRY_URL_CONFIG, schemaRegistryUrl);
 
         if (task.getValueSubjectNameStrategy().isPresent()) {
-            builder.put(AbstractKafkaAvroSerDeConfig.VALUE_SUBJECT_NAME_STRATEGY, task.getValueSubjectNameStrategy().get());
+            builder.put(AbstractKafkaSchemaSerDeConfig.VALUE_SUBJECT_NAME_STRATEGY, task.getValueSubjectNameStrategy().get());
         }
 
         Map<String, String> avroSerializerConfigs = builder.build();

data/src/test/java/org/embulk/output/kafka/TestKafkaOutputPlugin.java

@@ -1,5 +1,389 @@
 package org.embulk.output.kafka;
 
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.hasItem;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.collect.ImmutableList;
+import com.google.common.io.Resources;
+import com.salesforce.kafka.test.KafkaTestUtils;
+import com.salesforce.kafka.test.junit4.SharedKafkaTestResource;
+import io.confluent.kafka.schemaregistry.ParsedSchema;
+import io.confluent.kafka.schemaregistry.avro.AvroSchema;
+import io.confluent.kafka.schemaregistry.client.SchemaRegistryClient;
+import io.confluent.kafka.schemaregistry.client.rest.exceptions.RestClientException;
+import io.confluent.kafka.schemaregistry.testutil.MockSchemaRegistry;
+import io.confluent.kafka.serializers.KafkaAvroDeserializer;
+import java.io.IOException;
+import java.nio.file.Paths;
+import java.time.Instant;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.kafka.clients.consumer.ConsumerRecord;
+import org.apache.kafka.common.serialization.StringDeserializer;
+import org.embulk.config.ConfigSource;
+import org.embulk.spi.OutputPlugin;
+import org.embulk.test.TestingEmbulk;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+
 public class TestKafkaOutputPlugin
 {
+    @ClassRule
+    public static final SharedKafkaTestResource sharedKafkaTestResource = new SharedKafkaTestResource()
+        .withBrokers(3);
+
+    @Rule
+    public TestingEmbulk embulk = TestingEmbulk.builder()
+            .registerPlugin(OutputPlugin.class, "kafka", KafkaOutputPlugin.class)
+            .build();
+
+    private KafkaTestUtils kafkaTestUtils;
+    private final static ObjectMapper objectMapper = new ObjectMapper();
+
+    @Before
+    public void setUp() {
+        kafkaTestUtils = sharedKafkaTestResource.getKafkaTestUtils();
+        kafkaTestUtils.createTopic("json-topic", 8, (short) 1);
+        kafkaTestUtils.createTopic("json-complex-topic", 8, (short) 1);
+        kafkaTestUtils.createTopic("avro-simple-topic", 8, (short) 1);
+        kafkaTestUtils.createTopic("avro-complex-topic", 8, (short) 1);
+    }
+
+    @After
+    public void tearDown() {
+        kafkaTestUtils.getAdminClient().deleteTopics(ImmutableList.of(
+            "json-topic", "json-complex-topic", "avro-simple-topic", "avro-complex-topic"
+        ));
+    }
+
+    @Test
+    public void testSimpleJson() throws IOException
+    {
+        ConfigSource configSource = embulk.loadYamlResource("config_simple.yml");
+        configSource.set("brokers", ImmutableList.of(sharedKafkaTestResource.getKafkaBrokers().getBrokerById(1).getConnectString()));
+        embulk.runOutput(configSource, Paths.get(Resources.getResource("in1.csv").getPath()));
+        List<ConsumerRecord<String, String>> consumerRecords = kafkaTestUtils
+            .consumeAllRecordsFromTopic("json-topic", StringDeserializer.class,
+                StringDeserializer.class);
+
+        assertEquals(3, consumerRecords.size());
+        List<JsonNode> deserializedRecords = new ArrayList<>();
+        for (ConsumerRecord<String, String> record : consumerRecords) {
+            deserializedRecords.add(objectMapper.readTree(record.value()));
+        }
+        List<String> ids = deserializedRecords.stream()
+            .map(r -> r.get("id").asText())
+            .collect(Collectors.toList());
+        List<Integer> intItems = deserializedRecords.stream()
+            .map(r -> r.get("int_item").asInt())
+            .collect(Collectors.toList());
+        List<String> varcharItems = deserializedRecords.stream()
+            .map(r -> r.get("varchar_item").asText())
+            .collect(Collectors.toList());
+
+        assertThat(ids, hasItem("A001"));
+        assertThat(ids, hasItem("A002"));
+        assertThat(ids, hasItem("A003"));
+        assertThat(intItems, hasItem(1));
+        assertThat(intItems, hasItem(2));
+        assertThat(intItems, hasItem(3));
+        assertThat(varcharItems, hasItem("a"));
+        assertThat(varcharItems, hasItem("b"));
+        assertThat(varcharItems, hasItem("c"));
+    }
+
+    @Test
+    public void testComplexJson() throws IOException
+    {
+        ConfigSource configSource = embulk.loadYamlResource("config_complex.yml");
+        configSource.set("brokers", ImmutableList.of(sharedKafkaTestResource.getKafkaBrokers().getBrokerById(1).getConnectString()));
+
+        embulk.runOutput(configSource, Paths.get(Resources.getResource("in_complex.csv").getPath()));
+        List<ConsumerRecord<String, String>> consumerRecords = kafkaTestUtils
+            .consumeAllRecordsFromTopic("json-complex-topic", StringDeserializer.class,
+                StringDeserializer.class);
+
+        assertEquals(3, consumerRecords.size());
+        List<JsonNode> deserializedRecords = new ArrayList<>();
+        for (ConsumerRecord<String, String> record : consumerRecords) {
+            deserializedRecords.add(objectMapper.readTree(record.value()));
+        }
+        List<String> ids = deserializedRecords.stream()
+            .map(r -> r.get("id").asText())
+            .collect(Collectors.toList());
+        List<Integer> intItems = deserializedRecords.stream()
+            .map(r -> r.get("int_item").asInt())
+            .collect(Collectors.toList());
+        List<List<Integer>> arrayItems = deserializedRecords.stream()
+            .map(r -> ImmutableList.of(
+                r.get("array").get(0).asInt(),
+                r.get("array").get(1).asInt(),
+                r.get("array").get(2).asInt()
+            ))
+            .collect(Collectors.toList());
+
+        assertThat(ids, hasItem("A001"));
+        assertThat(ids, hasItem("A002"));
+        assertThat(ids, hasItem("A003"));
+        assertThat(intItems, hasItem(9));
+        assertThat(intItems, hasItem(0));
+        assertThat(arrayItems.get(0), hasItem(1));
+        assertThat(arrayItems.get(0), hasItem(2));
+        assertThat(arrayItems.get(0), hasItem(3));
+    }
+
+    @Test
+    public void testSimpleAvro() throws IOException {
+        ConfigSource configSource = embulk.loadYamlResource("config_simple_avro.yml");
+        configSource.set("brokers", ImmutableList
+            .of(sharedKafkaTestResource.getKafkaBrokers().getBrokerById(1).getConnectString()));
+
+        embulk.runOutput(configSource, Paths.get(Resources.getResource("in1.csv").getPath()));
+
+        SchemaRegistryClient schemaRegistryClient = MockSchemaRegistry
+            .getClientForScope("embulk-output-kafka");
+        KafkaAvroDeserializer kafkaAvroDeserializer = new KafkaAvroDeserializer(schemaRegistryClient);
+
+        List<ConsumerRecord<byte[], byte[]>> consumerRecords = kafkaTestUtils
+            .consumeAllRecordsFromTopic("avro-simple-topic");
+
+        assertEquals(3, consumerRecords.size());
+        List<GenericRecord> genericRecords = consumerRecords.stream().map(r -> (GenericRecord) kafkaAvroDeserializer
+            .deserialize("avro-simple-topic", r.value())).collect(Collectors.toList());
+
+        List<String> ids = genericRecords.stream()
+            .map(r -> String.valueOf(r.get("id")))
+            .collect(Collectors.toList());
+        List<Long> intItems = genericRecords.stream()
+            .map(r -> (Long) r.get("int_item"))
+            .collect(Collectors.toList());
+        List<String> varcharItems = genericRecords.stream()
+            .map(r -> String.valueOf(r.get("varchar_item")))
+            .collect(Collectors.toList());
+
+        assertThat(ids, hasItem("A001"));
+        assertThat(ids, hasItem("A002"));
+        assertThat(ids, hasItem("A003"));
+        assertThat(intItems, hasItem(1L));
+        assertThat(intItems, hasItem(2L));
+        assertThat(intItems, hasItem(3L));
+        assertThat(varcharItems, hasItem("a"));
+        assertThat(varcharItems, hasItem("b"));
+        assertThat(varcharItems, hasItem("c"));
+    }
+
+    @Test
+    public void testSimpleAvroSchemaFromRegistry() throws IOException, RestClientException
+    {
+        ConfigSource configSource = embulk.loadYamlResource("config_simple_avro.yml");
+        Object avsc = configSource.get(Object.class, "avsc");
+        String avscString = objectMapper.writeValueAsString(avsc);
+        configSource.set("avsc", null);
+        ParsedSchema parsedSchema = new AvroSchema(avscString);
+        MockSchemaRegistry.getClientForScope("embulk-output-kafka")
+                .register("avro-simple-topic-value", parsedSchema);
+        configSource.set("brokers", ImmutableList
+                .of(sharedKafkaTestResource.getKafkaBrokers().getBrokerById(1).getConnectString()));
+
+        embulk.runOutput(configSource, Paths.get(Resources.getResource("in1.csv").getPath()));
+
+        SchemaRegistryClient schemaRegistryClient = MockSchemaRegistry
+                .getClientForScope("embulk-output-kafka");
+        KafkaAvroDeserializer kafkaAvroDeserializer = new KafkaAvroDeserializer(schemaRegistryClient);
+
+        List<ConsumerRecord<byte[], byte[]>> consumerRecords = kafkaTestUtils
+                .consumeAllRecordsFromTopic("avro-simple-topic");
+
+        assertEquals(3, consumerRecords.size());
+        List<GenericRecord> genericRecords = consumerRecords.stream().map(r -> (GenericRecord) kafkaAvroDeserializer
+                .deserialize("avro-simple-topic", r.value())).collect(Collectors.toList());
+
+        List<String> ids = genericRecords.stream()
+                .map(r -> String.valueOf(r.get("id")))
+                .collect(Collectors.toList());
+        List<Long> intItems = genericRecords.stream()
+                .map(r -> (Long) r.get("int_item"))
+                .collect(Collectors.toList());
+        List<String> varcharItems = genericRecords.stream()
+                .map(r -> String.valueOf(r.get("varchar_item")))
+                .collect(Collectors.toList());
+
+        assertThat(ids, hasItem("A001"));
+        assertThat(ids, hasItem("A002"));
+        assertThat(ids, hasItem("A003"));
+        assertThat(intItems, hasItem(1L));
+        assertThat(intItems, hasItem(2L));
+        assertThat(intItems, hasItem(3L));
+        assertThat(varcharItems, hasItem("a"));
+        assertThat(varcharItems, hasItem("b"));
+        assertThat(varcharItems, hasItem("c"));
+    }
+
+    @Test
+    public void testSimpleAvroAvscFile() throws IOException {
+        ConfigSource configSource = embulk.loadYamlResource("config_simple_avro_avsc_file.yml");
+        configSource.set("brokers", ImmutableList
+            .of(sharedKafkaTestResource.getKafkaBrokers().getBrokerById(1).getConnectString()));
+
+        embulk.runOutput(configSource, Paths.get(Resources.getResource("in1.csv").getPath()));
+
+        SchemaRegistryClient schemaRegistryClient = MockSchemaRegistry
+            .getClientForScope("embulk-output-kafka");
+        KafkaAvroDeserializer kafkaAvroDeserializer = new KafkaAvroDeserializer(schemaRegistryClient);
+
+        List<ConsumerRecord<byte[], byte[]>> consumerRecords = kafkaTestUtils
+            .consumeAllRecordsFromTopic("avro-simple-topic");
+
+        assertEquals(3, consumerRecords.size());
+        List<GenericRecord> genericRecords = consumerRecords.stream().map(r -> (GenericRecord) kafkaAvroDeserializer
+            .deserialize("avro-simple-topic", r.value())).collect(Collectors.toList());
+
+        List<String> ids = genericRecords.stream()
+            .map(r -> String.valueOf(r.get("id")))
+            .collect(Collectors.toList());
+        List<Long> intItems = genericRecords.stream()
+            .map(r -> (Long) r.get("int_item"))
+            .collect(Collectors.toList());
+        List<String> varcharItems = genericRecords.stream()
+            .map(r -> String.valueOf(r.get("varchar_item")))
+            .collect(Collectors.toList());
+
+        assertThat(ids, hasItem("A001"));
+        assertThat(ids, hasItem("A002"));
+        assertThat(ids, hasItem("A003"));
+        assertThat(intItems, hasItem(1L));
+        assertThat(intItems, hasItem(2L));
+        assertThat(intItems, hasItem(3L));
+        assertThat(varcharItems, hasItem("a"));
+        assertThat(varcharItems, hasItem("b"));
+        assertThat(varcharItems, hasItem("c"));
+    }
+
+    @Test
+    public void testSimpleAvroComplex() throws IOException {
+        ConfigSource configSource = embulk.loadYamlResource("config_complex_avro.yml");
+        configSource.set("brokers", ImmutableList
+            .of(sharedKafkaTestResource.getKafkaBrokers().getBrokerById(1).getConnectString()));
+
+        embulk.runOutput(configSource, Paths.get(Resources.getResource("in_complex.csv").getPath()));
+
+        SchemaRegistryClient schemaRegistryClient = MockSchemaRegistry
+            .getClientForScope("embulk-output-kafka");
+        KafkaAvroDeserializer kafkaAvroDeserializer = new KafkaAvroDeserializer(schemaRegistryClient);
+
+        List<ConsumerRecord<byte[], byte[]>> consumerRecords = kafkaTestUtils
+            .consumeAllRecordsFromTopic("avro-complex-topic");
+
+        assertEquals(3, consumerRecords.size());
+        List<GenericRecord> genericRecords = consumerRecords.stream().map(r -> (GenericRecord) kafkaAvroDeserializer
+            .deserialize("avro-complex-topic", r.value())).collect(Collectors.toList());
+
+        List<String> ids = genericRecords.stream()
+            .map(r -> String.valueOf(r.get("id")))
+            .collect(Collectors.toList());
+        List<Long> intItems = genericRecords.stream()
+            .map(r -> (Long) r.get("int_item"))
+            .collect(Collectors.toList());
+        List<Instant> timeItems = genericRecords.stream()
+            .map(r -> Instant.ofEpochMilli((long) r.get("time")))
+            .collect(Collectors.toList());
+
+        assertThat(ids, hasItem("A001"));
+        assertThat(ids, hasItem("A002"));
+        assertThat(ids, hasItem("A003"));
+        assertThat(intItems, hasItem(9L));
+        assertThat(intItems, hasItem(0L));
+        assertThat(timeItems, hasItem(Instant.parse("2018-02-01T12:15:18.000Z")));
+        assertThat(timeItems, hasItem(Instant.parse("2018-02-02T12:15:18.000Z")));
+        assertThat(timeItems, hasItem(Instant.parse("2018-02-03T12:15:18.000Z")));
+    }
+
+    @Test
+    public void testKeyColumnConfig() throws IOException
+    {
+        ConfigSource configSource = embulk.loadYamlResource("config_with_key_column.yml");
+        configSource.set("brokers", ImmutableList.of(sharedKafkaTestResource.getKafkaBrokers().getBrokerById(1).getConnectString()));
+        embulk.runOutput(configSource, Paths.get(Resources.getResource("in1.csv").getPath()));
+        List<ConsumerRecord<String, String>> consumerRecords = kafkaTestUtils
+            .consumeAllRecordsFromTopic("json-topic", StringDeserializer.class,
+                StringDeserializer.class);
+
+        assertEquals(3, consumerRecords.size());
+        List<String> keys = new ArrayList<>();
+        for (ConsumerRecord<String, String> record : consumerRecords) {
+          keys.add(record.key());
+        }
+
+        assertThat(keys, hasItem("A001"));
+        assertThat(keys, hasItem("A002"));
+        assertThat(keys, hasItem("A003"));
+    }
+
+    @Test
+    public void testPartitionColumnConfig() throws IOException
+    {
+        ConfigSource configSource = embulk.loadYamlResource("config_with_partition_column.yml");
+        configSource.set("brokers", ImmutableList.of(sharedKafkaTestResource.getKafkaBrokers().getBrokerById(1).getConnectString()));
+        embulk.runOutput(configSource, Paths.get(Resources.getResource("in1.csv").getPath()));
+        List<ConsumerRecord<String, String>> consumerRecords = kafkaTestUtils
+            .consumeAllRecordsFromTopic("json-topic", StringDeserializer.class,
+                StringDeserializer.class);
+
+        assertEquals(3, consumerRecords.size());
+        List<Integer> partitions = new ArrayList<>();
+        for (ConsumerRecord<String, String> record : consumerRecords) {
+            partitions.add(record.partition());
+        }
+
+        assertThat(partitions, hasItem(1));
+        assertThat(partitions, hasItem(2));
+        assertThat(partitions, hasItem(3));
+    }
+
+    @Test
+    public void testColumnForDeletion() throws IOException
+    {
+        ConfigSource configSource = embulk.loadYamlResource("config_with_column_for_deletion.yml");
+        configSource.set("brokers", ImmutableList.of(sharedKafkaTestResource.getKafkaBrokers().getBrokerById(1).getConnectString()));
+        embulk.runOutput(configSource, Paths.get(Resources.getResource("in_with_deletion.csv").getPath()));
+        List<ConsumerRecord<String, String>> consumerRecords = kafkaTestUtils
+                .consumeAllRecordsFromTopic("json-topic", StringDeserializer.class,
+                        StringDeserializer.class);
+
+        assertEquals(3, consumerRecords.size());
+        HashMap<String, String> recordMap = new HashMap<>();
+        consumerRecords.forEach(record -> recordMap.put(record.key(), record.value()));
+        assertNotNull(recordMap.get("A001"));
+        assertNotNull(recordMap.get("A003"));
+        assertNull(recordMap.get("A002"));
+    }
+
+    @Test
+    public void testColumnForDeletionAvro() throws IOException
+    {
+        ConfigSource configSource = embulk.loadYamlResource("config_with_column_for_deletion_avro.yml");
+        configSource.set("brokers", ImmutableList.of(sharedKafkaTestResource.getKafkaBrokers().getBrokerById(1).getConnectString()));
+        embulk.runOutput(configSource, Paths.get(Resources.getResource("in_with_deletion.csv").getPath()));
+        List<ConsumerRecord<String, String>> consumerRecords = kafkaTestUtils
+                .consumeAllRecordsFromTopic("avro-simple-topic", StringDeserializer.class,
+                        StringDeserializer.class);
+
+        assertEquals(3, consumerRecords.size());
+        HashMap<String, String> recordMap = new HashMap<>();
+        consumerRecords.forEach(record -> recordMap.put(record.key(), record.value()));
+        assertNotNull(recordMap.get("A001"));
+        assertNotNull(recordMap.get("A003"));
+        assertNull(recordMap.get("A002"));
+    }
 }