embulk-output-kafka 0.1.3 → 0.1.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: fb1edf4ce79bd490f0f01662547d36c2e6d2dd6e
-  data.tar.gz: 44b9affb8e9d1385314274b5d1ac083fd5a4848e
+  metadata.gz: d02c758c81651e54d09421f19c1ba4970ddb04c4
+  data.tar.gz: 18c7cbb212fa721f11427250c6132d6cf937a292
 SHA512:
-  metadata.gz: 5c76d2cca3d141b7f44208449ba092ee732993183de6e6f19ec4e5280f302f0e3a4b5746414332103aa492b0cb7ea7df55159fa936248839bb68a8280ecfc060
-  data.tar.gz: 28dcd4f04aceec97e78c427b2b859266b9adb5147ae8cd28a63c5df86840d5c4051af049b225eefe01d410007afdd17222ec96fc3525cf0cd28e60c30414f6a2
+  metadata.gz: 4f6cd307587c3edc55b87404b708a9a21ad904e9685ca1e28eebb662a1cead956175173ce6d8d04821d1a5d3866c143beb10a67bb67ad8e71416b207cbdf31da
+  data.tar.gz: 13de4897add59955cab7a7adbdb252a1a843b8042f94c5f23faca63d4bda1e094c4003f0486244c3e50ec30f2ec5004728bce496713ab8051edf6fed4ff636f8

data/README.md

@@ -18,6 +18,7 @@
 - **avsc**: inline avro schema config (json, default: `null`)
 - **ignore_columns**: remove columns from output  (array(string), default: `[]`)
 - **key_column_name**: use column value as record key (string, default: `null`, if this parameter is null, set random number as record key, and it can use column in `ignore_columns`)
+- **partition_column_name**: use column value as partition id (string, default: `null`, this value is prefer to `key_column_name`, and if partition_column value is null, use key_column for partitioning)
 - **record_batch_size**: kafka producer record batch size (integer, default: `1000`)
 - **acks**: kafka producer require acks (string, default: `"1"`)
 - **retries**: kafka producer max retry count (integer, default: `1`)

data/build.gradle

@@ -17,7 +17,7 @@ configurations {
     provided
 }
 
-version = "0.1.3"
+version = "0.1.4"
 
 sourceCompatibility = 1.8
 targetCompatibility = 1.8

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

@@ -17,13 +17,24 @@ import java.util.stream.Collectors;
 public class AvroFormatColumnVisitor extends KafkaOutputColumnVisitor
 {
     private Schema avroSchema;
-    public GenericRecord genericRecord;
+    private GenericRecord genericRecord;
 
-    public AvroFormatColumnVisitor(KafkaOutputPlugin.PluginTask task, PageReader pageReader, Schema avroSchema, GenericRecord genericRecord)
+    AvroFormatColumnVisitor(KafkaOutputPlugin.PluginTask task, PageReader pageReader, Schema avroSchema)
     {
         super(task, pageReader);
         this.avroSchema = avroSchema;
-        this.genericRecord = genericRecord;
+    }
+
+    GenericRecord getGenericRecord()
+    {
+        return genericRecord;
+    }
+
+    @Override
+    void reset()
+    {
+        super.reset();
+        this.genericRecord = new GenericData.Record(avroSchema);
     }
 
     @Override

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

@@ -13,14 +13,25 @@ import java.time.format.DateTimeFormatter;
 public class JsonFormatColumnVisitor extends KafkaOutputColumnVisitor
 {
     private ObjectMapper objectMapper;
-    public ObjectNode jsonNode;
+    private ObjectNode jsonNode;
 
     private static DateTimeFormatter timestampFormatter = DateTimeFormatter.ISO_INSTANT;
 
-    public JsonFormatColumnVisitor(KafkaOutputPlugin.PluginTask task, PageReader pageReader, ObjectMapper objectMapper)
+    JsonFormatColumnVisitor(KafkaOutputPlugin.PluginTask task, PageReader pageReader, ObjectMapper objectMapper)
     {
         super(task, pageReader);
         this.objectMapper = objectMapper;
+    }
+
+    ObjectNode getJsonNode()
+    {
+        return jsonNode;
+    }
+
+    @Override
+    void reset()
+    {
+        super.reset();
         this.jsonNode = objectMapper.createObjectNode();
     }
 
@@ -42,6 +53,8 @@ public class JsonFormatColumnVisitor extends KafkaOutputColumnVisitor
     @Override
     public void longColumn(Column column)
     {
+        super.longColumn(column);
+
         if (isIgnoreColumn(column)) {
             return;
         }
@@ -58,6 +71,8 @@ public class JsonFormatColumnVisitor extends KafkaOutputColumnVisitor
     @Override
     public void doubleColumn(Column column)
     {
+        super.doubleColumn(column);
+
         if (isIgnoreColumn(column)) {
             return;
         }
@@ -68,12 +83,13 @@ public class JsonFormatColumnVisitor extends KafkaOutputColumnVisitor
         }
 
         jsonNode.put(column.getName(), pageReader.getDouble(column));
-        super.doubleColumn(column);
     }
 
     @Override
     public void stringColumn(Column column)
     {
+        super.stringColumn(column);
+
         if (isIgnoreColumn(column)) {
             return;
         }
@@ -84,7 +100,6 @@ public class JsonFormatColumnVisitor extends KafkaOutputColumnVisitor
         }
 
         jsonNode.put(column.getName(), pageReader.getString(column));
-        super.stringColumn(column);
     }
 
     @Override

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

@@ -6,32 +6,57 @@ import org.embulk.spi.PageReader;
 
 public abstract class KafkaOutputColumnVisitor implements ColumnVisitor
 {
-    KafkaOutputPlugin.PluginTask task;
+    private KafkaOutputPlugin.PluginTask task;
     PageReader pageReader;
+    private String partitionColumnName;
 
-    public Object recordKey = null;
-    public String topicName = null;
+    private Object recordKey = null;
+    private String topicName = null;
+    private Integer partition = null;
 
-    public KafkaOutputColumnVisitor(KafkaOutputPlugin.PluginTask task, PageReader pageReader)
+    KafkaOutputColumnVisitor(KafkaOutputPlugin.PluginTask task, PageReader pageReader)
     {
         this.task = task;
         this.pageReader = pageReader;
+        this.partitionColumnName = task.getPartitionColumnName().orElse(null);
     }
 
-    void setRecordKey(Column column, Object value)
+    Object getRecordKey()
+    {
+        return recordKey;
+    }
+
+    private void setRecordKey(Column column, Object value)
     {
         if (task.getKeyColumnName().isPresent() && task.getKeyColumnName().get().equals(column.getName())) {
             recordKey = value;
         }
     }
 
-    void setTopicName(Column column, String value)
+    String getTopicName()
+    {
+        return topicName;
+    }
+
+    private void setTopicName(Column column, String value)
     {
         if (task.getTopicColumn().isPresent() && task.getTopicColumn().get().equals(column.getName())) {
             topicName = value;
         }
     }
 
+    Integer getPartition()
+    {
+        return partition;
+    }
+
+    void reset()
+    {
+        this.recordKey = null;
+        this.topicName = null;
+        this.partition = null;
+    }
+
     boolean isIgnoreColumn(Column column)
     {
         return task.getIgnoreColumns().stream().anyMatch(name -> name.equals(column.getName()));
@@ -41,7 +66,12 @@ public abstract class KafkaOutputColumnVisitor implements ColumnVisitor
     public void longColumn(Column column)
     {
         if (!pageReader.isNull(column)) {
-            setRecordKey(column, pageReader.getLong(column));
+            long value = pageReader.getLong(column);
+            setRecordKey(column, value);
+
+            if (partitionColumnName != null && partitionColumnName.equals(column.getName())) {
+                partition = Long.valueOf(value).intValue();
+            }
         }
     }
 

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

@@ -4,7 +4,10 @@ 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 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;
@@ -15,7 +18,6 @@ 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;
@@ -32,7 +34,11 @@ import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
 import java.util.PrimitiveIterator;
+import java.util.Properties;
 import java.util.Random;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
 
 public class KafkaOutputPlugin
@@ -96,6 +102,10 @@ 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")
         public int getRecordBatchSize();
@@ -124,17 +134,31 @@ public class KafkaOutputPlugin
     private static ObjectMapper objectMapper = new ObjectMapper();
     private Logger logger = LoggerFactory.getLogger(getClass());
 
+    private AdminClient getKafkaAdminClient(PluginTask task)
+    {
+        Properties properties = new Properties();
+        properties.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, task.getBrokers());
+        AdminClient adminClient = AdminClient.create(properties);
+        return adminClient;
+    }
+
     @Override
     public ConfigDiff transaction(ConfigSource config,
             Schema schema, int taskCount,
             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();
     }
@@ -179,22 +203,24 @@ public class KafkaOutputPlugin
         AtomicInteger recordLoggingCount = new AtomicInteger(1);
 
         return new TransactionalPageOutput() {
+            private JsonFormatColumnVisitor columnVisitor = new JsonFormatColumnVisitor(task, pageReader, objectMapper);
+
             @Override
             public void add(Page page)
             {
                 pageReader.setPage(page);
                 while (pageReader.nextRecord()) {
-                    JsonFormatColumnVisitor columnVisitor = new JsonFormatColumnVisitor(task, pageReader, objectMapper);
+                    columnVisitor.reset();
 
                     pageReader.getSchema().visitColumns(columnVisitor);
 
-                    Object recordKey = columnVisitor.recordKey;
+                    Object recordKey = columnVisitor.getRecordKey();
                     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);
+                    String targetTopic = columnVisitor.getTopicName() != null ? columnVisitor.getTopicName() : task.getTopic();
+                    ProducerRecord<Object, ObjectNode> producerRecord = new ProducerRecord<>(targetTopic, recordKey, columnVisitor.getJsonNode());
                     producer.send(producerRecord, (metadata, exception) -> {
                         if (exception != null) {
                             logger.error("produce error", exception);
@@ -269,23 +295,25 @@ public class KafkaOutputPlugin
         final org.apache.avro.Schema finalAvroSchema = avroSchema;
         return new TransactionalPageOutput()
         {
+            private AvroFormatColumnVisitor columnVisitor = new AvroFormatColumnVisitor(task, pageReader, finalAvroSchema);
+
             @Override
             public void add(Page page)
             {
                 pageReader.setPage(page);
                 while (pageReader.nextRecord()) {
-                    AvroFormatColumnVisitor columnVisitor = new AvroFormatColumnVisitor(task, pageReader, finalAvroSchema, new GenericData.Record(finalAvroSchema));
+                    columnVisitor.reset();
 
                     pageReader.getSchema().visitColumns(columnVisitor);
 
-                    Object recordKey = columnVisitor.recordKey;
+                    Object recordKey = columnVisitor.getRecordKey();
                     if (recordKey == null) {
                         recordKey = randomLong.next();
                     }
 
-                    String targetTopic = columnVisitor.topicName != null ? columnVisitor.topicName : task.getTopic();
+                    String targetTopic = columnVisitor.getTopicName() != null ? columnVisitor.getTopicName() : task.getTopic();
 
-                    ProducerRecord<Object, Object> producerRecord = new ProducerRecord<>(targetTopic, recordKey, columnVisitor.genericRecord);
+                    ProducerRecord<Object, Object> producerRecord = new ProducerRecord<>(targetTopic, columnVisitor.getPartition(), recordKey, columnVisitor.getGenericRecord());
                     producer.send(producerRecord, (metadata, exception) -> {
                         if (exception != null) {
                             logger.error("produce error", exception);

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: embulk-output-kafka
 version: !ruby/object:Gem::Version
-  version: 0.1.3
+  version: 0.1.4
 platform: ruby
 authors:
 - joker1007
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2019-09-13 00:00:00.000000000 Z
+date: 2019-10-29 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   requirement: !ruby/object:Gem::Requirement
@@ -54,7 +54,7 @@ files:
 - classpath/common-config-5.3.0.jar
 - classpath/common-utils-5.3.0.jar
 - classpath/commons-compress-1.18.jar
-- classpath/embulk-output-kafka-0.1.3.jar
+- classpath/embulk-output-kafka-0.1.4.jar
 - classpath/jackson-annotations-2.9.0.jar
 - classpath/jackson-core-2.9.9.jar
 - classpath/jackson-databind-2.9.9.jar