diff --git a/CHANGELOG.md b/CHANGELOG.md index 0c8c653..cf89f59 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/Gemfile b/Gemfile index b4e2a20..02afd72 100644 --- a/Gemfile +++ b/Gemfile @@ -1,3 +1,5 @@ source "https://rubygems.org" gemspec + +gem "csv" diff --git a/clickhouse.gemspec b/clickhouse.gemspec index 68f24af..91b390c 100644 --- a/clickhouse.gemspec +++ b/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" diff --git a/lib/clickhouse/connection/client.rb b/lib/clickhouse/connection/client.rb index 9f6610f..80c2727 100644 --- a/lib/clickhouse/connection/client.rb +++ b/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) diff --git a/test/test_helper.rb b/test/test_helper.rb index 89e71d3..bd55a49 100644 --- a/test/test_helper.rb +++ b/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__ diff --git a/test/test_helper/minitest.rb b/test/test_helper/minitest.rb index a15925c..61f2445 100644 --- a/test/test_helper/minitest.rb +++ b/test/test_helper/minitest.rb @@ -1,4 +1,4 @@ -class MiniTest::Test +class Minitest::Test def teardown Clickhouse.instance_variables.each do |name| Clickhouse.instance_variable_set name, nil @@ -6,7 +6,7 @@ def teardown end end -class MiniTest::Spec +class Minitest::Spec def assert_query(expected, actual) assert_equal(expected.strip.gsub(/^\s+/, ""), actual) end diff --git a/test/unit/connection/test_client.rb b/test/unit/connection/test_client.rb index 26de6e8..2834c60 100644 --- a/test/unit/connection/test_client.rb +++ b/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,8 +63,9 @@ 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 @@ -71,8 +73,9 @@ class Connection < SimpleConnection 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