1 files changed, 27 insertions, 24 deletions

diff --git a/lib/gitlab/prometheus_client.rb b/lib/gitlab/prometheus_client.rb
index aa94614bf18..10527972663 100644
--- a/lib/gitlab/prometheus_client.rb
+++ b/lib/gitlab/prometheus_client.rb
@@ -3,10 +3,10 @@ module Gitlab
 
   # Helper methods to interact with Prometheus network services & resources
   class PrometheusClient
-    attr_reader :api_url
+    attr_reader :rest_client, :headers
 
-    def initialize(api_url:)
-      @api_url = api_url
+    def initialize(rest_client)
+      @rest_client = rest_client
     end
 
     def ping
@@ -40,37 +40,40 @@ module Gitlab
     private
 
     def json_api_get(type, args = {})
-      get(join_api_url(type, args))
+      path = ['api', 'v1', type].join('/')
+      get(path, args)
+    rescue JSON::ParserError
+      raise PrometheusError, 'Parsing response failed'
     rescue Errno::ECONNREFUSED
       raise PrometheusError, 'Connection refused'
     end
 
-    def join_api_url(type, args = {})
-      url = URI.parse(api_url)
-    rescue URI::Error
-      raise PrometheusError, "Invalid API URL: #{api_url}"
-    else
-      url.path = [url.path.sub(%r{/+\z}, ''), 'api', 'v1', type].join('/')
-      url.query = args.to_query
-
-      url.to_s
-    end
-
-    def get(url)
-      handle_response(HTTParty.get(url))
+    def get(path, args)
+      response = rest_client[path].get(params: args)
+      handle_response(response)
     rescue SocketError
-      raise PrometheusError, "Can't connect to #{url}"
+      raise PrometheusError, "Can't connect to #{rest_client.url}"
     rescue OpenSSL::SSL::SSLError
-      raise PrometheusError, "#{url} contains invalid SSL data"
-    rescue HTTParty::Error
+      raise PrometheusError, "#{rest_client.url} contains invalid SSL data"
+    rescue RestClient::ExceptionWithResponse => ex
+      handle_exception_response(ex.response)
+    rescue RestClient::Exception
       raise PrometheusError, "Network connection error"
     end
 
     def handle_response(response)
-      if response.code == 200 && response['status'] == 'success'
-        response['data'] || {}
-      elsif response.code == 400
-        raise PrometheusError, response['error'] || 'Bad data received'
+      json_data = JSON.parse(response.body)
+      if response.code == 200 && json_data['status'] == 'success'
+        json_data['data'] || {}
+      else
+        raise PrometheusError, "#{response.code} - #{response.body}"
+      end
+    end
+
+    def handle_exception_response(response)
+      if response.code == 400
+        json_data = JSON.parse(response.body)
+        raise PrometheusError, json_data['error'] || 'Bad data received'
       else
         raise PrometheusError, "#{response.code} - #{response.body}"
       end