embulk-output-kafka 0.1.3 → 0.1.8

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

@@ -4,9 +4,17 @@ 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 com.google.common.collect.ImmutableList;
+import io.confluent.kafka.schemaregistry.client.CachedSchemaRegistryClient;
+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.subject.TopicNameStrategy;
+import io.confluent.kafka.serializers.subject.strategy.SubjectNameStrategy;
+import org.apache.kafka.clients.admin.AdminClient;
+import org.apache.kafka.clients.admin.AdminClientConfig;
+import org.apache.kafka.clients.admin.DescribeTopicsResult;
 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;
@@ -15,15 +23,11 @@ import org.embulk.config.ConfigSource;
 import org.embulk.config.Task;
 import org.embulk.config.TaskReport;
 import org.embulk.config.TaskSource;
-import org.embulk.spi.ColumnConfig;
 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;
@@ -31,9 +35,10 @@ 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;
+import java.util.Properties;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 public class KafkaOutputPlugin
         implements OutputPlugin
@@ -96,8 +101,12 @@ public class KafkaOutputPlugin
         @ConfigDefault("null")
         public Optional<String> getKeyColumnName();
 
+        @Config("partition_column_name")
+        @ConfigDefault("null")
+        public Optional<String> getPartitionColumnName();
+
         @Config("record_batch_size")
-        @ConfigDefault("1000")
+        @ConfigDefault("16384")
         public int getRecordBatchSize();
 
         @Config("acks")
@@ -119,10 +128,22 @@ public class KafkaOutputPlugin
         @Config("value_subject_name_strategy")
         @ConfigDefault("null")
         public Optional<String> getValueSubjectNameStrategy();
+
+        @Config("column_for_deletion")
+        @ConfigDefault("null")
+        public Optional<String> getColumnForDeletion();
     }
 
     private static ObjectMapper objectMapper = new ObjectMapper();
-    private Logger logger = LoggerFactory.getLogger(getClass());
+
+    private static final int SCHEMA_REGISTRY_IDENTITY_MAP_CAPACITY = 1000;
+
+    private AdminClient getKafkaAdminClient(PluginTask task)
+    {
+        Properties properties = new Properties();
+        properties.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, task.getBrokers());
+        return AdminClient.create(properties);
+    }
 
     @Override
     public ConfigDiff transaction(ConfigSource config,
@@ -130,11 +151,17 @@ public class KafkaOutputPlugin
             Control control)
     {
         PluginTask task = config.loadConfig(PluginTask.class);
+        AdminClient adminClient = getKafkaAdminClient(task);
+        DescribeTopicsResult result = adminClient.describeTopics(ImmutableList.of(task.getTopic()));
+        try {
+            if (result.all().get(30, TimeUnit.SECONDS).size() == 0) {
+                throw new RuntimeException("target topic is not found");
+            }
+        }
+        catch (InterruptedException | ExecutionException | TimeoutException e) {
+            throw new RuntimeException("failed to connect kafka brokers");
+        }
 
-        // retryable (idempotent) output:
-        // return resume(task.dump(), schema, taskCount, control);
-
-        // non-retryable (non-idempotent) output:
         control.run(task.dump());
         return Exec.newConfigDiff();
     }
@@ -165,95 +192,48 @@ public class KafkaOutputPlugin
             case AVRO_WITH_SCHEMA_REGISTRY:
                 return buildPageOutputForAvroWithSchemaRegistry(task, schema, taskIndex);
             default:
-                throw new ConfigException("Unknow serialize format");
+                throw new ConfigException("Unknown 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);
-        AtomicInteger recordLoggingCount = new AtomicInteger(1);
-
-        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.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();
-            }
+        KafkaOutputColumnVisitor<ObjectNode> columnVisitor = new JsonFormatColumnVisitor(task, pageReader, objectMapper);
 
-            @Override
-            public void abort()
-            {
-                producer.flush();
-                producer.close();
-            }
-
-            @Override
-            public TaskReport commit()
-            {
-                return null;
-            }
-        };
+        return new JsonFormatTransactionalPageOutput(producer, pageReader, columnVisitor, task.getTopic(), taskIndex);
     }
 
     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 = getAvroSchema(task);
+        AvroFormatColumnVisitor avroFormatColumnVisitor = new AvroFormatColumnVisitor(task, pageReader, avroSchema);
 
+        return new AvroFormatTransactionalPageOutput(producer, pageReader, avroFormatColumnVisitor, task.getTopic(), taskIndex);
+    }
+
+    private org.apache.avro.Schema getAvroSchema(PluginTask task)
+    {
         org.apache.avro.Schema avroSchema = null;
-        if (!task.getAvsc().isPresent() && !task.getAvscFile().isPresent() || task.getAvsc().isPresent() == task.getAvscFile().isPresent()) {
+        if (!task.getSchemaRegistryUrl().isPresent()) {
+            throw new ConfigException("avro_with_schema_registry format needs schema_registry_url");
+        }
+
+        if (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());
+            return avroSchema;
         }
         if (task.getAvscFile().isPresent()) {
             try {
                 avroSchema = new org.apache.avro.Schema.Parser().parse(task.getAvscFile().get());
+                return avroSchema;
             }
             catch (IOException e) {
                 e.printStackTrace();
@@ -261,71 +241,28 @@ public class KafkaOutputPlugin
             }
         }
 
-        PrimitiveIterator.OfLong randomLong = new Random().longs(1, Long.MAX_VALUE).iterator();
-
-        AtomicInteger counter = new AtomicInteger(0);
-        AtomicInteger recordLoggingCount = new AtomicInteger(1);
-
-        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.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();
-            }
+        SchemaRegistryClient schemaRegistryClient = getSchemaRegistryClient(task.getSchemaRegistryUrl().get());
+        SubjectNameStrategy subjectNameStrategy = new TopicNameStrategy();
+        String subjectName = subjectNameStrategy.subjectName(task.getTopic(), false, null);
+        try {
+            String schema = schemaRegistryClient.getLatestSchemaMetadata(subjectName).getSchema();
+            avroSchema = new org.apache.avro.Schema.Parser().parse(schema);
+            return avroSchema;
+        }
+        catch (IOException | RestClientException e) {
+            throw new ConfigException("cannot fetch latest schema from schema registry.", e);
+        }
+    }
 
-            @Override
-            public TaskReport commit()
-            {
-                return null;
-            }
-        };
+    private static final String MOCK_SCHEMA_REGISTRY_PREFIX = "mock://";
+    private SchemaRegistryClient getSchemaRegistryClient(String url)
+    {
+        if (url.startsWith(MOCK_SCHEMA_REGISTRY_PREFIX)) {
+            String mockScope = url.substring(MOCK_SCHEMA_REGISTRY_PREFIX.length());
+            return MockSchemaRegistry.getClientForScope(mockScope);
+        }
+        else {
+            return new CachedSchemaRegistryClient(url, SCHEMA_REGISTRY_IDENTITY_MAP_CAPACITY);
+        }
     }
 }

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"));
+    }
 }