Add a proxy method to PrometheusClient

- Also refactor the get and json_api_get methods so that the get method
can be reused by the new proxy method.
- The new proxy method makes no changes to the request to the prometheus
server and response from the prometheus server. This allows it to be
used as a proxy to the Prometheus server, hence the name.

2 files changed, 115 insertions, 20 deletions

diff --git a/lib/gitlab/prometheus_client.rb b/lib/gitlab/prometheus_client.rb
index b4de7cd2bce..09609969afc 100644
--- a/lib/gitlab/prometheus_client.rb
+++ b/lib/gitlab/prometheus_client.rb
@@ -24,6 +24,19 @@ module Gitlab
       json_api_get('query', query: '1')
     end
 
+    def proxy(type, args)
+      path = api_path(type)
+      get(path, args)
+    rescue RestClient::ExceptionWithResponse => ex
+      if ex.response
+        ex.response
+      else
+        raise PrometheusClient::Error, "Network connection error"
+      end
+    rescue RestClient::Exception
+      raise PrometheusClient::Error, "Network connection error"
+    end
+
     def query(query, time: Time.now)
       get_result('vector') do
         json_api_get('query', query: query, time: time.to_f)
@@ -64,22 +77,16 @@ module Gitlab
 
     private
 
-    def json_api_get(type, args = {})
-      path = ['api', 'v1', type].join('/')
-      get(path, args)
-    rescue JSON::ParserError
-      raise PrometheusClient::Error, 'Parsing response failed'
-    rescue Errno::ECONNREFUSED
-      raise PrometheusClient::Error, 'Connection refused'
+    def api_path(type)
+      ['api', 'v1', type].join('/')
     end
 
-    def get(path, args)
-      response = rest_client[path].get(params: args)
+    def json_api_get(type, args = {})
+      path = api_path(type)
+      response = get(path, args)
       handle_response(response)
-    rescue SocketError
-      raise PrometheusClient::Error, "Can't connect to #{rest_client.url}"
-    rescue OpenSSL::SSL::SSLError
-      raise PrometheusClient::Error, "#{rest_client.url} contains invalid SSL data"
+    rescue JSON::ParserError
+      raise PrometheusClient::Error, 'Parsing response failed'
     rescue RestClient::ExceptionWithResponse => ex
       if ex.response
         handle_exception_response(ex.response)
@@ -90,6 +97,17 @@ module Gitlab
       raise PrometheusClient::Error, "Network connection error"
     end
 
+    def get(path, args)
+      response = rest_client[path].get(params: args)
+      response
+    rescue SocketError
+      raise PrometheusClient::Error, "Can't connect to #{rest_client.url}"
+    rescue OpenSSL::SSL::SSLError
+      raise PrometheusClient::Error, "#{rest_client.url} contains invalid SSL data"
+    rescue Errno::ECONNREFUSED
+      raise PrometheusClient::Error, 'Connection refused'
+    end
+
     def handle_response(response)
       json_data = JSON.parse(response.body)
       if response.code == 200 && json_data['status'] == 'success'
@@ -108,6 +126,8 @@ module Gitlab
       else
         raise PrometheusClient::Error, "#{response.code} - #{response.body}"
       end
+    rescue JSON::ParserError
+      raise PrometheusClient::Error, 'Parsing response failed'
     end
 
     def get_result(expected_type)
diff --git a/spec/lib/gitlab/prometheus_client_spec.rb b/spec/lib/gitlab/prometheus_client_spec.rb
index 2517ee71f24..d3a28b53da9 100644
--- a/spec/lib/gitlab/prometheus_client_spec.rb
+++ b/spec/lib/gitlab/prometheus_client_spec.rb
@@ -60,15 +60,13 @@ describe Gitlab::PrometheusClient do
   end
 
   describe 'failure to reach a provided prometheus url' do
-    let(:prometheus_url) {"https://prometheus.invalid.example.com"}
+    let(:prometheus_url) {"https://prometheus.invalid.example.com/api/v1/query?query=1"}
 
-    subject { described_class.new(RestClient::Resource.new(prometheus_url)) }
-
-    context 'exceptions are raised' do
+    shared_examples 'exceptions are raised' do
       it 'raises a Gitlab::PrometheusClient::Error error when a SocketError is rescued' do
         req_stub = stub_prometheus_request_with_exception(prometheus_url, SocketError)
 
-        expect { subject.send(:get, '/', {}) }
+        expect { subject }
           .to raise_error(Gitlab::PrometheusClient::Error, "Can't connect to #{prometheus_url}")
         expect(req_stub).to have_been_requested
       end
@@ -76,7 +74,7 @@ describe Gitlab::PrometheusClient do
       it 'raises a Gitlab::PrometheusClient::Error error when a SSLError is rescued' do
         req_stub = stub_prometheus_request_with_exception(prometheus_url, OpenSSL::SSL::SSLError)
 
-        expect { subject.send(:get, '/', {}) }
+        expect { subject }
           .to raise_error(Gitlab::PrometheusClient::Error, "#{prometheus_url} contains invalid SSL data")
         expect(req_stub).to have_been_requested
       end
@@ -84,11 +82,23 @@ describe Gitlab::PrometheusClient do
       it 'raises a Gitlab::PrometheusClient::Error error when a RestClient::Exception is rescued' do
         req_stub = stub_prometheus_request_with_exception(prometheus_url, RestClient::Exception)
 
-        expect { subject.send(:get, '/', {}) }
+        expect { subject }
           .to raise_error(Gitlab::PrometheusClient::Error, "Network connection error")
         expect(req_stub).to have_been_requested
       end
     end
+
+    context 'ping' do
+      subject { described_class.new(RestClient::Resource.new(prometheus_url)).ping }
+
+      it_behaves_like 'exceptions are raised'
+    end
+
+    context 'proxy' do
+      subject { described_class.new(RestClient::Resource.new(prometheus_url)).proxy('query', { query: '1' }) }
+
+      it_behaves_like 'exceptions are raised'
+    end
   end
 
   describe '#query' do
@@ -258,4 +268,69 @@ describe Gitlab::PrometheusClient do
       it { is_expected.to eq(step) }
     end
   end
+
+  describe 'proxy' do
+    context 'query' do
+      let(:prometheus_query) { prometheus_cpu_query('env-slug') }
+      let(:query_url) { prometheus_query_url(prometheus_query) }
+
+      around do |example|
+        Timecop.freeze { example.run }
+      end
+
+      it 'returns full response from the API call' do
+        req_stub = stub_prometheus_request(query_url, body: prometheus_value_body('vector'))
+
+        response = subject.proxy('query', { query: prometheus_query })
+        json_response = JSON.parse(response.body)
+
+        expect(response.code).to eq(200)
+        expect(json_response).to eq({
+          'status' => 'success',
+          'data' => {
+            'resultType' => 'vector',
+            'result' => [{ "metric" => {}, "value" => [1488772511.004, "0.000041021495238095323"] }]
+          }
+        })
+        expect(req_stub).to have_been_requested
+      end
+    end
+
+    context 'query_range' do
+      let(:prometheus_query) { prometheus_memory_query('env-slug') }
+      let(:query_url) { prometheus_query_range_url(prometheus_query, start: 2.hours.ago) }
+
+      around do |example|
+        Timecop.freeze { example.run }
+      end
+
+      it 'returns full response' do
+        req_stub = stub_prometheus_request(query_url, body: prometheus_values_body('vector'))
+
+        response = subject.proxy('query_range', {
+          query: prometheus_query,
+          start: 2.hours.ago.to_f,
+          end: Time.now.to_f,
+          step: 60
+        })
+        json_response = JSON.parse(response.body)
+
+        expect(response.code).to eq(200)
+        expect(json_response).to eq({
+          "status" => "success",
+          "data" => {
+            "resultType" => "vector",
+            "result" => [{
+              "metric" => {},
+              "values" => [
+                [1488758662.506, "0.00002996364761904785"],
+                [1488758722.506, "0.00003090239047619091"]
+              ]
+            }]
+          }
+        })
+        expect(req_stub).to have_been_requested
+      end
+    end
+  end
 end