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.



6
7
8
9
10
 
# File 'lib/dbx/databricks/sql_response.rb', line 6

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.



12
13
14
 
# File 'lib/dbx/databricks/sql_response.rb', line 12

def body
  @body
end

#data_array ⇒ Object

Returns the value of attribute data_array.



12
13
14
 
# File 'lib/dbx/databricks/sql_response.rb', line 12

def data_array
  @data_array
end

#raw_response ⇒ Object

Returns the value of attribute raw_response.



12
13
14
 
# File 'lib/dbx/databricks/sql_response.rb', line 12

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) 


45
46
47
48
 
# File 'lib/dbx/databricks/sql_response.rb', line 45

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>) 


87
88
89
 
# File 'lib/dbx/databricks/sql_response.rb', line 87

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

#error_message ⇒ String | nil

Dig out the error message from the response body.

Returns:

  • (String | nil) 


71
72
73
 
# File 'lib/dbx/databricks/sql_response.rb', line 71

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>) 


94
95
96
 
# File 'lib/dbx/databricks/sql_response.rb', line 94

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) 


65
66
67
 
# File 'lib/dbx/databricks/sql_response.rb', line 65

def failed?
  status == "FAILED"
end

#more_chunks? ⇒ Boolean

Determine if the response contains multiple chunks.

Returns:

  • (Boolean) 


32
33
34
35
 
# File 'lib/dbx/databricks/sql_response.rb', line 32

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) 


39
40
41
 
# File 'lib/dbx/databricks/sql_response.rb', line 39

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

#parse_body ⇒ Object

Parse the response body as JSON.



17
18
19
20
21
 
# File 'lib/dbx/databricks/sql_response.rb', line 17

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) 


60
61
62
 
# File 'lib/dbx/databricks/sql_response.rb', line 60

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

#results(symbolize_keys: false) ⇒ Array<Hash>

Return the results of the query as an array of hashes with string keys.

Returns:

  • (Array<Hash>) 


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

def results(symbolize_keys: false)
  return [] if failed?

  data_array.map do |row|
    hash = {}
    columns.each do |column|
      key = symbolize_keys ? column["name"].to_sym : column["name"]
      hash[key] = row[column["position"]]
    end
    hash
  end
end

#statement_id ⇒ String | nil

Dig out the statement_id from the response body.

Returns:

  • (String | nil) 


25
26
27
 
# File 'lib/dbx/databricks/sql_response.rb', line 25

def statement_id
  body["statement_id"]
end

#status ⇒ String

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

Returns:

  • (String) 


77
78
79
80
81
 
# File 'lib/dbx/databricks/sql_response.rb', line 77

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) 


53
54
55
 
# File 'lib/dbx/databricks/sql_response.rb', line 53

def success?
  status == "SUCCEEDED"
end