queuery_client 1.0.1 → 1.1.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: abbdfbc43e113b76608bc2c0a4555bd5e194e9b10c4b7d89f6fe3e78509245f4
-  data.tar.gz: 15cd76d0c3f02ed04e59254585428255686020e71603fa6633ffb022bd5f5357
+  metadata.gz: 680cc3940c5e82d5ae0772d89cd25ffc00310236f9d92b6a686f900dd46f0a88
+  data.tar.gz: 468821770777a14d147958b69a3f21b8a88a827e677b787ec2eba6b06b6d1007
 SHA512:
-  metadata.gz: 95d02421bdbe1c30b995f9a7a40d86901339380c71dac23979a5653ea0f4038c98e9f8011c6d43a32dbbee482d5ef061639a9b4c3da9a718f36b5eaa46518f9c
-  data.tar.gz: b558a205e45be08e7401a29aa1c94db29b181c01e950019b58773b9b511dd47fb62601e520da8474dca71dc53aa688cedc73245e4b5f347c0008607a7d12ae08
+  metadata.gz: f42b47de00d2f718ce017a5e7a343b1e84a9f681ab3195796e81f7da3c85f5bf21b23cb5f5d4a151c2b0f3a8feba1e8fb01d2645fa17a9393803f3f6f6e0826e
+  data.tar.gz: 9139744e2b4f16ea8654e6e32e961270bb67b732eac50dd3633f227aa4d97b74249d3ff8ad91e9d47f2fbade3a47dcdda716b1f7a252ec58d90e7352d755cb0d

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,38 @@
+
+module QueueryClient
+  module RedshiftDataType
+    FALSE_VALUES = [
+      false, 0,
+      "0", :"0",
+      "f", :f,
+      "F", :F,
+      "false", :false,
+      "FALSE", :FALSE,
+      "off", :off,
+      "OFF", :OFF,
+    ].to_set.freeze
+
+    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'
+          FALSE_VALUES.include?(value) ? false : true
+        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'
 
@@ -22,7 +23,23 @@ module QueueryClient
 
     def data_files
       b = Aws::S3::Resource.new(client: @s3_client).bucket(@bucket)
-      b.objects(prefix: @prefix).map {|obj| S3DataFile.new(obj) }
+      b.objects(prefix: @prefix)
+        .select {|obj| obj.key.include?('_part_') }
+        .map {|obj| S3DataFile.new(obj) }
+    end
+
+    def manifest_file
+      b = Aws::S3::Resource.new(client: @s3_client).bucket(@bucket)
+      obj = b.object("#{@prefix}manifest")
+      if obj.exists?
+        S3ManifestFile.new(obj)
+      else
+        nil
+      end
+    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,18 +1,21 @@
 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
 
@@ -29,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.1"
+  VERSION = "1.1.2"
 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.1
+  version: 1.1.2
 platform: ruby
 authors:
 - Hidekazu Kobayashi
-autorequire: 
+autorequire:
 bindir: exe
 cert_chain: []
-date: 2019-08-30 00:00:00.000000000 Z
+date: 2021-11-10 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.1
-signing_key: 
+rubygems_version: 3.1.6
+signing_key:
 specification_version: 4
 summary: Client library for Queuery Redshift HTTP API
 test_files: []