queuery_client 1.0.0 → 1.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c873a091add34619c21ef63bd324a41205dfdc4f509a9b4574a2a31c2c121e4b
-  data.tar.gz: 5b184dba90213ba0f022302ae2dc99e59ec491189e4c7f65c5c05b24d3497ec3
+  metadata.gz: 2a5b41cb805e7fc0864fd695adc519c6b7212c294383935af40301d88b5b8d63
+  data.tar.gz: d4e970371d23e36bf40d4cafdf9e7da0d886c25bb0882ce9e12fc2a971a0e2c0
 SHA512:
-  metadata.gz: 59fb860a84e25e604f160cd8209aecebd742b1a4e36fd84e3591bc4b7a0bc4f6c8ecb3482eb001c22eae379bec109ebf78ebe78a9f435a1c99b05f5e86299a56
-  data.tar.gz: a56489e0fa77068fc548f9b4c98fc9215f962c80ec2fe197f1958e05538992884ddc7aaa76bd2f86c43980db9c31437f24485f8371ea4739c867ce7837bd7c45
+  metadata.gz: f61f86445e64fca2863ba474a6cfb27e87af761516b3fbb5394adbe6f61e23e4f9d81778cded02dc512bf495bb8b711099962b5fd63154880bcc3e746bf25f61
+  data.tar.gz: ee0f09b1e92729056857117a7fdf4fa422d13bae5168f521e6d2cd8ca8adc2152334c59fd826aff916369ee8db811c792a49087ddd08dfdb283c0a31445bab78

data/lib/queuery_client/client.rb

@@ -4,17 +4,19 @@ module QueueryClient
       @options = options
     end
 
-    def execute_query(select_stmt, values)
-      garage_client.post("/v1/queries", q: select_stmt, values: values)
+    def execute_query(select_stmt, values, query_options)
+      garage_client.post("/v1/queries", q: select_stmt, values: values, enable_metadata: query_options[:enable_cast])
     end
+    alias start_query execute_query
 
-    def get_query(id)
-      garage_client.get("/v1/queries/#{id}", fields: '__default__,s3_prefix')
+    def get_query(id, query_options)
+      query_option_fields = build_query_option_fields(query_options)
+      garage_client.get("/v1/queries/#{id}", fields: '__default__,s3_prefix' + query_option_fields)
     end
 
-    def wait_for(id)
+    def wait_for(id, query_options)
       loop do
-        query = get_query(id)
+        query = get_query(id, query_options)
         case query.status
         when 'success', 'failed'
           return query
@@ -23,17 +25,19 @@ module QueueryClient
       end
     end
 
-    def query_and_wait(select_stmt, values)
-      query = execute_query(select_stmt, values)
-      wait_for(query.id)
+    def query_and_wait(select_stmt, values, query_options)
+      query = execute_query(select_stmt, values, query_options)
+      wait_for(query.id, query_options)
     end
 
-    def query(select_stmt, values)
-      query = query_and_wait(select_stmt, values)
+    def query(select_stmt, values, **query_options)
+      query = query_and_wait(select_stmt, values, query_options)
+      manifest_file_url = query.manifest_file_url if query_options[:enable_cast]
       case query.status
       when 'success'
         UrlDataFileBundle.new(
           query.data_file_urls,
+          manifest_file_url,
           s3_prefix: query.s3_prefix,
         )
       when 'failed'
@@ -41,13 +45,21 @@ module QueueryClient
       end
     end
 
+    # poll_result returns the results only if the query has already successed.
+    def poll_result(id)
+      query = get_query(id)
+      get_query_result(query)
+    end
+
     def garage_client
       @garage_client ||= BasicAuthGarageClient.new(
         endpoint: options.endpoint,
         path_prefix: '/',
         login: options.token,
         password: options.token_secret
-      )
+      ).tap do |client|
+        client.headers['Host'] = options.host_header if options.host_header
+      end
     end
 
     def options
@@ -57,5 +69,37 @@ module QueueryClient
     def default_options
       QueueryClient.configuration
     end
+
+    private
+
+    def get_query_result(query)
+      case query.status
+      when 'pending', 'running'
+        nil
+      when 'success'
+        UrlDataFileBundle.new(
+          query.data_file_urls,
+          nil,
+          s3_prefix: query.s3_prefix,
+        )
+      when 'failure'
+        raise QueryError.new(query.error)
+      end
+    end
+
+    def build_query_option_fields(query_options)
+      enable_query_options = query_options.select{ |name, v| name if v }.keys
+      return '' if enable_query_options.empty?
+      query_option_fields = enable_query_options.map{ |option_name| convert_field_name(option_name) }
+      ',' + query_option_fields.join(',')
+    end
+
+    def convert_field_name(option_name)
+      case option_name
+        when :enable_cast
+          'manifest_file_url'
+        # add another option here if you need
+      end
+    end
   end
 end

data/lib/queuery_client/configuration.rb

@@ -1,5 +1,8 @@
 module QueueryClient
   class Configuration
+    REQUIRED_KEYS = [:endpoint, :token, :token_secret]
+    OPTIONAL_KEYS = [:host_header]
+
     def initialize(options = {})
       @options = options
     end
@@ -12,11 +15,7 @@ module QueueryClient
       @options = nil
     end
 
-    [
-      :endpoint,
-      :token,
-      :token_secret,
-    ].each do |key|
+    REQUIRED_KEYS.each do |key|
       define_method(key) do
         options.fetch(key)
       end
@@ -26,6 +25,16 @@ module QueueryClient
       end
     end
 
+    OPTIONAL_KEYS.each do |key|
+      define_method(key) do
+        options[key]
+      end
+
+      define_method("#{key}=") do |value|
+        options[key] = value
+      end
+    end
+
     def merge(other)
       Configuration.new(to_h.merge(other.to_h))
     end

data/lib/queuery_client/data_file.rb

@@ -2,7 +2,7 @@ require 'redshift_csv_file'
 require 'zlib'
 
 module QueueryClient
-  class DataFile
+  class DataFile # abstract class
     def data_object?
       /\.csv(?:\.|\z)/ =~ File.basename(key)
     end
@@ -12,15 +12,21 @@ module QueueryClient
     end
 
     def each_row(&block)
+      return enum_for(:each_row) if !block_given?
+
       f = open
       begin
         if gzipped_object?
           f = Zlib::GzipReader.new(f)
         end
-        RedshiftCsvFile.new(f).each(&block)
+        RedshiftCsvFile.new(f).each do |row|
+          yield row
+        end
       ensure
         f.close
       end
+
+      self
     end
   end
 end

data/lib/queuery_client/data_file_bundle.rb

@@ -1,15 +1,31 @@
+require 'queuery_client/redshift_data_type'
+require 'date'
+require 'time'
+
 module QueueryClient
   class DataFileBundle
     # abstract data_files :: [DataFile]
+    # abstract manifest_file :: ManifestFile
+    # abstract def has_manifest?
 
     def each_row(&block)
+      return enum_for(:each_row) if !block_given?
+
       data_files.each do |file|
         if file.data_object?
-          file.each_row(&block)
+          file.each_row do |row|
+            if has_manifest?
+              yield RedshiftDataType.type_cast(row, manifest_file)
+            else
+              yield row
+            end
+          end
         end
       end
-    end
 
+      self
+    end
     alias each each_row
+
   end
 end

data/lib/queuery_client/manifest_file.rb

@@ -0,0 +1,20 @@
+require 'json'
+
+module QueueryClient
+  class ManifestFile # abstract class
+    def manifest_object?
+      /\.manifest(?:\.|\z)/ =~ File.basename(key)
+    end
+
+    def column_types
+      @column_types ||=
+        begin
+          f = open
+          j = JSON.load(f)
+          j['schema']['elements'].map{|x| x['type']['base']}
+        ensure
+          f.close
+        end
+    end
+  end
+end

data/lib/queuery_client/redshift_data_type.rb

@@ -0,0 +1,27 @@
+
+module QueueryClient
+  module RedshiftDataType
+    def self.type_cast(row, manifest_file)
+      row.zip(manifest_file.column_types).map do |value, type|
+        next nil if (value == '' and type != 'character varing') # null becomes '' on unload
+
+        case type
+        when 'smallint', 'integer', 'bigint'
+          value.to_i
+        when 'numeric', 'double precision'
+          value.to_f
+        when 'character', 'character varying'
+          value
+        when 'timestamp without time zone', 'timestamp with time zone'
+          Time.parse(value)
+        when 'date'
+          Date.parse(value)
+        when 'boolean'
+          value == 'true' ? true : false
+        else
+          raise "not support data type: #{type}"
+        end
+      end
+    end
+  end
+end

data/lib/queuery_client/s3_data_file_bundle.rb

@@ -1,5 +1,6 @@
 require 'queuery_client/data_file_bundle'
 require 'queuery_client/s3_data_file'
+require 'queuery_client/s3_manifest_file'
 require 'aws-sdk-s3'
 require 'logger'
 
@@ -24,5 +25,15 @@ module QueueryClient
       b = Aws::S3::Resource.new(client: @s3_client).bucket(@bucket)
       b.objects(prefix: @prefix).map {|obj| S3DataFile.new(obj) }
     end
+
+    def manifest_file
+      b = Aws::S3::Resource.new(client: @s3_client).bucket(@bucket)
+      obj = b.object("#{@prefix}manifest")
+      S3ManifestFile.new(obj)
+    end
+
+    def has_manifest?
+      !manifest_file.nil?
+    end
   end
 end

data/lib/queuery_client/s3_manifest_file.rb

@@ -0,0 +1,18 @@
+require 'queuery_client/manifest_file'
+require 'forwardable'
+
+module QueueryClient
+  class S3ManifestFile < ManifestFile
+    extend Forwardable
+
+    def initialize(object)
+      @object = object
+    end
+
+    def_delegators '@object', :url, :key, :presigned_url
+
+    def open
+      @object.get.body
+    end
+  end
+end

data/lib/queuery_client/url_data_file_bundle.rb

@@ -1,19 +1,23 @@
 require 'queuery_client/data_file_bundle'
 require 'queuery_client/url_data_file'
+require 'queuery_client/url_manifest_file'
 require 'uri'
 require 'logger'
 
 module QueueryClient
   class UrlDataFileBundle < DataFileBundle
-    def initialize(urls, s3_prefix:, logger: Logger.new($stderr))
-      raise ArgumentError, 'no URL given' if urls.empty?
-      @data_files = urls.map {|url| UrlDataFile.new(URI.parse(url)) }
+    def initialize(data_urls, manifest_url, s3_prefix:, logger: Logger.new($stderr))
+      raise ArgumentError, 'no URL given' if data_urls.empty?
+      @data_files = data_urls.map {|url| UrlDataFile.new(URI.parse(url)) }
+      @manifest_file = UrlManifestFile.new(URI.parse(manifest_url)) if manifest_url
       @s3_prefix = s3_prefix
       @logger = logger
     end
 
     attr_reader :data_files
+    attr_reader :manifest_file
     attr_reader :s3_prefix
+    attr_reader :logger
 
     def url
       uri = data_files.first.url.dup
@@ -28,5 +32,9 @@ module QueueryClient
       prefix = s3_uri.path[1..-1] # trim heading slash
       S3DataFileBundle.new(bucket, prefix)
     end
+
+    def has_manifest?
+      !@manifest_file.nil?
+    end
   end
 end

data/lib/queuery_client/url_manifest_file.rb

@@ -0,0 +1,28 @@
+require 'queuery_client/manifest_file'
+require 'net/http'
+require 'stringio'
+require 'json'
+
+module QueueryClient
+  class UrlManifestFile < ManifestFile
+    def initialize(url)
+      @url = url
+    end
+
+    attr_reader :url
+
+    def key
+      @url.path
+    end
+
+    def open
+      http = Net::HTTP.new(@url.host, @url.port)
+      http.use_ssl = (@url.scheme.downcase == 'https')
+      content = http.start {
+        res = http.get(@url.request_uri)
+        res.body
+      }
+      StringIO.new(content)
+    end
+  end
+end

data/lib/queuery_client/version.rb

@@ -1,3 +1,3 @@
 module QueueryClient
-  VERSION = "1.0.0"
+  VERSION = "1.1.0"
 end

data/lib/queuery_client.rb

@@ -16,8 +16,8 @@ module QueueryClient
       configuration.instance_eval(&block)
     end
 
-    def query(select_stmt, values = [])
-      Client.new.query(select_stmt, values)
+    def query(select_stmt, values = [], enable_cast: false)
+      Client.new.query(select_stmt, values, enable_cast: enable_cast)
     end
   end
 end

data/queuery_client.gemspec

@@ -32,6 +32,6 @@ Gem::Specification.new do |spec|
   spec.add_dependency 'redshift_csv_file'
   spec.add_dependency 'aws-sdk-s3'
   spec.add_development_dependency "bundler", "~> 1.13"
-  spec.add_development_dependency "rake", "~> 10.0"
+  spec.add_development_dependency "rake", "~> 12.0"
   spec.add_development_dependency "pry"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: queuery_client
 version: !ruby/object:Gem::Version
-  version: 1.0.0
+  version: 1.1.0
 platform: ruby
 authors:
 - Hidekazu Kobayashi
-autorequire: 
+autorequire:
 bindir: exe
 cert_chain: []
-date: 2019-07-23 00:00:00.000000000 Z
+date: 2021-10-07 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: garage_client
@@ -72,14 +72,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '10.0'
+        version: '12.0'
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '10.0'
+        version: '12.0'
 - !ruby/object:Gem::Dependency
   name: pry
   requirement: !ruby/object:Gem::Requirement
@@ -94,7 +94,7 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
-description: 
+description:
 email:
 - hidekazu-kobayashi@cookpad.com
 executables: []
@@ -113,11 +113,15 @@ files:
 - lib/queuery_client/configuration.rb
 - lib/queuery_client/data_file.rb
 - lib/queuery_client/data_file_bundle.rb
+- lib/queuery_client/manifest_file.rb
 - lib/queuery_client/query_error.rb
+- lib/queuery_client/redshift_data_type.rb
 - lib/queuery_client/s3_data_file.rb
 - lib/queuery_client/s3_data_file_bundle.rb
+- lib/queuery_client/s3_manifest_file.rb
 - lib/queuery_client/url_data_file.rb
 - lib/queuery_client/url_data_file_bundle.rb
+- lib/queuery_client/url_manifest_file.rb
 - lib/queuery_client/version.rb
 - queuery_client.gemspec
 homepage: https://github.com/bricolages/queuery_client
@@ -125,7 +129,7 @@ licenses:
 - MIT
 metadata:
   allowed_push_host: https://rubygems.org
-post_install_message: 
+post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -140,8 +144,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.0.3
-signing_key: 
+rubygems_version: 3.1.6
+signing_key:
 specification_version: 4
 summary: Client library for Queuery Redshift HTTP API
 test_files: []