Class: Blazer::Adapters::Snowflake2Adapter

Inherits:
BaseAdapter
  • Object
show all
Defined in:
lib/blazer/adapters/snowflake2_adapter.rb

Instance Attribute Summary

Attributes inherited from BaseAdapter

#data_source

Instance Method Summary collapse

Methods inherited from BaseAdapter

#cachable?, #cohort_analysis_statement, #cost, #explain, #initialize, #reconnect, #supports_cohort_analysis?

Constructor Details

This class inherits a constructor from Blazer::Adapters::BaseAdapter

Instance Method Details

#cancel(run_id) ⇒ Object



182
183
184
# File 'lib/blazer/adapters/snowflake2_adapter.rb', line 182

def cancel(run_id)
  # TODO
end

#parameter_bindingObject



177
178
179
# File 'lib/blazer/adapters/snowflake2_adapter.rb', line 177

def parameter_binding
  :positional
end

#preview_statementObject



167
168
169
# File 'lib/blazer/adapters/snowflake2_adapter.rb', line 167

def preview_statement
  "SELECT * FROM {table} LIMIT 10"
end

#quotingObject



172
173
174
# File 'lib/blazer/adapters/snowflake2_adapter.rb', line 172

def quoting
  :backslash_escape
end

#run_statement(statement, comment, bind_params) ⇒ Object



4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
# File 'lib/blazer/adapters/snowflake2_adapter.rb', line 4

def run_statement(statement, comment, bind_params)
  require "net/http"

  columns = []
  rows = []
  error = nil

  api_prefix = "https://#{settings.fetch("account_id")}.snowflakecomputing.com/api/v2/statements/"
  authorization = "Bearer #{settings.fetch("access_token")}"

  submit_uri = URI(api_prefix)
  # for testing
  # submit_uri.query = URI.encode_www_form({"async" => true})

  timeout = data_source.timeout || 300
  stop_at = Time.now + timeout

  post_data = {
    statement: "#{statement} /*#{comment}*/",
    timeout: timeout
  }
  # use empty? since any? doesn't work for [nil]
  unless bind_params.empty?
    post_data[:bindings] =
      bind_params.map.with_index.to_h do |v, i|
        type =
          case v
          when Integer
            "FIXED"
          when Float
            "REAL"
          when ActiveSupport::TimeWithZone
            "TIMESTAMP_NTZ"
          else
            "TEXT"
          end
        v = v.to_i * 1000000000 + v.nsec if v.is_a?(ActiveSupport::TimeWithZone)
        [i + 1, {type: type, value: v&.to_s}]
      end
  end
  post_data[:database] = settings["database"] if settings["database"]
  post_data[:schema] = settings["schema"] if settings["schema"]
  post_data[:warehouse] = settings["warehouse"] if settings["warehouse"]
  post_data[:role] = settings["role"] if settings["role"]

  req = Net::HTTP::Post.new(submit_uri)
  req["Authorization"] = authorization
  req.body = post_data.to_json

  options = {
    use_ssl: true,
    open_timeout: 3,
    read_timeout: timeout + 3
  }

  begin
    res = Net::HTTP.start(submit_uri.hostname, submit_uri.port, options) do |http|
      http.request(req)
    end

    while res.is_a?(Net::HTTPAccepted)
      if Time.now > stop_at
        error = Blazer::TIMEOUT_MESSAGE
        break
      end

      sleep(3)

      data = JSON.parse(res.body)
      statement_uri = URI("#{api_prefix}#{CGI.escape(data["statementHandle"])}")
      req = Net::HTTP::Get.new(statement_uri)
      req["Authorization"] = authorization

      res = Net::HTTP.start(statement_uri.hostname, statement_uri.port, options) do |http|
        http.request(req)
      end
    end

    if res.is_a?(Net::HTTPSuccess)
      data = JSON.parse(res.body)
       = data["resultSetMetaData"]
      columns = ["rowType"].map { |v| v["name"].downcase }
      column_types = ["rowType"].map { |v| v["type"] }
      rows = data["data"]

      if ["partitionInfo"]
        1.upto(["partitionInfo"].size - 1) do |i|
          statement_uri.query = URI.encode_www_form({"partition" => i})
          req = Net::HTTP::Get.new(statement_uri)
          req["Authorization"] = authorization

          res = Net::HTTP.start(statement_uri.hostname, statement_uri.port, options) do |http|
            http.request(req)
          end

          if res.is_a?(Net::HTTPSuccess)
            data = JSON.parse(res.body)
            rows += data["data"]
          else
            data = JSON.parse(res.body)
            error = data["message"]
            break
          end
        end
      end

      if error
        columns.clear
        rows.clear
      else
        column_types.each_with_index do |c, i|
          # TODO handle more types
          case c
          when "fixed"
            rows.each do |row|
              row[i] &&= row[i].to_i
            end
          when "real"
            rows.each do |row|
              row[i] &&= row[i].to_f
            end
          when "timestamp_ntz"
            utc = ActiveSupport::TimeZone["Etc/UTC"]
            rows.each do |row|
              row[i] &&= utc.strptime(row[i], "%s.%N")
            end
          end
        end
      end
    elsif !error
      data = JSON.parse(res.body)
      error = data["message"]
      error = Blazer::TIMEOUT_MESSAGE if error.include?("Statement reached its statement or warehouse timeout")
    end
  rescue Errno::ECONNREFUSED => e
    error = e.message
  end

  [columns, rows, error]
end

#schemaObject



161
162
163
164
165
# File 'lib/blazer/adapters/snowflake2_adapter.rb', line 161

def schema
  sql = "SELECT table_schema, table_name, column_name, data_type, ordinal_position FROM information_schema.columns WHERE table_schema != 'INFORMATION_SCHEMA'"
  result = data_source.run_statement(sql)
  result.rows.group_by { |r| [r[0], r[1]] }.sort_by { |k, _| [k[0] == default_schema ? "" : k[0], k[1]] }.map { |k, vs| {schema: k[0].downcase, table: k[1].downcase, columns: vs.sort_by { |v| v[2] }.map { |v| {name: v[2].downcase, data_type: v[3].downcase} }} }
end

#tablesObject



145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
# File 'lib/blazer/adapters/snowflake2_adapter.rb', line 145

def tables
  sql = "SELECT table_schema, table_name FROM information_schema.tables WHERE table_schema != 'INFORMATION_SCHEMA'"
  result = data_source.run_statement(sql)
  result.rows.sort_by { |r| [r[0] == default_schema ? "" : r[0], r[1]] }.map do |row|
    table =
      if row[0] == default_schema
        row[1]
      else
        "#{row[0]}.#{row[1]}"
      end

    # TODO quote if needed
    table.downcase
  end
end