queuery_client 1.0.2 → 1.1.3

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 893d3e23c931f4f67621e03306cb3523e2bf86affe3a4c7bd6958e5cd4a7aa12
-  data.tar.gz: 1fbcebdefffabad096aee653ab255329696b8e5b6e2056f4742a619ec15e3810
+  metadata.gz: 4c739774fd292daacd00b2b803065c89d8244a309f2ef09936b107360b4b2cbf
+  data.tar.gz: 7a2aac170572e82baaf78a8923465e083b6c485e75191842c7a7cad39922e1e1
 SHA512:
-  metadata.gz: 4e7ac1fc3d9b850fc0d7d9261c1db44664511997ec17789dd153d9faf9296af7caa689561d30cfde7e07918d43331fbd11b92a1704c9781b489bb3b17704a4e3
-  data.tar.gz: 59ca2f54c33e271ffa93d59c224fa98ae11b729152407b5b7691413bce1ca2b4a2ec6b14f9a6c6253cade2fe1bee290aa12f6496daa821441e9d7e9092b3cf91
+  metadata.gz: a9c74e88514542dba44ca726bd1c085d89ad308b56e2e094c4459055db4e93ae8589d090cb87e2df05d980c4a8d312b1985ab365015c19cbad57ea3ce0e30838
+  data.tar.gz: cd8d07630fc414c51aa7bec61d8abf5c3dabc93239db9e824577def9a56f7affcc39dd0e38190644f32081e188d5b13d4551b0e93999c76d90698358a75ed5e9

data/lib/queuery_client/client.rb

@@ -4,18 +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
@@ -24,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'
@@ -54,7 +57,9 @@ module QueueryClient
         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
@@ -74,11 +79,27 @@ module QueueryClient
       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'
+          value # Ruby does not have a class without timezone
+        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.2"
+  VERSION = "1.1.3"
 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

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: queuery_client
 version: !ruby/object:Gem::Version
-  version: 1.0.2
+  version: 1.1.3
 platform: ruby
 authors:
 - Hidekazu Kobayashi
-autorequire: 
+autorequire:
 bindir: exe
 cert_chain: []
-date: 2020-07-22 00:00:00.000000000 Z
+date: 2021-11-12 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: garage_client
@@ -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.1.2
-signing_key: 
+rubygems_version: 3.1.6
+signing_key:
 specification_version: 4
 summary: Client library for Queuery Redshift HTTP API
 test_files: []