presto-client-legacy 0.4.17

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,176 @@
1
+ #
2
+ # Presto client for Ruby
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ #
16
+ module PrestoLegacy::Client
17
+
18
+ require 'faraday'
19
+ require 'presto/client/models'
20
+ require 'presto/client/errors'
21
+ require 'presto/client/statement_client'
22
+
23
+ class Query
24
+ def self.start(query, options)
25
+ new StatementClient.new(faraday_client(options), query, options)
26
+ end
27
+
28
+ def self.resume(next_uri, options)
29
+ new StatementClient.new(faraday_client(options), nil, options, next_uri)
30
+ end
31
+
32
+ def self.faraday_client(options)
33
+ server = options[:server]
34
+ unless server
35
+ raise ArgumentError, ":server option is required"
36
+ end
37
+
38
+ ssl = faraday_ssl_options(options)
39
+
40
+ url = "#{ssl ? "https" : "http"}://#{server}"
41
+ proxy = options[:http_proxy] || options[:proxy] # :proxy is obsoleted
42
+
43
+ faraday = Faraday.new(url: url, proxy: "#{proxy}", ssl: ssl) do |faraday|
44
+ #faraday.request :url_encoded
45
+ faraday.response :logger if options[:http_debug]
46
+ faraday.adapter Faraday.default_adapter
47
+ end
48
+
49
+ return faraday
50
+ end
51
+
52
+ def self.faraday_ssl_options(options)
53
+ ssl = options[:ssl]
54
+
55
+ case ssl
56
+ when true
57
+ ssl = {verify: true}
58
+
59
+ when Hash
60
+ verify = ssl.fetch(:verify, true)
61
+ case verify
62
+ when true
63
+ # detailed SSL options. pass through to faraday
64
+ when nil, false
65
+ ssl = {verify: false}
66
+ else
67
+ raise ArgumentError, "Can't convert #{verify.class} of :verify option of :ssl option to true or false"
68
+ end
69
+
70
+ when nil, false
71
+ ssl = false
72
+
73
+ else
74
+ raise ArgumentError, "Can't convert #{ssl.class} of :ssl option to true, false, or Hash"
75
+ end
76
+
77
+ return ssl
78
+ end
79
+
80
+ private_class_method :faraday_client, :faraday_ssl_options
81
+
82
+ def initialize(api)
83
+ @api = api
84
+ end
85
+
86
+ def current_results
87
+ @api.current_results
88
+ end
89
+
90
+ def advance
91
+ @api.advance
92
+ end
93
+
94
+ def wait_for_columns
95
+ while @api.current_results.columns == nil && @api.advance
96
+ end
97
+ end
98
+
99
+ def wait_for_data
100
+ while @api.current_results.data == nil && @api.advance
101
+ end
102
+ end
103
+
104
+ private :wait_for_columns
105
+ private :wait_for_data
106
+
107
+ def columns
108
+ wait_for_columns
109
+
110
+ raise_if_failed
111
+
112
+ return @api.current_results.columns
113
+ end
114
+
115
+ def rows
116
+ rows = []
117
+ each_row_chunk {|chunk|
118
+ rows.concat(chunk)
119
+ }
120
+ return rows
121
+ end
122
+
123
+ def each_row(&block)
124
+ each_row_chunk {|chunk|
125
+ chunk.each(&block)
126
+ }
127
+ end
128
+
129
+ def each_row_chunk(&block)
130
+ wait_for_data
131
+
132
+ raise_if_failed
133
+
134
+ if self.columns == nil
135
+ raise PrestoError, "Query #{@api.current_results.id} has no columns"
136
+ end
137
+
138
+ begin
139
+ if data = @api.current_results.data
140
+ block.call(data)
141
+ end
142
+ end while @api.advance
143
+ end
144
+
145
+ def query_info
146
+ @api.query_info
147
+ end
148
+
149
+ def next_uri
150
+ @api.current_results.next_uri
151
+ end
152
+
153
+ def cancel
154
+ @api.cancel_leaf_stage
155
+ end
156
+
157
+ def close
158
+ @api.cancel_leaf_stage
159
+ nil
160
+ end
161
+
162
+ def raise_if_failed
163
+ if @api.closed?
164
+ raise PrestoClientError, "Query aborted by user"
165
+ elsif @api.exception?
166
+ # query is gone
167
+ raise @api.exception
168
+ elsif @api.query_failed?
169
+ results = @api.current_results
170
+ error = results.error
171
+ raise PrestoQueryError.new("Query #{results.id} failed: #{error.message}", results.id, error.error_code, error.failure_info)
172
+ end
173
+ end
174
+ end
175
+
176
+ end
@@ -0,0 +1,259 @@
1
+ #
2
+ # Presto client for Ruby
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ #
16
+ module PrestoLegacy::Client
17
+
18
+ require 'multi_json'
19
+ require 'presto/client/models'
20
+ require 'presto/client/errors'
21
+
22
+ module PrestoHeaders
23
+ PRESTO_USER = "X-Presto-User"
24
+ PRESTO_SOURCE = "X-Presto-Source"
25
+ PRESTO_CATALOG = "X-Presto-Catalog"
26
+ PRESTO_SCHEMA = "X-Presto-Schema"
27
+ PRESTO_TIME_ZONE = "X-Presto-Time-Zone"
28
+ PRESTO_LANGUAGE = "X-Presto-Language"
29
+ PRESTO_SESSION = "X-Presto-Session"
30
+
31
+ PRESTO_CURRENT_STATE = "X-Presto-Current-State"
32
+ PRESTO_MAX_WAIT = "X-Presto-Max-Wait"
33
+ PRESTO_MAX_SIZE = "X-Presto-Max-Size"
34
+ PRESTO_PAGE_SEQUENCE_ID = "X-Presto-Page-Sequence-Id"
35
+ end
36
+
37
+ class StatementClient
38
+ HEADERS = {
39
+ "User-Agent" => "presto-ruby/#{VERSION}",
40
+ }
41
+
42
+ def initialize(faraday, query, options, next_uri=nil)
43
+ @faraday = faraday
44
+ @faraday.headers.merge!(HEADERS)
45
+
46
+ @options = options
47
+ @query = query
48
+ @closed = false
49
+ @exception = nil
50
+
51
+ @faraday.headers.merge!(optional_headers)
52
+
53
+ if next_uri
54
+ body = faraday_get_with_retry(next_uri)
55
+ @results = Models::QueryResults.decode(MultiJson.load(body))
56
+ else
57
+ post_query_request!
58
+ end
59
+ end
60
+
61
+ def optional_headers
62
+ headers = {}
63
+ if v = @options[:user]
64
+ headers[PrestoHeaders::PRESTO_USER] = v
65
+ end
66
+ if v = @options[:source]
67
+ headers[PrestoHeaders::PRESTO_SOURCE] = v
68
+ end
69
+ if v = @options[:catalog]
70
+ headers[PrestoHeaders::PRESTO_CATALOG] = v
71
+ end
72
+ if v = @options[:schema]
73
+ headers[PrestoHeaders::PRESTO_SCHEMA] = v
74
+ end
75
+ if v = @options[:time_zone]
76
+ headers[PrestoHeaders::PRESTO_TIME_ZONE] = v
77
+ end
78
+ if v = @options[:language]
79
+ headers[PrestoHeaders::PRESTO_LANGUAGE] = v
80
+ end
81
+ if v = @options[:properties]
82
+ headers[PrestoHeaders::PRESTO_SESSION] = encode_properties(v)
83
+ end
84
+ headers
85
+ end
86
+
87
+ private :optional_headers
88
+
89
+ def init_request(req)
90
+ req.options.timeout = @options[:http_timeout] || 300
91
+ req.options.open_timeout = @options[:http_open_timeout] || 60
92
+ end
93
+
94
+ private :init_request
95
+
96
+ def post_query_request!
97
+ uri = "/v1/statement"
98
+ response = @faraday.post do |req|
99
+ req.url uri
100
+
101
+ req.body = @query
102
+ init_request(req)
103
+ end
104
+
105
+ # TODO error handling
106
+ if response.status != 200
107
+ raise PrestoHttpError.new(response.status, "Failed to start query: #{response.body}")
108
+ end
109
+
110
+ body = response.body
111
+ @results = load_json(uri, body, Models::QueryResults)
112
+ end
113
+
114
+ private :post_query_request!
115
+
116
+ attr_reader :query
117
+
118
+ def debug?
119
+ !!@options[:debug]
120
+ end
121
+
122
+ def closed?
123
+ @closed
124
+ end
125
+
126
+ attr_reader :exception
127
+
128
+ def exception?
129
+ @exception
130
+ end
131
+
132
+ def query_failed?
133
+ @results.error != nil
134
+ end
135
+
136
+ def query_succeeded?
137
+ @results.error == nil && !@exception && !@closed
138
+ end
139
+
140
+ def current_results
141
+ @results
142
+ end
143
+
144
+ def has_next?
145
+ !!@results.next_uri
146
+ end
147
+
148
+ def advance
149
+ if closed? || !has_next?
150
+ return false
151
+ end
152
+ uri = @results.next_uri
153
+
154
+ body = faraday_get_with_retry(uri)
155
+ @results = load_json(uri, body, Models::QueryResults)
156
+
157
+ return true
158
+ end
159
+
160
+ def query_info
161
+ uri = "/v1/query/#{@results.id}"
162
+ body = faraday_get_with_retry(uri)
163
+ load_json(uri, body, Models::QueryInfo)
164
+ end
165
+
166
+ def load_json(uri, body, body_class)
167
+ hash = MultiJson.load(body)
168
+ begin
169
+ body_class.decode(hash)
170
+ rescue => e
171
+ if body.size > 1024 + 3
172
+ body = "#{body[0, 1024]}..."
173
+ end
174
+ @exception = PrestoHttpError.new(500, "Presto API returned unexpected structure at #{uri}. Expected #{body_class} but got #{body}: #{e}")
175
+ raise @exception
176
+ end
177
+ end
178
+
179
+ private :load_json
180
+
181
+ def faraday_get_with_retry(uri, &block)
182
+ start = Time.now
183
+ attempts = 0
184
+
185
+ begin
186
+ begin
187
+ response = @faraday.get(uri)
188
+ rescue Faraday::Error::TimeoutError, Faraday::Error::ConnectionFailed
189
+ # temporally error to retry
190
+ response = nil
191
+ rescue => e
192
+ @exception = e
193
+ raise @exception
194
+ end
195
+
196
+ if response
197
+ if response.status == 200 && !response.body.to_s.empty?
198
+ return response.body
199
+ end
200
+
201
+ if response.status != 503 # retry only if 503 Service Unavailable
202
+ # deterministic error
203
+ @exception = PrestoHttpError.new(response.status, "Presto API error at #{uri} returned #{response.status}: #{response.body}")
204
+ raise @exception
205
+ end
206
+ end
207
+
208
+ attempts += 1
209
+ sleep attempts * 0.1
210
+ end while (Time.now - start) < 2*60*60 && !@closed
211
+
212
+ @exception = PrestoHttpError.new(408, "Presto API error due to timeout")
213
+ raise @exception
214
+ end
215
+
216
+ def cancel_leaf_stage
217
+ if uri = @results.next_uri
218
+ response = @faraday.delete do |req|
219
+ req.url uri
220
+ end
221
+ return response.status / 100 == 2
222
+ end
223
+ return false
224
+ end
225
+
226
+ HTTP11_SEPARATOR = ["(", ")", "<", ">", "@", ",", ";", ":", "\\", "<", ">", "/", "[", "]", "?", "=", "{", "}", " ", "\v"]
227
+ HTTP11_TOKEN_CHARSET = (32..126).map {|x| x.chr } - HTTP11_SEPARATOR
228
+ HTTP11_TOKEN_REGEXP = /^[#{Regexp.escape(HTTP11_TOKEN_CHARSET.join)}]+\z/
229
+ HTTP11_CTL_CHARSET = (0..31).map {|x| x.chr } + [127.chr]
230
+ HTTP11_CTL_CHARSET_REGEXP = /[#{Regexp.escape(HTTP11_CTL_CHARSET.join)}]/
231
+
232
+ def encode_properties(properties)
233
+ # this is a hack to set same header multiple times.
234
+ properties.map do |k, v|
235
+ token = k.to_s
236
+ field_value = v.to_s # TODO LWS encoding is not implemented
237
+ unless k =~ HTTP11_TOKEN_REGEXP
238
+ raise Faraday::ClientError, "Key of properties can't include HTTP/1.1 control characters or separators (#{HTTP11_SEPARATOR.map {|c| c =~ /\s/ ? c.dump : c }.join(' ')})"
239
+ end
240
+ if field_value =~ HTTP11_CTL_CHARSET_REGEXP
241
+ raise Faraday::ClientError, "Value of properties can't include HTTP/1.1 control characters"
242
+ end
243
+ "#{token}=#{field_value}"
244
+ end.join("\r\n#{PrestoHeaders::PRESTO_SESSION}: ")
245
+ end
246
+
247
+ def close
248
+ return if @closed
249
+
250
+ # cancel running statement
251
+ # TODO make async reqeust and ignore response?
252
+ cancel_leaf_stage
253
+
254
+ @closed = true
255
+ nil
256
+ end
257
+ end
258
+
259
+ end
@@ -0,0 +1,20 @@
1
+ #
2
+ # Presto client for Ruby
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ #
16
+ module PrestoLegacy
17
+ module Client
18
+ VERSION = "0.4.17"
19
+ end
20
+ end
@@ -0,0 +1,23 @@
1
+ #
2
+ # Presto client for Ruby
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ #
16
+ module PrestoLegacy
17
+ module Client
18
+
19
+ require 'presto/client/version'
20
+ require 'presto/client/client'
21
+
22
+ end
23
+ end
@@ -0,0 +1 @@
1
+ require 'presto/client'
@@ -0,0 +1,91 @@
1
+
2
+ if ARGV.length != 3
3
+ puts "usage: <presto-source-dir> <template.erb> <output.rb>"
4
+ end
5
+
6
+ source_dir, template_path, output_path = *ARGV
7
+
8
+ require_relative 'presto_models'
9
+
10
+ require 'erb'
11
+ erb = ERB.new(File.read(template_path))
12
+
13
+ source_path = source_dir
14
+
15
+ predefined_simple_classes = %w[StageId TaskId ConnectorSession]
16
+ predefined_models = %w[DistributionSnapshot PlanNode EquiJoinClause WriterTarget PageBufferInfo DeleteHandle]
17
+
18
+ assume_primitive = %w[Object Type Long Symbol QueryId PlanNodeId PlanFragmentId MemoryPoolId TransactionId URI Duration DataSize DateTime ColumnHandle ConnectorTableHandle ConnectorOutputTableHandle ConnectorIndexHandle ConnectorColumnHandle ConnectorInsertTableHandle ConnectorTableLayoutHandle Expression FunctionCall TimeZoneKey Locale TypeSignature Frame TupleDomain<ColumnHandle> SerializableNativeValue ConnectorTransactionHandle]
19
+ enum_types = %w[QueryState StageState TaskState QueueState PlanDistribution OutputPartitioning Step SortOrder BufferState NullPartitioning BlockedReason ParameterKind FunctionKind PartitionFunctionHandle]
20
+
21
+ root_models = %w[QueryResults QueryInfo] + %w[
22
+ OutputNode
23
+ ProjectNode
24
+ TableScanNode
25
+ ValuesNode
26
+ AggregationNode
27
+ MarkDistinctNode
28
+ FilterNode
29
+ WindowNode
30
+ RowNumberNode
31
+ TopNRowNumberNode
32
+ LimitNode
33
+ DistinctLimitNode
34
+ TopNNode
35
+ SampleNode
36
+ SortNode
37
+ RemoteSourceNode
38
+ JoinNode
39
+ SemiJoinNode
40
+ IndexJoinNode
41
+ IndexSourceNode
42
+ TableWriterNode
43
+ DeleteNode
44
+ MetadataDeleteNode
45
+ TableFinishNode
46
+ UnnestNode
47
+ ExchangeNode
48
+ UnionNode
49
+ EnforceSingleRowNode
50
+ ] + %w[InsertTableHandle OutputTableHandle TableHandle]
51
+
52
+ name_mapping = Hash[*%w[
53
+ StatementStats StageStats ClientStageStats
54
+ ClientStageStats StageStats ClientStageStats
55
+ QueryResults Column ClientColumn
56
+ ].each_slice(3).map { |x, y, z| [[x,y], z] }.flatten(1)]
57
+
58
+ path_mapping = Hash[*%w[
59
+ ClientColumn presto-client/src/main/java/com/facebook/presto/client/Column.java
60
+ ClientStageStats presto-client/src/main/java/com/facebook/presto/client/StageStats.java
61
+ Column presto-main/src/main/java/com/facebook/presto/execution/Column.java
62
+ QueryStats presto-main/src/main/java/com/facebook/presto/execution/QueryStats.java
63
+ StageStats presto-main/src/main/java/com/facebook/presto/execution/StageStats.java
64
+ ].map.with_index { |v,i| i % 2 == 0 ? v : (source_path + "/" + v) }]
65
+
66
+ analyzer = PrestoModels::ModelAnalyzer.new(
67
+ source_path,
68
+ skip_models: predefined_models + predefined_simple_classes + assume_primitive + enum_types,
69
+ path_mapping: path_mapping,
70
+ name_mapping: name_mapping
71
+ )
72
+ analyzer.analyze(root_models)
73
+ models = analyzer.models
74
+ skipped_models = analyzer.skipped_models
75
+
76
+ formatter = PrestoModels::ModelFormatter.new(
77
+ base_indent_count: 2,
78
+ struct_class: "Base",
79
+ special_struct_initialize_method: "initialize_struct",
80
+ primitive_types: assume_primitive,
81
+ skip_types: skipped_models,
82
+ simple_classes: predefined_simple_classes,
83
+ enum_types: enum_types,
84
+ )
85
+ formatter.format(models)
86
+
87
+ @contents = formatter.contents
88
+
89
+ data = erb.result
90
+ File.open(output_path, 'w') {|f| f.write data }
91
+