embulk-output-kafka 0.1.0

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

@@ -0,0 +1,323 @@
+package org.embulk.output.kafka;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonValue;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import org.apache.avro.generic.GenericData;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerRecord;
+import org.embulk.config.Config;
+import org.embulk.config.ConfigDefault;
+import org.embulk.config.ConfigDiff;
+import org.embulk.config.ConfigException;
+import org.embulk.config.ConfigSource;
+import org.embulk.config.Task;
+import org.embulk.config.TaskReport;
+import org.embulk.config.TaskSource;
+import org.embulk.spi.Exec;
+import org.embulk.spi.OutputPlugin;
+import org.embulk.spi.Page;
+import org.embulk.spi.PageReader;
+import org.embulk.spi.Schema;
+import org.embulk.spi.TransactionalPageOutput;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.List;
+import java.util.Locale;
+import java.util.Map;
+import java.util.Optional;
+import java.util.PrimitiveIterator;
+import java.util.Random;
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class KafkaOutputPlugin
+        implements OutputPlugin
+{
+    public enum RecordSerializeFormat
+    {
+        JSON,
+        AVRO_WITH_SCHEMA_REGISTRY;
+
+        @JsonValue
+        public String toString()
+        {
+            return name().toLowerCase(Locale.ENGLISH);
+        }
+
+        @JsonCreator
+        public static RecordSerializeFormat ofString(String name)
+        {
+            switch (name.toLowerCase(Locale.ENGLISH)) {
+                case "json":
+                    return JSON;
+                case "avro_with_schema_registry":
+                    return AVRO_WITH_SCHEMA_REGISTRY;
+                default:
+            }
+
+            throw new ConfigException(String.format("Unknown serialize format '%s'. Supported modes are json, avro_with_schema_registry", name));
+        }
+    }
+
+    public interface PluginTask
+            extends Task
+    {
+        @Config("brokers")
+        public List<String> getBrokers();
+
+        @Config("topic")
+        public String getTopic();
+
+        @Config("topic_column")
+        @ConfigDefault("null")
+        public Optional<String> getTopicColumn();
+
+        @Config("schema_registry_url")
+        @ConfigDefault("null")
+        public Optional<String> getSchemaRegistryUrl();
+
+        @Config("serialize_format")
+        public RecordSerializeFormat getRecordSerializeFormat();
+
+        @Config("avsc_file")
+        @ConfigDefault("null")
+        public Optional<File> getAvscFile();
+
+        @Config("avsc")
+        @ConfigDefault("null")
+        public Optional<ObjectNode> getAvsc();
+
+        @Config("key_column_name")
+        @ConfigDefault("null")
+        public Optional<String> getKeyColumnName();
+
+        @Config("record_batch_size")
+        @ConfigDefault("1000")
+        public int getRecordBatchSize();
+
+        @Config("acks")
+        @ConfigDefault("\"1\"")
+        public String getAcks();
+
+        @Config("retries")
+        @ConfigDefault("1")
+        public int getRetries();
+
+        @Config("other_producer_configs")
+        @ConfigDefault("{}")
+        public Map<String, String> getOtherProducerConfigs();
+    }
+
+    private static ObjectMapper objectMapper = new ObjectMapper();
+    private Logger logger = LoggerFactory.getLogger(getClass());
+    private int recordLoggingCount = 1;
+
+    @Override
+    public ConfigDiff transaction(ConfigSource config,
+            Schema schema, int taskCount,
+            Control control)
+    {
+        PluginTask task = config.loadConfig(PluginTask.class);
+
+        // retryable (idempotent) output:
+        // return resume(task.dump(), schema, taskCount, control);
+
+        // non-retryable (non-idempotent) output:
+        control.run(task.dump());
+        return Exec.newConfigDiff();
+    }
+
+    @Override
+    public ConfigDiff resume(TaskSource taskSource,
+            Schema schema, int taskCount,
+            Control control)
+    {
+        throw new UnsupportedOperationException("kafka output plugin does not support resuming");
+    }
+
+    @Override
+    public void cleanup(TaskSource taskSource,
+            Schema schema, int taskCount,
+            List<TaskReport> successTaskReports)
+    {
+    }
+
+    @Override
+    public TransactionalPageOutput open(TaskSource taskSource, Schema schema, int taskIndex)
+    {
+        PluginTask task = taskSource.loadTask(PluginTask.class);
+
+        switch (task.getRecordSerializeFormat()) {
+            case JSON:
+                return buildPageOutputForJson(task, schema, taskIndex);
+            case AVRO_WITH_SCHEMA_REGISTRY:
+                return buildPageOutputForAvroWithSchemaRegistry(task, schema, taskIndex);
+            default:
+                throw new ConfigException("Unknow serialize format");
+        }
+    }
+
+    private TransactionalPageOutput buildPageOutputForJson(PluginTask task, Schema schema, int taskIndex)
+    {
+        KafkaProducer<Object, ObjectNode> producer = RecordProducerFactory.getForJson(task, schema, task.getOtherProducerConfigs());
+
+        PageReader pageReader = new PageReader(schema);
+        PrimitiveIterator.OfLong randomLong = new Random().longs(1, Long.MAX_VALUE).iterator();
+        AtomicInteger counter = new AtomicInteger(0);
+
+        return new TransactionalPageOutput() {
+            @Override
+            public void add(Page page)
+            {
+                pageReader.setPage(page);
+                while (pageReader.nextRecord()) {
+                    JsonFormatColumnVisitor columnVisitor = new JsonFormatColumnVisitor(task, pageReader, objectMapper);
+
+                    pageReader.getSchema().visitColumns(columnVisitor);
+
+                    Object recordKey = columnVisitor.recordKey;
+                    if (recordKey == null) {
+                        recordKey = randomLong.next();
+                    }
+
+                    String targetTopic = columnVisitor.topicName != null ? columnVisitor.topicName : task.getTopic();
+                    ProducerRecord<Object, ObjectNode> producerRecord = new ProducerRecord<>(targetTopic, recordKey, columnVisitor.jsonNode);
+                    producer.send(producerRecord, (metadata, exception) -> {
+                        if (exception != null) {
+                            logger.error("produce error", exception);
+                        }
+
+                        logger.debug("sent record: {key: {}, value: {}}", producerRecord.key(), producerRecord.value());
+
+                        int current = counter.incrementAndGet();
+                        if (current >= recordLoggingCount) {
+                            logger.info("[task-{}] Producer sent {} records", String.format("%04d", taskIndex), current);
+                            recordLoggingCount = recordLoggingCount * 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;
+            }
+        };
+    }
+
+    private TransactionalPageOutput buildPageOutputForAvroWithSchemaRegistry(PluginTask task, Schema schema, int taskIndex)
+    {
+        KafkaProducer<Object, Object> producer = RecordProducerFactory.getForAvroWithSchemaRegistry(task, schema, task.getOtherProducerConfigs());
+
+        PageReader pageReader = new PageReader(schema);
+
+        org.apache.avro.Schema avroSchema = null;
+        if (!task.getAvsc().isPresent() && !task.getAvscFile().isPresent() || task.getAvsc().isPresent() == task.getAvscFile().isPresent()) {
+            throw new ConfigException("avro_with_schema_registry format needs either one of avsc and avsc_file");
+        }
+        if (task.getAvsc().isPresent()) {
+            avroSchema = new org.apache.avro.Schema.Parser().parse(task.getAvsc().get().toString());
+        }
+        if (task.getAvscFile().isPresent()) {
+            try {
+                avroSchema = new org.apache.avro.Schema.Parser().parse(task.getAvscFile().get());
+            }
+            catch (IOException e) {
+                e.printStackTrace();
+                throw new ConfigException("avsc_file cannot read");
+            }
+        }
+
+        final Object[] key = new Object[1];
+        final String[] topicName = new String[1];
+        PrimitiveIterator.OfLong randomLong = new Random().longs(1, Long.MAX_VALUE).iterator();
+
+        AtomicInteger counter = new AtomicInteger(0);
+
+        final org.apache.avro.Schema finalAvroSchema = avroSchema;
+        return new TransactionalPageOutput()
+        {
+            @Override
+            public void add(Page page)
+            {
+                pageReader.setPage(page);
+                while (pageReader.nextRecord()) {
+                    AvroFormatColumnVisitor columnVisitor = new AvroFormatColumnVisitor(task, pageReader, finalAvroSchema, new GenericData.Record(finalAvroSchema));
+
+                    pageReader.getSchema().visitColumns(columnVisitor);
+
+                    Object recordKey = columnVisitor.recordKey;
+                    if (recordKey == null) {
+                        recordKey = randomLong.next();
+                    }
+
+                    String targetTopic = columnVisitor.topicName != null ? columnVisitor.topicName : task.getTopic();
+
+                    ProducerRecord<Object, Object> producerRecord = new ProducerRecord<>(targetTopic, recordKey, columnVisitor.genericRecord);
+                    producer.send(producerRecord, (metadata, exception) -> {
+                        if (exception != null) {
+                            logger.error("produce error", exception);
+                        }
+
+                        logger.debug("sent record: {key: {}, value: {}}", producerRecord.key(), producerRecord.value());
+
+                        int current = counter.incrementAndGet();
+                        if (current >= recordLoggingCount) {
+                            logger.info("[task-{}] Producer sent {} records", String.format("%04d", taskIndex), current);
+                            recordLoggingCount = recordLoggingCount * 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

@@ -0,0 +1,105 @@
+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.KafkaAvroSerializer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.clients.producer.ProducerConfig;
+import org.apache.kafka.common.serialization.DoubleSerializer;
+import org.apache.kafka.common.serialization.LongSerializer;
+import org.apache.kafka.common.serialization.StringSerializer;
+import org.embulk.config.ConfigException;
+import org.embulk.spi.Column;
+import org.embulk.spi.ColumnVisitor;
+import org.embulk.spi.Schema;
+
+import java.util.Map;
+import java.util.Properties;
+
+class RecordProducerFactory
+{
+    private RecordProducerFactory() {}
+
+    private static Properties buildProperties(KafkaOutputPlugin.PluginTask task, Schema schema, Map<String, String> configs)
+    {
+        Properties kafkaProps = new Properties();
+
+        kafkaProps.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, task.getBrokers());
+        kafkaProps.put(ProducerConfig.ACKS_CONFIG, task.getAcks());
+        kafkaProps.put(ProducerConfig.RETRIES_CONFIG, task.getRetries());
+        kafkaProps.put(ProducerConfig.BATCH_SIZE_CONFIG, task.getRecordBatchSize());
+
+        configs.forEach(kafkaProps::setProperty);
+
+        if (task.getKeyColumnName().isPresent()) {
+            String keyColumnName = task.getKeyColumnName().get();
+            Column column = schema.getColumns().stream()
+                    .filter(c -> c.getName().equals(keyColumnName))
+                    .findFirst()
+                    .orElseThrow(() -> new ConfigException("key column is not found"));
+
+            column.visit(new ColumnVisitor()
+            {
+                @Override
+                public void booleanColumn(Column column)
+                {
+                    throw new RuntimeException("boolean column is not supported for key_column_name");
+                }
+
+                @Override
+                public void longColumn(Column column)
+                {
+                    kafkaProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, LongSerializer.class);
+                }
+
+                @Override
+                public void doubleColumn(Column column)
+                {
+                    kafkaProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, DoubleSerializer.class);
+                }
+
+                @Override
+                public void stringColumn(Column column)
+                {
+                    kafkaProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class);
+                }
+
+                @Override
+                public void timestampColumn(Column column)
+                {
+                    throw new RuntimeException("timestamp column is not supported for key_column_name");
+                }
+
+                @Override
+                public void jsonColumn(Column column)
+                {
+                    throw new RuntimeException("json column is not supported for key_column_name");
+                }
+            });
+        }
+        else {
+            kafkaProps.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, LongSerializer.class);
+        }
+
+        return kafkaProps;
+    }
+
+    static KafkaProducer<Object, ObjectNode> getForJson(KafkaOutputPlugin.PluginTask task, Schema schema, Map<String, String> configs)
+    {
+        return new KafkaProducer<>(buildProperties(task, schema, configs), null, new KafkaJsonSerializer());
+    }
+
+    static KafkaProducer<Object, Object> getForAvroWithSchemaRegistry(KafkaOutputPlugin.PluginTask task, Schema schema, Map<String, String> configs)
+    {
+        KafkaAvroSerializer kafkaAvroSerializer = new KafkaAvroSerializer();
+        String schemaRegistryUrl = task.getSchemaRegistryUrl().orElseThrow(() -> new ConfigException("avro_with_schema_registry format needs schema_registry_url"));
+
+        Map<String, String> avroSerializerConfigs = ImmutableMap.<String, String>builder()
+                .put(AbstractKafkaAvroSerDeConfig.SCHEMA_REGISTRY_URL_CONFIG, schemaRegistryUrl)
+                .build();
+        kafkaAvroSerializer.configure(avroSerializerConfigs, false);
+
+        return new KafkaProducer<>(buildProperties(task, schema, configs), null, kafkaAvroSerializer);
+    }
+}

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

@@ -0,0 +1,5 @@
+package org.embulk.output.kafka;
+
+public class TestKafkaOutputPlugin
+{
+}

data/src/test/resources/SimpleRecord.avsc

@@ -0,0 +1,9 @@
+{
+  "type": "record",
+  "name": "SimpleRecord",
+  "fields": [
+    {"name": "id", "type": "string"},
+    {"name": "int_item", "type": "long"},
+    {"name": "varchar_item", "type": "string"}
+  ]
+}

data/src/test/resources/config_complex.yml

@@ -0,0 +1,26 @@
+in:
+  type: file
+  path_prefix: ./src/test/resources/in_complex
+  parser:
+    charset: UTF-8
+    newline: CRLF
+    type: csv
+    delimiter: "\t"
+    quote: "\0"
+    escape: "\0"
+    null_string: 'NULL'
+    skip_header_lines: 1
+    columns:
+    - {name: 'id', type: string}
+    - {name: 'int_item', type: long}
+    - {name: 'time', type: timestamp, format: "%Y-%m-%dT%H:%M:%S"}
+    - {name: 'array', type: json}
+    - {name: 'data', type: json}
+out:
+  type: kafka
+  topic: "json-topic"
+  serialize_format: json
+  brokers:
+    - "localhost:9092"
+  other_producer_configs:
+    buffer.memory: "67108864"

data/src/test/resources/config_complex_avro.yml

@@ -0,0 +1,43 @@
+in:
+  type: file
+  path_prefix: ./src/test/resources/in_complex
+  parser:
+    charset: UTF-8
+    newline: CRLF
+    type: csv
+    delimiter: "\t"
+    quote: "\0"
+    escape: "\0"
+    null_string: 'NULL'
+    skip_header_lines: 1
+    columns:
+    - {name: 'id', type: string}
+    - {name: 'int_item', type: long}
+    - {name: 'time', type: timestamp, format: "%Y-%m-%dT%H:%M:%S"}
+    - {name: 'array', type: json}
+    - {name: 'data', type: json}
+out:
+  type: kafka
+  topic: "avro-complex-topic"
+  acks: all
+  retries: 3
+  brokers:
+    - "localhost:9092"
+  schema_registry_url: "http://localhost:48081/"
+  serialize_format: avro_with_schema_registry
+  other_producer_configs:
+    buffer.memory: "67108864"
+  avsc:
+    type: record
+    name: ComplexRecord
+    fields: [
+      {name: "id", type: "string"},
+      {name: "int_item", type: "long"},
+      {name: "time", type: "long", logicalType: "timestamp-milli"},
+      {name: "array", type: {type: "array", items: "long"}},
+      {name: "data", type: {type: "record", name: "InnerData", fields: [
+        {name: "hoge", type: "string"},
+        {name: "aaa", type: ["null", "string"]},
+        {name: "array", type: {type: "array", items: "long"}},
+      ]}},
+    ]

data/src/test/resources/config_simple.yml

@@ -0,0 +1,22 @@
+in:
+  type: file
+  path_prefix: ./src/test/resources/in1
+  parser:
+    charset: UTF-8
+    newline: CRLF
+    type: csv
+    delimiter: ','
+    quote: '"'
+    escape: '"'
+    null_string: 'NULL'
+    skip_header_lines: 1
+    columns:
+    - {name: 'id', type: string}
+    - {name: 'int_item', type: long}
+    - {name: 'varchar_item', type: string}
+out:
+  type: kafka
+  topic: "json-topic"
+  serialize_format: json
+  brokers:
+    - "localhost:9092"

data/src/test/resources/config_simple_avro.yml

@@ -0,0 +1,32 @@
+in:
+  type: file
+  path_prefix: ./src/test/resources/in1
+  parser:
+    charset: UTF-8
+    newline: CRLF
+    type: csv
+    delimiter: ','
+    quote: '"'
+    escape: '"'
+    null_string: 'NULL'
+    skip_header_lines: 1
+    columns:
+    - {name: 'id', type: string}
+    - {name: 'int_item', type: long}
+    - {name: 'varchar_item', type: string}
+
+out:
+  type: kafka
+  topic: "avro-simple-topic"
+  brokers:
+    - "localhost:9092"
+  schema_registry_url: "http://localhost:48081/"
+  serialize_format: avro_with_schema_registry
+  avsc:
+    type: record
+    name: SimpleRecord
+    fields: [
+      {name: "id", type: "string"},
+      {name: "int_item", type: "long"},
+      {name: "varchar_item", type: "string"},
+    ]

data/src/test/resources/config_simple_avro_avsc_file.yml

@@ -0,0 +1,25 @@
+in:
+  type: file
+  path_prefix: ./src/test/resources/in1
+  parser:
+    charset: UTF-8
+    newline: CRLF
+    type: csv
+    delimiter: ','
+    quote: '"'
+    escape: '"'
+    null_string: 'NULL'
+    skip_header_lines: 1
+    columns:
+    - {name: 'id', type: string}
+    - {name: 'int_item', type: long}
+    - {name: 'varchar_item', type: string}
+
+out:
+  type: kafka
+  topic: "avro-simple-topic"
+  brokers:
+    - "localhost:9092"
+  schema_registry_url: "http://localhost:48081/"
+  serialize_format: avro_with_schema_registry
+  avsc_file: ./src/test/resources/SimpleRecord.avsc

data/src/test/resources/config_with_key_column.yml

@@ -0,0 +1,23 @@
+in:
+  type: file
+  path_prefix: ./src/test/resources/in1
+  parser:
+    charset: UTF-8
+    newline: CRLF
+    type: csv
+    delimiter: ','
+    quote: '"'
+    escape: '"'
+    null_string: 'NULL'
+    skip_header_lines: 1
+    columns:
+    - {name: 'id', type: string}
+    - {name: 'int_item', type: long}
+    - {name: 'varchar_item', type: string}
+out:
+  type: kafka
+  topic: "json-topic"
+  serialize_format: json
+  brokers:
+    - "localhost:9092"
+  key_column_name: id

data/src/test/resources/in1.csv

@@ -0,0 +1,4 @@
+id,int_item,varchar_item
+A001,9,a
+A002,0,b
+A003,9,c

data/src/test/resources/in_complex.csv

@@ -0,0 +1,5 @@
+id	int_item	time	array	data
+A001	9	2018-02-01T12:15:18	[1,2,3]	{"hoge": "fuga1", "aaa": "bbb1", "array": [1,2,3]}
+A002	0	2018-02-02T12:15:18	[1,2,3]	{"hoge": "fuga2", "aaa": null, "array": [4,5,6]}
+A003	9	2018-02-03T12:15:18	[1,2,3]	{"hoge": "fuga3", "aaa": "bbb3", "array": [7,8,9]}
+