Add support for gzip compression (#215)

* add support for gzip compression * change test Co-authored-by: Aleksander Dudek <[email protected]>
awslabs · Nov 14, 2021 · f309128 · f309128
1 parent a896721
commit f309128
Show file tree

Hide file tree

Showing 4 changed files with 44 additions and 1 deletion.
diff --git a/README.md b/README.md
@@ -353,7 +353,7 @@ Also, there are some format related options below:
 If your record contains a field whose string should be sent to Amazon Kinesis directly (without formatter), use this parameter to specify the field. In that case, other fields than **data_key** are thrown away and never sent to Amazon Kinesis. Default `nil`, which means whole record will be formatted and sent.
 
 ### compression
-Specifing compression way for data of each record. Current accepted options are `zlib`. Otherwise, no compression will be preformed.
+Specifying compression way for data of each record. Current accepted options are `zlib` and `gzip`. Otherwise, no compression will be preformed.
 
 ### log_truncate_max_size
 Integer, default 1024. When emitting the log entry, the message will be truncated by this size to avoid infinite loop when the log is also sent to Kinesis. The value 0 means no truncation.

diff --git a/lib/fluent/plugin/kinesis.rb b/lib/fluent/plugin/kinesis.rb
@@ -17,6 +17,7 @@
 require 'fluent/plugin/output'
 require 'fluent/plugin/kinesis_helper/client'
 require 'fluent/plugin/kinesis_helper/api'
+require 'fluent/plugin/kinesis_helper/compression'
 require 'zlib'
 
 module Fluent
@@ -121,6 +122,8 @@ def compressor_create
         case @compression
         when "zlib"
           ->(data) { Zlib::Deflate.deflate(data) }
+        when "gzip"
+          ->(data) { Gzip.compress(data) }
         else
           ->(data) { data }
         end

diff --git a/lib/fluent/plugin/kinesis_helper/compression.rb b/lib/fluent/plugin/kinesis_helper/compression.rb
@@ -0,0 +1,27 @@
+require "stringio"
+require "zlib"
+
+class Stream < StringIO
+  def initialize(*)
+    super
+    set_encoding "BINARY"
+  end
+
+  def close
+    rewind;
+  end
+end
+
+class Gzip
+  def self.compress(string, level = Zlib::DEFAULT_COMPRESSION, strategy = Zlib::DEFAULT_STRATEGY)
+    output = Stream.new
+    gz = Zlib::GzipWriter.new(output, level, strategy)
+    gz.write(string)
+    gz.close
+    output.string
+  end
+
+  def self.decompress(string)
+    Zlib::GzipReader.wrap(StringIO.new(string), &:read)
+  end
+end
diff --git a/test/plugin/test_kinesis.rb b/test/plugin/test_kinesis.rb
@@ -14,6 +14,7 @@
 
 require_relative '../helper'
 require 'fluent/plugin/kinesis'
+require 'fluent/plugin/kinesis_helper/compression'
 
 module Fluent
   module Plugin
@@ -87,6 +88,18 @@ def test_format_compression(data)
     assert_equal expected, MessagePack.unpack(result).first
   end
 
+  data(
+    'gzip' => ['gzip', Gzip.compress("foo2")],
+    )
+  def test_format_compression_gzip(data)
+    compression, expected = data
+    d = create_driver(default_config + "data_key a\ncompression #{compression}")
+    driver_run(d, [{"a"=>"foo2"}])
+    result = d.formatted.first
+    unpacked = MessagePack.unpack(result).first
+    assert_equal Gzip.decompress(expected), Gzip.decompress(unpacked)
+  end
+
   data(
     'not_exceeded'      => [{"a"=>"a"*30}, ["a".b*30].to_msgpack],
     'exceeded'          => [{"a"=>"a"*31}, ''],