lib/drill-sergeant.rb in drill-sergeant-0.1.3 vs lib/drill-sergeant.rb in drill-sergeant-0.1.4

@@ -4,10 +4,11 @@
 require "net/http"
 
 # modules
 require "drill/version"
 
+# TODO make module and move to client
 class Drill
   class Error < StandardError; end
 
   HEADERS = {
     "Content-Type" => "application/json",
@@ -15,32 +16,23 @@
   }
 
   def initialize(url: nil, open_timeout: 3, read_timeout: nil)
     url ||= ENV["DRILL_URL"] || "http://localhost:8047"
     # remove trailing slash
-    @uri = URI.parse(url.sub(/\/\z/, ""))
+    @uri = URI.parse(url.chomp("/"))
     @http = Net::HTTP.new(@uri.host, @uri.port)
     @http.use_ssl = true if @uri.scheme == "https"
     @http.open_timeout = open_timeout if open_timeout
     @http.read_timeout = read_timeout if read_timeout
   end
 
   def query(statement)
-    data = {
-      queryType: "sql",
-      query: statement
-    }
-
-    body = post("query.json", data)
-
-    # return columns in order
-    result = []
-    columns = body["columns"]
-    body["rows"].each do |row|
-      result << columns.each_with_object({}) { |c, memo| memo[c] = row[c] }
+    options = {}
+    if Gem::Version.new(server_version) >= Gem::Version.new("1.19.0")
+      options["drill.exec.http.rest.errors.verbose"] = true
     end
-    result
+    run_query(statement, options)
   end
 
   def profiles(query_id = nil)
     path = query_id ? "profiles/#{escape_path(query_id)}.json" : "profiles.json"
     get(path)
@@ -62,12 +54,39 @@
 
   def options
     get("options.json")
   end
 
+  def server_version
+    @server_version ||= run_query("SELECT version FROM sys.version", {})[0]["version"]
+  end
+
   private
 
+  def run_query(statement, options)
+    data = {
+      queryType: "sql",
+      query: statement,
+      options: options
+    }
+
+    body = post("query.json", data)
+
+    # errors return 200 with Drill 1.19+
+    if body["queryState"] != "COMPLETED"
+      raise Error, body["errorMessage"] || "Bad state: #{body["queryState"]}"
+    end
+
+    # return columns in order
+    result = []
+    columns = body["columns"]
+    body["rows"].each do |row|
+      result << columns.each_with_object({}) { |c, memo| memo[c] = row[c] }
+    end
+    result
+  end
+
   def get(path)
     handle_response do
       @http.get("#{@uri.request_uri}#{path}", HEADERS)
     end
   end
@@ -80,16 +99,16 @@
 
   def handle_response
     begin
       response = yield
     rescue Errno::ECONNREFUSED => e
-      raise Drill::Error, e.message
+      raise Error, e.message
     end
 
     unless response.kind_of?(Net::HTTPSuccess)
-      body = JSON.parse(response.body) rescue nil
+      body = JSON.parse(response.body) rescue {}
       message = body["errorMessage"] || "Bad response: #{response.code}"
-      raise Drill::Error, message
+      raise Error, message
     end
 
     JSON.parse(response.body)
   end