out_kafka2: Skip events with an invalid partition

lib/fluent/plugin/out_kafka2.rb

@@ -98,6 +98,7 @@ class Fluent::Kafka2Output < Output
     include Fluent::KafkaPluginUtil::AwsIamSettings
     include Fluent::KafkaPluginUtil::SSLSettings
     include Fluent::KafkaPluginUtil::SaslSettings
+    include Fluent::KafkaPluginUtil::PartitionSettings
 
     def initialize
       super
@@ -351,6 +352,7 @@ def write(chunk)
             partition_key = (@exclude_partition_key ? record.delete(@partition_key_key) : record[@partition_key_key]) || @default_partition_key
             partition = (@exclude_partition ? record.delete(@partition_key) : record[@partition_key]) || @default_partition
             message_key = (@exclude_message_key ? record.delete(@message_key_key) : record[@message_key_key]) || @default_message_key
+            partition = coerce_partition(partition) unless partition.nil?
 
             if mutate_headers
               headers = base_headers.clone

test/plugin/test_out_kafka2.rb

@@ -27,6 +27,32 @@ def create_driver(conf = config, tag='test')
     Fluent::Test::Driver::Output.new(Fluent::Kafka2Output).configure(conf)
   end
 
+  class DummyProducer
+    attr_reader :produced
+
+    # Coerces like ruby-kafka's Producer#produce, which is where an invalid
+    # partition raises today
+    def produce(value, key:, partition_key:, partition:, headers:, create_time:, topic:)
+      @produced ||= []
+      @produced << {value: value && value.to_s, key: key && key.to_s,
+                    partition: partition && Integer(partition)}
+    end
+
+    def deliver_messages
+    end
+
+    def clear_buffer
+    end
+  end
+
+  def create_chunk(records, tag: 'test')
+    metadata = Fluent::Plugin::Buffer::Metadata.new(nil, tag, nil)
+    chunk = Fluent::Plugin::Buffer::MemoryChunk.new(metadata)
+    chunk.extend(Fluent::ChunkMessagePackEventStreamer)
+    chunk.append(records.map { |record| [event_time, record].to_msgpack })
+    chunk
+  end
+
   def test_configure
     assert_nothing_raised(Fluent::ConfigError) {
       create_driver(base_config)
@@ -92,6 +118,32 @@ def test_configure_unsupported_scram_mechanism
     }
   end
 
+  data("non numeric partition" => "not-a-number",
+       "hash partition" => {"x" => 1},
+       "negative partition" => -2,
+       "out of int32 range" => 2**31)
+  def test_write_skips_event_with_invalid_partition(partition)
+    d = create_driver
+    producer = DummyProducer.new
+    stub(d.instance).get_producer { producer }
+
+    assert_nothing_raised {
+      d.instance.write(create_chunk([{"a" => "b"}, {"a" => "c", "partition" => partition}, {"a" => "d"}]))
+    }
+
+    assert_equal ['{"a":"b"}', '{"a":"d"}'], producer.produced.collect { |message| message[:value] }
+  end
+
+  def test_write_keeps_event_with_zero_padded_partition
+    d = create_driver
+    producer = DummyProducer.new
+    stub(d.instance).get_producer { producer }
+
+    d.instance.write(create_chunk([{"a" => "b", "partition" => "010"}]))
+
+    assert_equal [10], producer.produced.collect { |message| message[:partition] }
+  end
+
   data("crc32" => "crc32",
        "murmur2" => "murmur2")
   def test_partitioner_hash_function(data)