lib/drill.rb in drill-sergeant-0.1.0 vs lib/drill.rb in drill-sergeant-0.1.1

@@ -3,32 +3,43 @@
 require "drill/version"
 
 class Drill
   class Error < StandardError; end
 
+  HEADERS = {
+    "Content-Type" => "application/json",
+    "Accept" => "application/json"
+  }
+
   def initialize(url: nil)
     url ||= ENV["DRILL_URL"] || "http://localhost:8047"
-    @uri = URI.parse("#{url}/query.json")
+    # strip trailing slash if exists
+    @uri = URI.parse("#{url.sub(/\/\z/, "")}/query.json")
     @http = Net::HTTP.new(@uri.host, @uri.port)
   end
 
   def query(statement)
-    header = {"Content-Type" => "application/json", "Accept" => "application/json"}
     data = {
       queryType: "sql",
       query: statement
     }
 
     begin
-      response = @http.post(@uri.request_uri, data.to_json, header)
+      response = @http.post(@uri.request_uri, data.to_json, HEADERS)
     rescue Errno::ECONNREFUSED => e
       raise Drill::Error, e.message
     end
 
     body = JSON.parse(response.body)
     if body["errorMessage"]
       raise Drill::Error, body["errorMessage"].split("\n")[0]
     end
 
-    body["rows"]
+    # 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
 end