Class: DatabricksSQLResponse

Inherits:
Object
  • Object
show all
Defined in:
lib/dbx/databricks/sql_response.rb

Overview

This class represents a response from the Databricks SQL API. It is used by DatabricksSQL to handle http failures and parse the response body.

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(http_response) ⇒ DatabricksSQLResponse

Returns a new instance of DatabricksSQLResponse.



8
9
10
11
12
 
# File 'lib/dbx/databricks/sql_response.rb', line 8

def initialize(http_response)
  self.raw_response = http_response
  self.body = parse_body
  self.data_array = extract_data_array
end

Instance Attribute Details

#body ⇒ Object

Returns the value of attribute body.



14
15
16
 
# File 'lib/dbx/databricks/sql_response.rb', line 14

def body
  @body
end

#data_array ⇒ Object

Returns the value of attribute data_array.



14
15
16
 
# File 'lib/dbx/databricks/sql_response.rb', line 14

def data_array
  @data_array
end

#raw_response ⇒ Object

Returns the value of attribute raw_response.



14
15
16
 
# File 'lib/dbx/databricks/sql_response.rb', line 14

def raw_response
  @raw_response
end

Instance Method Details

#add_chunk_to_data(chunk_response) ⇒ Array

Combine the data from the chunk response into the data from the original response.

Returns:

  • (Array) 


47
48
49
50
 
# File 'lib/dbx/databricks/sql_response.rb', line 47

def add_chunk_to_data(chunk_response)
  chunk_data_array = chunk_response.data_array
  self.data_array = [*data_array, *chunk_data_array]
end

#columns ⇒ Array<String>

Dig out the columns array from the response body.

Returns:

  • (Array<String>) 


89
90
91
 
# File 'lib/dbx/databricks/sql_response.rb', line 89

def columns
  body.dig("manifest", "schema", "columns") || []
end

#error_message ⇒ String | nil

Dig out the error message from the response body.

Returns:

  • (String | nil) 


73
74
75
 
# File 'lib/dbx/databricks/sql_response.rb', line 73

def error_message
  body.dig("status", "error", "message")
end

#extract_data_array ⇒ Array<Array>

Dig out values array for the queried data. Chunks have a simpler hash structure than initial SQL responses.

Returns:

  • (Array<Array>) 


96
97
98
 
# File 'lib/dbx/databricks/sql_response.rb', line 96

def extract_data_array
  body.dig("result", "data_array") || body["data_array"] || []
end

#failed? ⇒ Boolean

Determine if the response from the API has failed.

Returns:

  • (Boolean) 


67
68
69
 
# File 'lib/dbx/databricks/sql_response.rb', line 67

def failed?
  status == "FAILED"
end

#more_chunks? ⇒ Boolean

Determine if the response contains multiple chunks.

Returns:

  • (Boolean) 


34
35
36
37
 
# File 'lib/dbx/databricks/sql_response.rb', line 34

def more_chunks?
  chunk_count = body&.dig("manifest", "total_chunk_count")&.to_i
  chunk_count && chunk_count > 1
end

#next_chunk ⇒ String | nil

Dig out the next_chunk_internal_link from the response body.

Returns:

  • (String | nil) 


41
42
43
 
# File 'lib/dbx/databricks/sql_response.rb', line 41

def next_chunk
  body.dig("result", "next_chunk_internal_link")
end

#parse_body ⇒ Object

Parse the response body as JSON.



19
20
21
22
23
 
# File 'lib/dbx/databricks/sql_response.rb', line 19

def parse_body
  return {} unless raw_response.is_a?(Net::HTTPSuccess)

  @body = JSON.parse(raw_response.body)
end

#pending? ⇒ Boolean

Determine if the response from the API is still executing. PENDING means the warehouse is starting up RUNNING means the query is still executing

Returns:

  • (Boolean) 


62
63
64
 
# File 'lib/dbx/databricks/sql_response.rb', line 62

def pending?
  %w[PENDING RUNNING].include?(status)
end

#results ⇒ Array<Hash>

Return the results of the query as an array of hashes.

Returns:

  • (Array<Hash>) 


102
103
104
105
106
107
108
109
110
111
112
 
# File 'lib/dbx/databricks/sql_response.rb', line 102

def results
  return [] if failed?

  data_array.map do |row|
    hash = {}
    columns.each do |column|
      hash[column["name"]] = row[column["position"]]
    end
    hash
  end
end

#statement_id ⇒ String | nil

Dig out the statement_id from the response body.

Returns:

  • (String | nil) 


27
28
29
 
# File 'lib/dbx/databricks/sql_response.rb', line 27

def statement_id
  body["statement_id"]
end

#status ⇒ String

Dig out the status of the query from the response body.

Returns:

  • (String) 


79
80
81
82
83
 
# File 'lib/dbx/databricks/sql_response.rb', line 79

def status
  return "FAILED" unless raw_response.is_a?(Net::HTTPSuccess)

  body.dig("status", "state")
end

#success? ⇒ Boolean

Determine if the response from the API has succeeded.

Returns:

  • (Boolean) 


55
56
57
 
# File 'lib/dbx/databricks/sql_response.rb', line 55

def success?
  status == "SUCCEEDED"
end