awslabs#38 Supporting gzip compression

lib/fluent/plugin/out_kinesis.rb

@@ -17,6 +17,7 @@
 require 'logger'
 require 'securerandom'
 require 'fluent/plugin/version'
+require 'zlib'
 
 module FluentPluginKinesis
   class OutputFilter < Fluent::BufferedOutput
@@ -56,6 +57,7 @@ class OutputFilter < Fluent::BufferedOutput
     config_param :order_events,           :bool,   default: false
     config_param :retries_on_putrecords,  :integer, default: 3
     config_param :use_yajl,               :bool,   default: false
+    config_param :gzip_compression,       :bool,   default: false
 
     config_param :debug, :bool, default: false
 
@@ -220,7 +222,11 @@ def get_key(name, record)
     end
 
     def build_data_to_put(data)
-      Hash[data.map{|k, v| [k.to_sym, v] }]
+      if @gzip_compression
+        Hash[data.map{|k, v| [k.to_sym, k=="data" ? Zlib::Deflate.deflate(v) : v] }]
+      else
+        Hash[data.map{|k, v| [k.to_sym, v] }]
+      end
     end
 
     def put_record_for_order_events(data_list)

test/plugin/test_out_kinesis.rb

@@ -30,6 +30,14 @@ def setup
     use_yajl true
   ]
 
+  CONFIG_WITH_COMPRESSION = CONFIG + %[
+    gzip_compression true
+  ]
+
+  CONFIG_YAJL_WITH_COMPRESSION = CONFIG_YAJL + %[
+    gzip_compression true
+  ]
+
   def create_driver(conf = CONFIG, tag='test')
     Fluent::Test::BufferedOutputTestDriver
       .new(FluentPluginKinesis::OutputFilter, tag).configure(conf)
@@ -237,7 +245,7 @@ def test_mode_configuration
 
 
   data("json"=>CONFIG, "yajl"=>CONFIG_YAJL)
-  def test_format(config)
+  def test_format_without_compression(config)
 
     d = create_driver(config)
 
@@ -276,6 +284,46 @@ def test_format(config)
     d.run
   end
 
+  data("json"=>CONFIG_WITH_COMPRESSION, "yajl"=>CONFIG_YAJL_WITH_COMPRESSION)
+  def test_format_with_compression(config)
+
+    d = create_driver(config)
+
+    data1 = {"test_partition_key"=>"key1","a"=>1,"time"=>"2011-01-02T13:14:15Z","tag"=>"test"}
+    data2 = {"test_partition_key"=>"key2","a"=>2,"time"=>"2011-01-02T13:14:15Z","tag"=>"test"}
+
+    time = Time.parse("2011-01-02 13:14:15 UTC").to_i
+    d.emit(data1, time)
+    d.emit(data2, time)
+
+    d.expect_format({
+      'data' => data1.to_json,
+      'partition_key' => 'key1' }.to_msgpack
+    )
+    d.expect_format({
+      'data' => data2.to_json,
+      'partition_key' => 'key2' }.to_msgpack
+    )
+
+    client = create_mock_client
+    client.describe_stream(stream_name: 'test_stream')
+    client.put_records(
+      stream_name: 'test_stream',
+      records: [
+        {
+          data: Zlib::Deflate.deflate(data1.to_json),
+          partition_key: 'key1'
+        },
+        {
+          data: Zlib::Deflate.deflate(data2.to_json),
+          partition_key: 'key2'
+        }
+      ]
+    ) { {} }
+
+    d.run
+  end
+
   def test_order_events
 
     d = create_driver(CONFIG + "\norder_events true")
@@ -349,7 +397,7 @@ def test_format_at_lowlevel_with_more_options
     )
   end
 
-  def test_multibyte_with_yajl
+  def test_multibyte_with_yajl_without_compression
 
     d = create_driver(CONFIG_YAJL)
 
@@ -380,6 +428,37 @@ def test_multibyte_with_yajl
     d.run
   end
 
+  def test_multibyte_with_yajl_with_compression
+
+    d = create_driver(CONFIG_YAJL_WITH_COMPRESSION)
+
+    data1 = {"test_partition_key"=>"key1","a"=>"\xE3\x82\xA4\xE3\x83\xB3\xE3\x82\xB9\xE3\x83\x88\xE3\x83\xBC\xE3\x83\xAB","time"=>"2011-01-02T13:14:15Z","tag"=>"test"}
+    json = Yajl.dump(data1)
+    data1["a"].force_encoding("ASCII-8BIT")
+
+    time = Time.parse("2011-01-02 13:14:15 UTC").to_i
+    d.emit(data1, time)
+
+    d.expect_format({
+      'data' => json,
+      'partition_key' => 'key1' }.to_msgpack
+    )
+
+    client = create_mock_client
+    client.describe_stream(stream_name: 'test_stream')
+    client.put_records(
+      stream_name: 'test_stream',
+      records: [
+        {
+          data: Zlib::Deflate.deflate(json),
+          partition_key: 'key1'
+        }
+      ]
+    ) { {} }
+
+    d.run
+  end
+
   def test_get_key
     d = create_driver
     assert_equal(