Class: DatabricksSql::Client

Inherits:
Object
  • Object
show all
Defined in:
lib/databricks_sql/client.rb

Constant Summary collapse

DEFAULT_STATEMENT_PATH = 
"/api/2.0/sql/statements"
TERMINAL_STATUSES = 
%w[SUCCEEDED FAILED CANCELED CLOSED].freeze

Instance Attribute Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(host: nil, token: nil, warehouse_id: nil, timeout: nil, open_timeout: nil, statement_path: nil) ⇒ Client

Returns a new instance of Client.



15
16
17
18
19
20
21
22
23
24
25
26
27
 
# File 'lib/databricks_sql/client.rb', line 15

def initialize(host: nil, token: nil, warehouse_id: nil, timeout: nil, open_timeout: nil, statement_path: nil)
  global_config = DatabricksSql.configuration

  @host = normalize_host(host || global_config.host)
  @token = (token || global_config.token).to_s
  @warehouse_id = (warehouse_id || global_config.warehouse_id).to_s
  @timeout = timeout || global_config.timeout || 30
  @open_timeout = open_timeout || global_config.open_timeout || 10
  @statement_path = statement_path || global_config.statement_path || DEFAULT_STATEMENT_PATH
  @external_link_require_https = global_config.external_link_require_https != false
  @external_link_allowed_hosts = normalize_allowed_hosts(global_config.external_link_allowed_hosts)
  validate_configuration!
end

Instance Attribute Details

#host ⇒ Object (readonly)

Returns the value of attribute host.



13
14
15
 
# File 'lib/databricks_sql/client.rb', line 13

def host
  @host
end

#warehouse_id ⇒ Object (readonly)

Returns the value of attribute warehouse_id.



13
14
15
 
# File 'lib/databricks_sql/client.rb', line 13

def warehouse_id
  @warehouse_id
end

Instance Method Details

#cancel_statement(statement_id:, request_timeout: nil) ⇒ Object 



100
101
102
103
104
105
 
# File 'lib/databricks_sql/client.rb', line 100

def cancel_statement(statement_id:, request_timeout: nil)
  path = "#{@statement_path}/#{statement_id}/cancel"
  request_json(:post, path, payload: {}, request_timeout: request_timeout)
rescue HTTPError
  nil
end

#execute_statement(statement:, parameters: nil, format: "JSON_ARRAY", disposition: "INLINE", catalog: nil, schema: nil, column_schema: nil, wait_timeout: "10s", on_wait_timeout: "CONTINUE", poll_interval: 1.0, max_wait: 300, request_timeout: nil, cancel_on_timeout: false, auto_fetch_chunks: true, byte_limit: nil, row_limit: nil, query_tags: nil) ⇒ Object 



29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
 
# File 'lib/databricks_sql/client.rb', line 29

def execute_statement(statement:, parameters: nil, format: "JSON_ARRAY", disposition: "INLINE", catalog: nil,
  schema: nil, column_schema: nil, wait_timeout: "10s", on_wait_timeout: "CONTINUE",
  poll_interval: 1.0, max_wait: 300, request_timeout: nil, cancel_on_timeout: false,
  auto_fetch_chunks: true, byte_limit: nil, row_limit: nil, query_tags: nil)
  response = execute_statement_async(
    statement: statement,
    parameters: parameters,
    format: format,
    disposition: disposition,
    catalog: catalog,
    schema: schema,
    wait_timeout: wait_timeout,
    on_wait_timeout: on_wait_timeout,
    request_timeout: request_timeout,
    byte_limit: byte_limit,
    row_limit: row_limit,
    query_tags: query_tags
  )

  statement_id = response.fetch("statement_id")
  if terminal_status?(statement_status(response))
    response = hydrate_terminal_response(
      response,
      statement_id: statement_id,
      request_timeout: request_timeout
    )

    return build_result(response, format: format, disposition: disposition, column_schema: column_schema,
                                  auto_fetch_chunks: auto_fetch_chunks,
                                  request_timeout: request_timeout)
  end

  wait_for_statement(
    statement_id: statement_id,
    format: format,
    disposition: disposition,
    column_schema: column_schema,
    poll_interval: poll_interval,
    max_wait: max_wait,
    request_timeout: request_timeout,
    cancel_on_timeout: cancel_on_timeout,
    auto_fetch_chunks: auto_fetch_chunks
  )
end

#execute_statement_async(statement:, parameters: nil, format: "JSON_ARRAY", disposition: "INLINE", catalog: nil, schema: nil, wait_timeout: "10s", on_wait_timeout: "CONTINUE", request_timeout: nil, byte_limit: nil, row_limit: nil, query_tags: nil) ⇒ Object 



74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
 
# File 'lib/databricks_sql/client.rb', line 74

def execute_statement_async(statement:, parameters: nil, format: "JSON_ARRAY", disposition: "INLINE", catalog: nil,
                            schema: nil, wait_timeout: "10s", on_wait_timeout: "CONTINUE", request_timeout: nil,
                            byte_limit: nil, row_limit: nil, query_tags: nil)
  payload = {
    statement: statement,
    warehouse_id: warehouse_id,
    disposition: disposition,
    format: format,
    wait_timeout: wait_timeout,
    on_wait_timeout: on_wait_timeout
  }
  payload[:parameters] = parameters if parameters
  payload[:catalog] = catalog if catalog
  payload[:schema] = schema if schema
  payload[:byte_limit] = byte_limit unless byte_limit.nil?
  payload[:row_limit] = row_limit unless row_limit.nil?
  payload[:query_tags] = query_tags if query_tags

  request_json(:post, @statement_path, payload: payload, request_timeout: request_timeout)
end

#fetch_next_chunk(statement_id:, next_chunk_internal_link:, columns:, disposition: "INLINE", column_schema: nil, request_timeout: nil) ⇒ Object 



117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
 
# File 'lib/databricks_sql/client.rb', line 117

def fetch_next_chunk(statement_id:, next_chunk_internal_link:, columns:, disposition: "INLINE", column_schema: nil,
                     request_timeout: nil)
  payload = get_statement_chunk_by_link(
    next_chunk_internal_link: next_chunk_internal_link,
    request_timeout: request_timeout
  )
  rows = chunk_rows(payload, columns, disposition: disposition, column_schema: column_schema,
                                      request_timeout: request_timeout)

  Result.new(
    statement_id: statement_id,
    status: "SUCCEEDED",
    disposition: normalize_disposition(disposition),
    format: "JSON_ARRAY",
    columns: columns,
    rows: rows,
    manifest: nil,
    raw_response: payload,
    next_chunk_internal_link: extract_next_chunk_internal_link(payload)
  )
end

#get_statement(statement_id:, request_timeout: nil) ⇒ Object 



95
96
97
98
 
# File 'lib/databricks_sql/client.rb', line 95

def get_statement(statement_id:, request_timeout: nil)
  path = "#{@statement_path}/#{statement_id}"
  request_json(:get, path, request_timeout: request_timeout)
end

#get_statement_chunk(statement_id:, chunk_index:, row_offset: nil, request_timeout: nil) ⇒ Object 



107
108
109
110
111
 
# File 'lib/databricks_sql/client.rb', line 107

def get_statement_chunk(statement_id:, chunk_index:, row_offset: nil, request_timeout: nil)
  path = "#{@statement_path}/#{statement_id}/result/chunks/#{chunk_index}"
  path = "#{path}?row_offset=#{row_offset}" unless row_offset.nil?
  request_json(:get, path, request_timeout: request_timeout)
end

#get_statement_chunk_by_link(next_chunk_internal_link:, request_timeout: nil) ⇒ Object 



113
114
115
 
# File 'lib/databricks_sql/client.rb', line 113

def get_statement_chunk_by_link(next_chunk_internal_link:, request_timeout: nil)
  request_json(:get, next_chunk_internal_link, request_timeout: request_timeout)
end

#inspect ⇒ Object 



169
170
171
172
173
174
175
176
177
178
179
 
# File 'lib/databricks_sql/client.rb', line 169

def inspect
  "#<#{self.class} " \
    "@external_link_allowed_hosts=#{@external_link_allowed_hosts.inspect}, " \
    "@external_link_require_https=#{@external_link_require_https.inspect}, " \
    "@host=#{@host.inspect}, " \
    "@open_timeout=#{@open_timeout.inspect}, " \
    "@statement_path=#{@statement_path.inspect}, " \
    "@timeout=#{@timeout.inspect}, " \
    "@token=\"[REDACTED]\", " \
    "@warehouse_id=#{@warehouse_id.inspect}>"
end

#wait_for_statement(statement_id:, format: "JSON_ARRAY", disposition: "INLINE", column_schema: nil, poll_interval: 1.0, max_wait: 300, request_timeout: nil, cancel_on_timeout: false, auto_fetch_chunks: true) ⇒ Object 



139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
 
# File 'lib/databricks_sql/client.rb', line 139

def wait_for_statement(statement_id:, format: "JSON_ARRAY", disposition: "INLINE", column_schema: nil,
             poll_interval: 1.0, max_wait: 300, request_timeout: nil, cancel_on_timeout: false,
             auto_fetch_chunks: true)
  started_at = Process.clock_gettime(Process::CLOCK_MONOTONIC)

  loop do
    response = get_statement(statement_id: statement_id, request_timeout: request_timeout)

    if terminal_status?(statement_status(response))
      response = hydrate_terminal_response(
        response,
        statement_id: statement_id,
        request_timeout: request_timeout
      )

      return build_result(response, format: format, disposition: disposition, column_schema: column_schema,
                                    auto_fetch_chunks: auto_fetch_chunks,
                                    request_timeout: request_timeout)
    end

    elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - started_at
    if elapsed >= max_wait
      cancel_statement(statement_id: statement_id, request_timeout: request_timeout) if cancel_on_timeout
      raise TimeoutError.new("Polling exceeded max_wait of #{max_wait}s for statement #{statement_id}", max_wait)
    end

    sleep(poll_interval)
  end
end