Merge pull request #10 from HatsuneMiku3939/feature/header-support

request/response header support

fluent-plugin-http-pull.gemspec

@@ -3,7 +3,7 @@ $LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
 
 Gem::Specification.new do |spec|
   spec.name    = "fluent-plugin-http-pull"
-  spec.version = "0.4.0"
+  spec.version = "0.5.0"
   spec.authors = ["filepang"]
   spec.email   = ["filepang@gmail.com"]
 

lib/fluent/plugin/in_http_pull.rb

@@ -45,11 +45,30 @@ def initialize
       desc 'password of basic auth'
       config_param :password, :string, default: nil
 
+      config_section :response_header, param_name: :response_headers, multi: true do
+        desc 'The name of header to cature from response'
+        config_param :header, :string
+      end
+
+      config_section :request_header, param_name: :request_headers, multi: true do
+        desc 'The name of request header'
+        config_param :header, :string
+
+        desc 'The value of request header'
+        config_param :value, :string
+      end
+
       def configure(conf)
         compat_parameters_convert(conf, :parser)
         super
 
         @parser = parser_create unless @status_only
+        @_request_headers = @request_headers.map do |section|
+          header = section["header"]
+          value = section["value"]
+
+          [header.to_sym, value]
+        end.to_h
       end
 
       def start
@@ -67,10 +86,20 @@ def on_timer
           request_options[:proxy] = @proxy if @proxy
           request_options[:user] = @user if @user
           request_options[:password] = @password if @password
+          request_options[:headers] = @_request_headers unless @request_headers.empty?
 
           res = RestClient::Request.execute request_options
+
           record["status"] = res.code
           record["body"] = res.body
+
+          record["header"] = {} unless @response_headers.empty?
+          @response_headers.each do |section|
+            name = section["header"]
+            symbolize_name = name.downcase.gsub(/-/, '_').to_sym
+
+            record["header"][name] = res.headers[symbolize_name]
+          end
         rescue StandardError => err
           if err.respond_to? :http_code
             record["status"] = err.http_code || 0

test/helper/stub_server.rb

@@ -11,6 +11,7 @@ def initialize
     @server.mount_proc '/internal_error', &method(:internal_error)
     @server.mount_proc '/redirect', &method(:redirect)
     @server.mount_proc '/protected', &method(:protected)
+    @server.mount_proc '/custom_header', &method(:custom_header)
   end
 
   def start
@@ -81,4 +82,12 @@ def protected(req, res)
     res['Content-Type'] = 'application/json'
     res.body = '{ "status": "OK" }'
   end
+
+  def custom_header(req, res)
+    res.header["HATSUNE-MIKU"] = req["HATSUNE-MIKU"] if req["HATSUNE-MIKU"]
+
+    res.status = 200
+    res['Content-Type'] = 'application/json'
+    res.body = '{ "status": "OK" }'
+  end
 end

test/plugin/test_in_http_pull.rb

@@ -453,6 +453,109 @@ class HttpPullInputTest < Test::Unit::TestCase
     end
   end
 
+  sub_test_case "capture response header" do
+    TEST_INTERVAL_3_RES_HEADER_CONFIG = %[
+      tag test
+      url http://127.0.0.1:3939
+
+      interval 3s
+      format json
+
+      <response_header>
+        header Content-Type
+      </response_header>
+    ]
+
+    TEST_INTERVAL_3_RES_MULTI_HEADER_CONFIG = %[
+      tag test
+      url http://127.0.0.1:3939
+
+      interval 3s
+      format json
+
+      <response_header>
+        header Content-Type
+      </response_header>
+
+      <response_header>
+        header Content-Length
+      </response_header>
+    ]
+
+    test 'interval 3 with single header' do
+      d = create_driver TEST_INTERVAL_3_RES_HEADER_CONFIG
+      assert_equal("test", d.instance.tag)
+      assert_equal(3, d.instance.interval)
+
+      d.run(timeout: 8) do
+        sleep 7
+      end
+      assert_equal(2, d.events.size)
+
+      d.events.each do |tag, time, record|
+        assert_equal("test", tag)
+
+        assert_equal({"url"=>"http://127.0.0.1:3939","status"=>200,"message"=>{"status"=>"OK"},"header"=>{"Content-Type"=>"application/json"}}, record)
+        assert(time.is_a?(Fluent::EventTime))
+      end
+    end
+
+    test 'interval 3 with multiple header' do
+      d = create_driver TEST_INTERVAL_3_RES_MULTI_HEADER_CONFIG
+      assert_equal("test", d.instance.tag)
+      assert_equal(3, d.instance.interval)
+
+      d.run(timeout: 8) do
+        sleep 7
+      end
+      assert_equal(2, d.events.size)
+
+      d.events.each do |tag, time, record|
+        assert_equal("test", tag)
+
+        assert_equal({"url"=>"http://127.0.0.1:3939","status"=>200,"message"=>{"status"=>"OK"},"header"=>{"Content-Type"=>"application/json","Content-Length"=>"18"}}, record)
+        assert(time.is_a?(Fluent::EventTime))
+      end
+    end
+  end
+
+  sub_test_case "custom request header" do
+    TEST_INTERVAL_3_CUSTOM_HEADER_CONFIG = %[
+      tag test
+      url http://127.0.0.1:3939/custom_header
+
+      interval 3s
+      format json
+
+      <request_header>
+        header HATSUNE-MIKU
+        value 3939
+      </request_header>
+
+      <response_header>
+        header HATSUNE-MIKU
+      </response_header>
+    ]
+
+    test 'interval 3 with single header' do
+      d = create_driver TEST_INTERVAL_3_CUSTOM_HEADER_CONFIG
+      assert_equal("test", d.instance.tag)
+      assert_equal(3, d.instance.interval)
+
+      d.run(timeout: 8) do
+        sleep 7
+      end
+      assert_equal(2, d.events.size)
+
+      d.events.each do |tag, time, record|
+        assert_equal("test", tag)
+
+        assert_equal({"url"=>"http://127.0.0.1:3939/custom_header","status"=>200,"message"=>{"status"=>"OK"},"header"=>{"HATSUNE-MIKU"=>"3939"}}, record)
+        assert(time.is_a?(Fluent::EventTime))
+      end
+    end
+  end
+
   private
 
   def create_driver(conf)