[PLAT-6535] Add support for custom headers

CHANGELOG.md

@@ -1,5 +1,7 @@
 ## Clickhouse CHANGELOG
 
+* Added support for sending custom headers
+
 ### Version 0.1.10 (January 13, 2017)
 
 * Fixed `erubis` dependency once and for all

Gemfile

@@ -1,3 +1,5 @@
 source "https://rubygems.org"
 
 gemspec
+
+gem "csv"

clickhouse.gemspec

@@ -17,7 +17,7 @@ Gem::Specification.new do |gem|
   gem.licenses      = ["MIT"]
 
   gem.add_dependency "bundler", ">= 1.13.4"
-  gem.add_dependency "faraday"
+  gem.add_dependency "faraday", "~> 0.17.6"
   gem.add_dependency "pond"
   gem.add_dependency "activesupport", ">= 4.1.8"
 

lib/clickhouse/connection/client.rb

@@ -3,6 +3,15 @@ class Connection
     module Client
       include ActiveSupport::NumberHelper
 
+      def resolve_headers
+        headers_config = @config[:headers]
+        case headers_config
+        when Proc then headers_config.call || {}
+        when Hash then headers_config
+        else {}
+        end
+      end
+
       def connect!
         ping! unless connected?
       end
@@ -62,7 +71,10 @@ def request(method, query, body = nil, optimized = false, query_params = {})
         query = query.strip
         start = Time.now
 
-        response = client.send(method, path(query, query_params), body)
+        headers = resolve_headers
+        response = client.send(method, path(query, query_params), body) do |req|
+          req.headers.merge!(headers)
+        end
         status = response.status
         t1 = Time.now
         if optimized
@@ -104,7 +116,10 @@ def request_post(method, query, body = nil, query_params = {})
         query = query.strip
         start = Time.now
 
-        response = client.send(method, path(query, query_params), body)
+        headers = resolve_headers
+        response = client.send(method, path(query, query_params), body) do |req|
+          req.headers.merge!(headers)
+        end
         status = response.status
         duration = Time.now - start
         query, format = Utils.extract_format(query)

test/test_helper.rb

@@ -2,7 +2,7 @@
 
 require "minitest"
 require "minitest/autorun"
-require "mocha/setup"
+require "mocha/minitest"
 
 def path(path)
   File.expand_path "../../#{path}", __FILE__

test/test_helper/minitest.rb

@@ -1,12 +1,12 @@
-class MiniTest::Test
+class Minitest::Test
   def teardown
     Clickhouse.instance_variables.each do |name|
       Clickhouse.instance_variable_set name, nil
     end
   end
 end
 
-class MiniTest::Spec
+class Minitest::Spec
   def assert_query(expected, actual)
     assert_equal(expected.strip.gsub(/^\s+/, ""), actual)
   end

test/unit/connection/test_client.rb

@@ -2,7 +2,8 @@
 
 module Unit
   module Connection
-    class TestClient < MiniTest::Test
+    class TestClient < Minitest::Test
+      EMPTY_JSON = '{"rows":0,"data":[],"meta":[],"statistics":{"elapsed":0.001,"rows_read":0,"bytes_read":0}}'
 
       class Connection < SimpleConnection
         include Clickhouse::Connection::Client
@@ -62,17 +63,19 @@ class Connection < SimpleConnection
         describe "#get" do
           it "sends a GET request the server" do
             @connection.instance_variable_set :@client, (client = mock)
-            client.expects(:get).with("/?query=foo&output_format_write_statistics=1", nil).returns(stub(:status => 200, :body => ""))
+            client.expects(:get).with("/?query=foo&output_format_write_statistics=1", nil).yields(stub(:headers => {})).returns(stub(:status => 200, :body => EMPTY_JSON))
             @connection.stubs(:log)
+            @connection.stubs(:parse_body).returns({})
             @connection.get("foo")
           end
         end
 
         describe "#post" do
           it "sends a POST request the server" do
             @connection.instance_variable_set :@client, (client = mock)
-            client.expects(:post).with("/?query=foo&output_format_write_statistics=1", "body").returns(stub(:status => 200, :body => ""))
+            client.expects(:post).with("/?query=foo&output_format_write_statistics=1", "body").yields(stub(:headers => {})).returns(stub(:status => 200, :body => EMPTY_JSON))
             @connection.stubs(:log)
+            @connection.stubs(:parse_body).returns({})
             @connection.post("foo", "body")
           end
         end
@@ -85,21 +88,22 @@ class Connection < SimpleConnection
           it "connects to the server first" do
             @connection.instance_variable_set :@client, (client = mock)
             @connection.expects(:connect!)
-            client.stubs(:get).returns(stub(:status => 200, :body => ""))
+            @connection.stubs(:parse_body).returns({})
+            client.stubs(:get).yields(stub(:headers => {})).returns(stub(:status => 200, :body => EMPTY_JSON))
             @connection.send :request, :get, "/", "query"
           end
 
           it "queries the server returning the response" do
             @connection.instance_variable_set :@client, (client = mock)
-            client.expects(:get).with("/?query=SELECT+1&output_format_write_statistics=1", nil).returns(stub(:status => 200, :body => ""))
-            @connection.expects(:parse_body).returns(data = mock)
+            client.expects(:get).with("/?query=SELECT+1&output_format_write_statistics=1", nil).yields(stub(:headers => {})).returns(stub(:status => 200, :body => ""))
+            @connection.expects(:parse_body).returns(data = {})
             assert_equal data, @connection.send(:request, :get, "SELECT 1")
           end
 
           describe "when not receiving status 200" do
             it "raises a Clickhouse::QueryError" do
               @connection.instance_variable_set :@client, (client = mock)
-              client.expects(:get).with("/?query=SELECT+1&output_format_write_statistics=1", nil).returns(stub(:status => 500, :body => ""))
+              client.expects(:get).with("/?query=SELECT+1&output_format_write_statistics=1", nil).yields(stub(:headers => {})).returns(stub(:status => 500, :body => ""))
               assert_raises Clickhouse::QueryError do
                 @connection.send(:request, :get, "SELECT 1")
               end
@@ -121,8 +125,9 @@ class Connection < SimpleConnection
               {"meta": []}
             JSON
             @connection.instance_variable_set :@client, (client = mock)
-            client.expects(:get).with("/?query=SELECT+1+FORMAT+JSONCompact&output_format_write_statistics=1", nil).returns(stub(:status => 200, :body => json))
-            assert_equal({"meta" => []}, @connection.send(:request, :get, "SELECT 1 FORMAT JSONCompact"))
+            client.expects(:get).with("/?query=SELECT+1+FORMAT+JSONCompact&output_format_write_statistics=1", nil).yields(stub(:headers => {})).returns(stub(:status => 200, :body => json))
+            result = @connection.send(:request, :get, "SELECT 1 FORMAT JSONCompact")
+            assert_equal [], result["meta"]
           end
         end
 
@@ -131,8 +136,8 @@ class Connection < SimpleConnection
             it "includes the database in the querystring" do
               @connection.instance_variable_get(:@config)[:database] = "system"
               @connection.instance_variable_set(:@client, (client = mock))
-              client.expects(:get).with("/?database=system&query=SELECT+1&output_format_write_statistics=1", nil).returns(stub(:status => 200, :body => ""))
-              @connection.expects(:parse_body).returns(data = mock)
+              client.expects(:get).with("/?database=system&query=SELECT+1&output_format_write_statistics=1", nil).yields(stub(:headers => {})).returns(stub(:status => 200, :body => ""))
+              @connection.expects(:parse_body).returns(data = {})
               assert_equal data, @connection.send(:request, :get, "SELECT 1")
             end
           end
@@ -165,9 +170,9 @@ class Connection < SimpleConnection
           it "parses the statistics" do
             @connection.stubs(:log)
             @connection.instance_variable_set :@client, (client = mock)
-            Time.expects(:now).returns(1882).twice
+            Time.stubs(:now).returns(1882)
 
-            client.expects(:get).with("/?query=SELECT+1+FORMAT+JSONCompact&output_format_write_statistics=1", nil).returns(stub(:status => 200, :body => @json))
+            client.expects(:get).with("/?query=SELECT+1+FORMAT+JSONCompact&output_format_write_statistics=1", nil).yields(stub(:headers => {})).returns(stub(:status => 200, :body => @json))
             @connection.expects(:write_log).with(
               0, "SELECT 1", {
                 "elapsed" => "188.2ms",
@@ -184,9 +189,9 @@ class Connection < SimpleConnection
 
           it "write the expected logs" do
             @connection.instance_variable_set :@client, (client = mock)
-            Time.expects(:now).returns(1882).twice
+            Time.stubs(:now).returns(1882)
 
-            client.expects(:get).with("/?query=SELECT+1+FORMAT+JSONCompact&output_format_write_statistics=1", nil).returns(stub(:status => 200, :body => @json))
+            client.expects(:get).with("/?query=SELECT+1+FORMAT+JSONCompact&output_format_write_statistics=1", nil).yields(stub(:headers => {})).returns(stub(:status => 200, :body => @json))
             log = "\n \e[1m\e[35mSQL (0.0ms)\e\e[0m  SELECT 1;\e\n  \e[1m\e[36m1947 rows in set. Elapsed: 188.2ms. Processed: 1.98 thousand rows, 1.96 KB (10.53 thousand rows/s, 10.39 KB/s)\e[0m "
 
             @connection.expects(:log).with(:debug, log)
@@ -195,12 +200,96 @@ class Connection < SimpleConnection
 
           describe "#number_to_human_duration" do
             it "returns in seconds when more than 1 seconds" do
-              assert_equal "2.0s", @connection.send(:number_to_human_duration, 2)
+              assert_equal "2s", @connection.send(:number_to_human_duration, 2)
             end
           end
         end
-      end
 
+        describe "#resolve_headers" do
+          it "returns empty hash when headers config is nil" do
+            @connection.instance_variable_get(:@config).delete(:headers)
+            assert_equal({}, @connection.resolve_headers)
+          end
+
+          it "returns the hash when headers config is a Hash" do
+            headers = {"X-Custom-Header" => "value"}
+            @connection.instance_variable_get(:@config)[:headers] = headers
+            assert_equal headers, @connection.resolve_headers
+          end
+
+          it "calls the proc and returns result when headers config is a Proc" do
+            headers = {"X-Dynamic-Header" => "dynamic-value"}
+            @connection.instance_variable_get(:@config)[:headers] = -> { headers }
+            assert_equal headers, @connection.resolve_headers
+          end
+
+          it "returns empty hash when proc returns nil" do
+            @connection.instance_variable_get(:@config)[:headers] = -> { nil }
+            assert_equal({}, @connection.resolve_headers)
+          end
+        end
+
+        describe "headers in requests" do
+          def json_response
+            '{"rows":0,"data":[],"meta":[],"statistics":{"elapsed":0.001,"rows_read":0,"bytes_read":0}}'
+          end
+
+          def connection_with_stubs(headers_config:, &block)
+            captured_headers = nil
+
+            stubs = Faraday::Adapter::Test::Stubs.new do |stub|
+              stub.get(/.*/) do |env|
+                captured_headers = env.request_headers.to_h
+                [200, {}, json_response]
+              end
+              stub.post(/.*/) do |env|
+                captured_headers = env.request_headers.to_h
+                [200, {}, json_response]
+              end
+            end
+
+            clickhouse = Clickhouse::Connection.new(url: "http://localhost:8123", headers: headers_config)
+            faraday_client = Faraday.new(url: "http://localhost:8123") { |f| f.adapter :test, stubs }
+            clickhouse.instance_variable_set(:@client, faraday_client)
+
+            yield clickhouse, -> { captured_headers }
+          end
+
+          describe "GET request" do
+            it "sends headers from Hash config to the server" do
+              connection_with_stubs(headers_config: {"X-Test" => "test-value"}) do |clickhouse, get_headers|
+                clickhouse.query("SELECT 1 FORMAT JSON")
+                assert_equal "test-value", get_headers.call["X-Test"]
+              end
+            end
+
+            it "sends headers from Proc config to the server" do
+              call_count = 0
+              headers_proc = -> {
+                call_count += 1
+                {"X-Dynamic" => "value-#{call_count}"}
+              }
+
+              connection_with_stubs(headers_config: headers_proc) do |clickhouse, get_headers|
+                clickhouse.query("SELECT 1 FORMAT JSON")
+                assert_equal "value-1", get_headers.call["X-Dynamic"]
+
+                clickhouse.query("SELECT 2 FORMAT JSON")
+                assert_equal "value-2", get_headers.call["X-Dynamic"]
+              end
+            end
+          end
+
+          describe "POST request" do
+            it "sends headers from config to the server" do
+              connection_with_stubs(headers_config: {"X-Post-Header" => "post-value"}) do |clickhouse, get_headers|
+                clickhouse.query_post("SELECT 1 FORMAT JSON")
+                assert_equal "post-value", get_headers.call["X-Post-Header"]
+              end
+            end
+          end
+        end
+      end
     end
   end
 end