redshift-connector 4.3.0

data/lib/redshift-connector/importer/rebuild_truncate.rb

@@ -0,0 +1,31 @@
+require 'redshift-connector/importer/activerecord-import'
+require 'redshift-connector/logger'
+
+module RedshiftConnector
+  class Importer::RebuildTruncate
+    def initialize(dao:, bundle:, columns:, logger: RedshiftConnector.logger)
+      @dao = dao
+      @bundle = bundle
+      @columns = columns
+      @logger = logger
+    end
+
+    def execute
+      truncate_table(@dao.table_name)
+      import
+    end
+
+    def truncate_table(table_name)
+      @logger.info "TRUNCATE #{table_name}"
+      @dao.connection.execute("truncate #{table_name}")
+      @logger.info "truncated."
+    end
+
+    def import
+      @logger.info "IMPORT #{@bundle.url}* -> #{@dao.table_name} (#{@columns.join(', ')})"
+      @bundle.each_batch do |rows|
+        @dao.import(@columns, rows)
+      end
+    end
+  end
+end

data/lib/redshift-connector/importer/upsert.rb

@@ -0,0 +1,25 @@
+require 'redshift-connector/importer/activerecord-import'
+require 'redshift-connector/logger'
+
+module RedshiftConnector
+  class Importer::Upsert
+    def initialize(dao:, bundle:, columns:, upsert_columns:, logger: RedshiftConnector.logger)
+      @dao = dao
+      @bundle = bundle
+      @columns = columns
+      @upsert_columns = upsert_columns
+      @logger = logger
+    end
+
+    def execute
+      import
+    end
+
+    def import
+      @logger.info "IMPORT #{@bundle.url}* -> #{@dao.table_name} (#{@columns.join(', ')}) upsert (#{@upsert_columns.join(', ')})"
+      @bundle.each_batch do |rows|
+        @dao.import(@columns, rows, on_duplicate_key_update: @upsert_columns)
+      end
+    end
+  end
+end

data/lib/redshift-connector/logger.rb

@@ -0,0 +1,20 @@
+module RedshiftConnector
+  @logger = nil
+
+  def RedshiftConnector.logger
+    # Defer to access Rails
+    @logger || Rails.logger
+  end
+
+  def RedshiftConnector.logger=(logger)
+    @logger = logger
+  end
+
+  class NullLogger
+    def noop(*args) end
+    alias error noop
+    alias warn noop
+    alias info noop
+    alias debug noop
+  end
+end

data/lib/redshift-connector/query.rb

@@ -0,0 +1,93 @@
+module RedshiftConnector
+  class DeltaQuery
+    def initialize(schema:, table:, columns:, condition: nil)
+      @schema = schema
+      @table = table
+      @columns = columns
+      @condition = condition
+    end
+
+    def table_spec
+      "#{@schema}.#{@table}"
+    end
+
+    def description
+      "#{table_spec} (#{@columns.join(', ')}) where (#{@condition})"
+    end
+
+    def to_sql
+      "select #{@columns.map {|c| %Q("#{c}") }.join(', ')}" \
+          + " from #{table_spec}" \
+          + (@condition ? " where #{@condition}" : '')
+    end
+  end
+
+  class SelectAllQuery
+    def initialize(schema:, table:, columns:)
+      @schema = schema
+      @table = table
+      @columns = columns
+    end
+
+    def table_spec
+      "#{@schema}.#{@table}"
+    end
+
+    def description
+      "#{table_spec} (#{@columns.join(', ')})"
+    end
+
+    def to_sql
+      "select #{@columns.map {|c| %Q("#{c}") }.join(', ')}" \
+          + " from #{table_spec}"
+    end
+  end
+
+  class UnloadQuery
+    def UnloadQuery.wrap(query:, bundle:)
+      new(query: ArbitraryQuery.new(query), bundle: bundle)
+    end
+
+    def initialize(query:, bundle:)
+      @query = query
+      @bundle = bundle
+    end
+
+    def table_spec
+      @query.table_spec
+    end
+
+    def description
+      @query.description
+    end
+
+    def to_sql
+      <<-EndSQL.gsub(/^\s+/, '')
+        unload ('#{escape_query(@query.to_sql)}')
+        to '#{@bundle.url}'
+        credentials '#{@bundle.credential_string}'
+        gzip
+        allowoverwrite
+        delimiter ',' escape addquotes
+      EndSQL
+    end
+
+    def escape_query(query)
+      query.gsub("'", "\\\\'")
+    end
+  end
+
+  class ArbitraryQuery
+    def initialize(query)
+      @query = query
+    end
+
+    def description
+      @query
+    end
+
+    def to_sql
+      @query
+    end
+  end
+end

data/lib/redshift-connector/reader.rb

@@ -0,0 +1,18 @@
+# create module
+module RedshiftConnector
+  module Reader
+  end
+end
+
+require 'redshift-connector/reader/redshift_csv'
+require 'redshift-connector/reader/csv'
+require 'redshift-connector/reader/tsv'
+require 'redshift-connector/reader/exception'
+
+module RedshiftConnector
+  module Reader
+    def Reader.get(id)
+      Abstract.get_reader_class(id)
+    end
+  end
+end

data/lib/redshift-connector/reader/abstract.rb

@@ -0,0 +1,18 @@
+module RedshiftConnector
+  class Reader::Abstract
+    READER_CLASSES = {}   # {Symbol => Class}
+
+    def self.declare_reader(id)
+      READER_CLASSES[id.to_sym] = self
+    end
+
+    def self.get_reader_class(id)
+      READER_CLASSES[id.to_sym] or
+          raise ArgumentError, "unknown data file reader type: #{id.inspect}"
+    end
+  end
+
+  def self.get_reader_class(id)
+    Reader::Abstract.get_reader_class(id)
+  end
+end

data/lib/redshift-connector/reader/csv.rb

@@ -0,0 +1,24 @@
+require 'redshift-connector/reader/abstract'
+require 'redshift-connector/reader/exception'
+require 'csv'
+
+module RedshiftConnector
+  # Parses (standard) CSV files.
+  # For UNLOAD-generated CSV, use RedshiftCSV class.
+  class Reader::CSV < Reader::Abstract
+    declare_reader :csv
+
+    def self.data_object?(obj)
+      /\.csv(?:\.|\z)/ =~ File.basename(obj.key)
+    end
+
+    def initialize(f)
+      @f = f
+    end
+
+    def each(&block)
+      csv = CSV.new(@f)
+      csv.each(&block)
+    end
+  end
+end

data/lib/redshift-connector/reader/exception.rb

@@ -0,0 +1,3 @@
+module RedshiftConnector
+  class Reader::MalformedCSVException < StandardError; end
+end

data/lib/redshift-connector/reader/redshift_csv.rb

@@ -0,0 +1,54 @@
+require 'redshift-connector/reader/abstract'
+require 'redshift-connector/reader/exception'
+
+module RedshiftConnector
+  # Reads CSV file generated by Redshift UNLOAD statement (with option ADDQUOTES ESCAPE).
+  # UNLOAD escapes data by '\' (backslash character), we cannot use standard CSV class.
+  class Reader::RedshiftCSV < Reader::Abstract
+    declare_reader :redshift_csv
+
+    def self.data_object?(obj)
+      /\.csv(?:\.|\z)/ =~ File.basename(obj.key)
+    end
+
+    # f :: IO
+    def initialize(f)
+      @f = f
+    end
+
+    def each
+      # We can use simple #each_line to read single row
+      # because line terminators are always escaped by UNLOAD.
+      @f.each_line do |line|
+        yield parse_row(line, @f.lineno)
+      end
+    end
+
+    def parse_row(line, lineno = nil)
+      row = []
+      s = StringScanner.new(line)
+      s.skip(/\s+/)
+      until s.eos?
+        col = s.scan(/"(?:\\.|[^"\\]+)*"/) or raise MalformedCSVException, "CSV parse error at line #{lineno}"
+        row.push unescape_column(col)
+        s.skip(/\s*/)    # skip line terminator on line ends
+        s.skip(/,\s*/)
+      end
+      row
+    end
+
+    UNESCAPE_MAP = {
+      '\\"' => '"',
+      "\\'" => "'",
+      '\\,' => ',',
+      '\\r' => "\r",
+      '\\n' => "\n",
+      '\\\\' => '\\'
+    }
+
+    def unescape_column(col)
+      charmap = UNESCAPE_MAP
+      col[1...-1].gsub(/\\./) {|s| charmap[s] }
+    end
+  end
+end

data/lib/redshift-connector/reader/tsv.rb

@@ -0,0 +1,24 @@
+require 'redshift-connector/reader/abstract'
+require 'redshift-connector/reader/exception'
+require 'csv'
+
+module RedshiftConnector
+  # Parses TSV (Tab Separated Format) files.
+  class Reader::TSV < Reader::Abstract
+    declare_reader :tsv
+
+    def self.data_object?(obj)
+      /\.tsv(?:\.|\z)/ =~ File.basename(obj.key)
+    end
+
+    def initialize(f)
+      @f = f
+    end
+
+    def each(&block)
+      @f.each_line do |line|
+        yield line.chomp.split("\t", -1)
+      end
+    end
+  end
+end

data/lib/redshift-connector/s3_bucket.rb

@@ -0,0 +1,72 @@
+require 'aws-sdk'
+
+module RedshiftConnector
+  class S3Bucket
+    @buckets = {}
+    @default = nil
+
+    def S3Bucket.add(name, default: false, **params)
+      instance = new(**params)
+      @buckets[name.to_s] = instance
+      if !@default or default
+        @default = instance
+      end
+    end
+
+    def S3Bucket.default
+      @default or raise ArgumentError, "no default S3 bucket configured"
+    end
+
+    def S3Bucket.get(name)
+      @buckets[name.to_s] or raise ArgumentError, "no such S3 bucket configured: #{name.inspect}"
+    end
+
+    def initialize(bucket:, prefix: nil, access_key_id: nil, secret_access_key: nil, iam_role: nil)
+      @name = bucket
+      @prefix = prefix
+      @access_key_id = access_key_id
+      @secret_access_key = secret_access_key
+      @iam_role = iam_role
+    end
+
+    attr_reader :name
+    attr_reader :prefix
+
+    def url
+      "s3://#{@bucket.name}/#{@prefix}/"
+    end
+
+    def client
+      @client ||= Aws::S3::Client.new(access_key_id: @access_key_id, secret_access_key: @secret_access_key)
+    end
+
+    def bucket
+      @bucket ||= begin
+        resource = Aws::S3::Resource.new(client: client)
+        resource.bucket(@name)
+      end
+    end
+
+    def object(key)
+      bucket.object(key)
+    end
+
+    def objects(prefix:)
+      bucket.objects(prefix: prefix)
+    end
+
+    def delete_objects(keys)
+      bucket.delete_objects(delete: {objects: keys.map {|k| {key: k} }})
+    end
+
+    def credential_string
+      if @iam_role
+        "aws_iam_role=#{@iam_role}"
+      elsif @access_key_id
+        "aws_access_key_id=#{@access_key_id};aws_secret_access_key=#{@secret_access_key}"
+      else
+        raise ArgumentError, "no credential given for Redshift S3 access"
+      end
+    end
+  end
+end

data/lib/redshift-connector/s3_data_file.rb

@@ -0,0 +1,34 @@
+require 'zlib'
+
+module RedshiftConnector
+  class S3DataFile
+    def initialize(object, reader_class:)
+      @object = object
+      @reader_class = reader_class
+    end
+
+    def key
+      @object.key
+    end
+
+    def each_row(&block)
+      response = @object.get
+      f = if gzipped_object?
+        Zlib::GzipReader.new(response.body)
+      else
+        response.body
+      end
+      @reader_class.new(f).each(&block)
+    ensure
+      response.body.close if response
+    end
+
+    def data_object?
+      @reader_class.data_object?(@object)
+    end
+
+    def gzipped_object?
+      File.extname(@object.key) == '.gz'
+    end
+  end
+end

data/lib/redshift-connector/s3_data_file_bundle.rb

@@ -0,0 +1,101 @@
+require 'redshift-connector/s3_bucket'
+require 'redshift-connector/s3_data_file'
+require 'redshift-connector/reader'
+require 'redshift-connector/logger'
+require 'aws-sdk'
+
+module RedshiftConnector
+  class S3DataFileBundle
+    def self.for_prefix(bucket: S3Bucket.default, prefix:, format:, filter: nil, batch_size: 1000, logger: RedshiftConnector.logger)
+      real_prefix = "#{bucket.prefix}/#{prefix}"
+      new(bucket, real_prefix, format: format, filter: filter, batch_size: batch_size, logger: logger)
+    end
+
+    def self.for_table(bucket: S3Bucket.default, schema:, table:, txn_id:, filter: nil, batch_size: 1000, logger: RedshiftConnector.logger)
+      prefix = "#{bucket.prefix}/#{schema}_export/#{table}/#{txn_id}/#{table}.csv."
+      new(bucket, prefix, format: :redshift_csv, filter: filter, batch_size: batch_size, logger: logger)
+    end
+
+    def initialize(bucket, prefix, format: :csv, filter: nil, batch_size: 1000, logger: RedshiftConnector.logger)
+      @bucket = bucket
+      @prefix = prefix
+      @format = format
+      @filter = filter || lambda {|*row| row }
+      @batch_size = batch_size
+      @logger = logger
+      @reader_class = Reader.get(format)
+    end
+
+    attr_reader :bucket
+    attr_reader :prefix
+
+    def url
+      "s3://#{@bucket.name}/#{@prefix}"
+    end
+
+    def credential_string
+      @bucket.credential_string
+    end
+
+    REPORT_SIZE = 10_0000
+
+    def each_batch(report: true)
+      @logger.info "reader: #{@reader_class}"
+      n = 0
+      reported = 0
+      do_each_batch(@batch_size) do |rows|
+        yield rows
+        n += rows.size
+        if n / REPORT_SIZE > reported
+          @logger.info "#{n} rows processed" if report
+          reported = n / REPORT_SIZE
+        end
+      end
+      @logger.info "total #{n} rows processed" if report
+    end
+
+    def do_each_batch(batch_size)
+      filter = @filter
+      buf = []
+      each_row do |row|
+        buf.push filter.(*row)
+        if buf.size == batch_size
+          yield buf
+          buf = []
+        end
+      end
+      yield buf unless buf.empty?
+    end
+    private :do_each_batch
+
+    def each_row(&block)
+      each_object do |obj|
+        obj.each_row(&block)
+      end
+    end
+
+    alias each each_row
+
+    def each_object(&block)
+      all_data_objects.each do |obj|
+        @logger.info "processing s3 object: #{obj.key}"
+        yield obj
+      end
+    end
+
+    def all_data_objects
+      @bucket.objects(prefix: @prefix)
+        .map {|obj| S3DataFile.new(obj, reader_class: @reader_class) }
+        .select {|obj| obj.data_object? }
+    end
+
+    def clear
+      pref = File.dirname(@prefix) + '/'
+      keys = @bucket.objects(prefix: pref).map(&:key)
+      unless keys.empty?
+        @logger.info "DELETE #{pref}*"
+        @bucket.delete_objects(keys)
+      end
+    end
+  end
+end