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.

4 files changed, 419 insertions, 22 deletions

diff --git a/app/services/prometheus/proxy_service.rb b/app/services/prometheus/proxy_service.rb
new file mode 100644
index 00000000000..c5d2b84878b
--- /dev/null
+++ b/app/services/prometheus/proxy_service.rb
@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+module Prometheus
+  class ProxyService < BaseService
+    include ReactiveCaching
+    include Gitlab::Utils::StrongMemoize
+
+    self.reactive_cache_key = ->(service) { service.cache_key }
+    self.reactive_cache_lease_timeout = 30.seconds
+    self.reactive_cache_refresh_interval = 30.seconds
+    self.reactive_cache_lifetime = 1.minute
+    self.reactive_cache_worker_finder = ->(_id, *args) { from_cache(*args) }
+
+    attr_accessor :proxyable, :method, :path, :params
+
+    PROXY_SUPPORT = {
+      'query' => {
+        method: ['GET'],
+        params: %w(query time timeout)
+      },
+      'query_range' => {
+        method: ['GET'],
+        params: %w(query start end step timeout)
+      }
+    }.freeze
+
+    def self.from_cache(proxyable_class_name, proxyable_id, method, path, params)
+      proxyable_class = begin
+        proxyable_class_name.constantize
+      rescue NameError
+        nil
+      end
+      return unless proxyable_class
+
+      proxyable = proxyable_class.find(proxyable_id)
+
+      new(proxyable, method, path, params)
+    end
+
+    # proxyable can be any model which responds to .prometheus_adapter
+    # like Environment.
+    def initialize(proxyable, method, path, params)
+      @proxyable = proxyable
+      @path = path
+
+      # Convert ActionController::Parameters to hash because reactive_cache_worker
+      # does not play nice with ActionController::Parameters.
+      @params = filter_params(params, path).to_hash
+
+      @method = method
+    end
+
+    def id
+      nil
+    end
+
+    def execute
+      return cannot_proxy_response unless can_proxy?
+      return no_prometheus_response unless can_query?
+
+      with_reactive_cache(*cache_key) do |result|
+        result
+      end
+    end
+
+    def calculate_reactive_cache(proxyable_class_name, proxyable_id, method, path, params)
+      return no_prometheus_response unless can_query?
+
+      response = prometheus_client_wrapper.proxy(path, params)
+
+      success(http_status: response.code, body: response.body)
+    rescue Gitlab::PrometheusClient::Error => err
+      service_unavailable_response(err)
+    end
+
+    def cache_key
+      [@proxyable.class.name, @proxyable.id, @method, @path, @params]
+    end
+
+    private
+
+    def service_unavailable_response(exception)
+      error(exception.message, :service_unavailable)
+    end
+
+    def no_prometheus_response
+      error('No prometheus server found', :service_unavailable)
+    end
+
+    def cannot_proxy_response
+      error('Proxy support for this API is not available currently')
+    end
+
+    def prometheus_adapter
+      strong_memoize(:prometheus_adapter) do
+        @proxyable.prometheus_adapter
+      end
+    end
+
+    def prometheus_client_wrapper
+      prometheus_adapter&.prometheus_client_wrapper
+    end
+
+    def can_query?
+      prometheus_adapter&.can_query?
+    end
+
+    def filter_params(params, path)
+      params.slice(*PROXY_SUPPORT.dig(path, :params))
+    end
+
+    def can_proxy?
+      PROXY_SUPPORT.dig(@path, :method)&.include?(@method)
+    end
+  end
+end
diff --git a/lib/gitlab/prometheus_client.rb b/lib/gitlab/prometheus_client.rb
index b4de7cd2bce..f13156f898e 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,14 @@ 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 RestClient::ExceptionWithResponse => ex
       if ex.response
         handle_exception_response(ex.response)
@@ -90,8 +95,18 @@ module Gitlab
       raise PrometheusClient::Error, "Network connection error"
     end
 
+    def get(path, args)
+      rest_client[path].get(params: args)
+    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)
+      json_data = parse_json(response.body)
       if response.code == 200 && json_data['status'] == 'success'
         json_data['data'] || {}
       else
@@ -103,7 +118,7 @@ module Gitlab
       if response.code == 200 && response['status'] == 'success'
         response['data'] || {}
       elsif response.code == 400
-        json_data = JSON.parse(response.body)
+        json_data = parse_json(response.body)
         raise PrometheusClient::QueryError, json_data['error'] || 'Bad data received'
       else
         raise PrometheusClient::Error, "#{response.code} - #{response.body}"
@@ -114,5 +129,11 @@ module Gitlab
       data = yield
       data['result'] if data['resultType'] == expected_type
     end
+
+    def parse_json(response_body)
+      JSON.parse(response_body)
+    rescue JSON::ParserError
+      raise PrometheusClient::Error, 'Parsing response failed'
+    end
   end
 end
diff --git a/spec/lib/gitlab/prometheus_client_spec.rb b/spec/lib/gitlab/prometheus_client_spec.rb
index 2517ee71f24..f15ae83a02c 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,59 @@ describe Gitlab::PrometheusClient do
       it { is_expected.to eq(step) }
     end
   end
+
+  describe 'proxy' do
+    context 'get API' 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
+
+      context 'when response status code is 200' do
+        it 'returns response object' 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 'when response status code is not 200' do
+        it 'returns response object' do
+          req_stub = stub_prometheus_request(query_url, status: 400, body: { error: 'error' })
+
+          response = subject.proxy('query', { query: prometheus_query })
+          json_response = JSON.parse(response.body)
+
+          expect(req_stub).to have_been_requested
+          expect(response.code).to eq(400)
+          expect(json_response).to eq('error' => 'error')
+        end
+      end
+
+      context 'when RestClient::Exception is raised' do
+        before do
+          stub_prometheus_request_with_exception(query_url, RestClient::Exception)
+        end
+
+        it 'raises PrometheusClient::Error' do
+          expect { subject.proxy('query', { query: prometheus_query }) }.to(
+            raise_error(Gitlab::PrometheusClient::Error, 'Network connection error')
+          )
+        end
+      end
+    end
+  end
 end
diff --git a/spec/services/prometheus/proxy_service_spec.rb b/spec/services/prometheus/proxy_service_spec.rb
new file mode 100644
index 00000000000..4bdb20de4c9
--- /dev/null
+++ b/spec/services/prometheus/proxy_service_spec.rb
@@ -0,0 +1,195 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+describe Prometheus::ProxyService do
+  include ReactiveCachingHelpers
+
+  set(:project) { create(:project) }
+  set(:environment) { create(:environment, project: project) }
+
+  describe '#initialize' do
+    let(:params) { ActionController::Parameters.new(query: '1').permit! }
+
+    it 'initializes attributes' do
+      result = described_class.new(environment, 'GET', 'query', params)
+
+      expect(result.proxyable).to eq(environment)
+      expect(result.method).to eq('GET')
+      expect(result.path).to eq('query')
+      expect(result.params).to eq('query' => '1')
+    end
+
+    it 'converts ActionController::Parameters into hash' do
+      result = described_class.new(environment, 'GET', 'query', params)
+
+      expect(result.params).to be_an_instance_of(Hash)
+    end
+
+    context 'with unknown params' do
+      let(:params) { ActionController::Parameters.new(query: '1', other_param: 'val').permit! }
+
+      it 'filters unknown params' do
+        result = described_class.new(environment, 'GET', 'query', params)
+
+        expect(result.params).to eq('query' => '1')
+      end
+    end
+  end
+
+  describe '#execute' do
+    let(:prometheus_adapter) { instance_double(PrometheusService) }
+    let(:params) { ActionController::Parameters.new(query: '1').permit! }
+
+    subject { described_class.new(environment, 'GET', 'query', params) }
+
+    context 'when prometheus_adapter is nil' do
+      before do
+        allow(environment).to receive(:prometheus_adapter).and_return(nil)
+      end
+
+      it 'returns error' do
+        expect(subject.execute).to eq(
+          status: :error,
+          message: 'No prometheus server found',
+          http_status: :service_unavailable
+        )
+      end
+    end
+
+    context 'when prometheus_adapter cannot query' do
+      before do
+        allow(environment).to receive(:prometheus_adapter).and_return(prometheus_adapter)
+        allow(prometheus_adapter).to receive(:can_query?).and_return(false)
+      end
+
+      it 'returns error' do
+        expect(subject.execute).to eq(
+          status: :error,
+          message: 'No prometheus server found',
+          http_status: :service_unavailable
+        )
+      end
+    end
+
+    context 'cannot proxy' do
+      subject { described_class.new(environment, 'POST', 'garbage', params) }
+
+      it 'returns error' do
+        expect(subject.execute).to eq(
+          message: 'Proxy support for this API is not available currently',
+          status: :error
+        )
+      end
+    end
+
+    context 'with caching', :use_clean_rails_memory_store_caching do
+      let(:return_value) { { 'http_status' => 200, 'body' => 'body' } }
+
+      let(:opts) do
+        [environment.class.name, environment.id, 'GET', 'query', { 'query' => '1' }]
+      end
+
+      before do
+        allow(environment).to receive(:prometheus_adapter)
+          .and_return(prometheus_adapter)
+        allow(prometheus_adapter).to receive(:can_query?).and_return(true)
+      end
+
+      context 'when value present in cache' do
+        before do
+          stub_reactive_cache(subject, return_value, opts)
+        end
+
+        it 'returns cached value' do
+          result = subject.execute
+
+          expect(result[:http_status]).to eq(return_value[:http_status])
+          expect(result[:body]).to eq(return_value[:body])
+        end
+      end
+
+      context 'when value not present in cache' do
+        it 'returns nil' do
+          expect(ReactiveCachingWorker)
+            .to receive(:perform_async)
+            .with(subject.class, subject.id, *opts)
+
+          result = subject.execute
+
+          expect(result).to eq(nil)
+        end
+      end
+    end
+
+    context 'call prometheus api' do
+      let(:prometheus_client) { instance_double(Gitlab::PrometheusClient) }
+
+      before do
+        synchronous_reactive_cache(subject)
+
+        allow(environment).to receive(:prometheus_adapter)
+          .and_return(prometheus_adapter)
+        allow(prometheus_adapter).to receive(:can_query?).and_return(true)
+        allow(prometheus_adapter).to receive(:prometheus_client_wrapper)
+          .and_return(prometheus_client)
+      end
+
+      context 'connection to prometheus server succeeds' do
+        let(:rest_client_response) { instance_double(RestClient::Response) }
+        let(:prometheus_http_status_code) { 400 }
+
+        let(:response_body) do
+          '{"status":"error","errorType":"bad_data","error":"parse error at char 1: no expression found in input"}'
+        end
+
+        before do
+          allow(prometheus_client).to receive(:proxy).and_return(rest_client_response)
+
+          allow(rest_client_response).to receive(:code)
+            .and_return(prometheus_http_status_code)
+          allow(rest_client_response).to receive(:body).and_return(response_body)
+        end
+
+        it 'returns the http status code and body from prometheus' do
+          expect(subject.execute).to eq(
+            http_status: prometheus_http_status_code,
+            body: response_body,
+            status: :success
+          )
+        end
+      end
+
+      context 'connection to prometheus server fails' do
+        context 'prometheus client raises Gitlab::PrometheusClient::Error' do
+          before do
+            allow(prometheus_client).to receive(:proxy)
+              .and_raise(Gitlab::PrometheusClient::Error, 'Network connection error')
+          end
+
+          it 'returns error' do
+            expect(subject.execute).to eq(
+              status: :error,
+              message: 'Network connection error',
+              http_status: :service_unavailable
+            )
+          end
+        end
+      end
+    end
+  end
+
+  describe '.from_cache' do
+    it 'initializes an instance of ProxyService class' do
+      result = described_class.from_cache(
+        environment.class.name, environment.id, 'GET', 'query', { 'query' => '1' }
+      )
+
+      expect(result).to be_an_instance_of(described_class)
+      expect(result.proxyable).to eq(environment)
+      expect(result.method).to eq('GET')
+      expect(result.path).to eq('query')
+      expect(result.params).to eq('query' => '1')
+    end
+  end
+end