redshift_connector 8.0.0

data/lib/redshift_connector/data_file.rb

@@ -0,0 +1,32 @@
+require 'redshift_connector/logger'
+require 'zlib'
+
+module RedshiftConnector
+  class DataFile
+    def initialize(reader_class:)
+      @reader_class = reader_class
+    end
+
+    def each_row(&block)
+      f = open
+      begin
+        if gzipped_object?
+          f = Zlib::GzipReader.new(f)
+        end
+        @reader_class.new(f).each(&block)
+      ensure
+        f.close
+      end
+    end
+
+    # abstract open
+
+    def data_object?
+      @reader_class.data_object?(key)
+    end
+
+    def gzipped_object?
+      File.extname(key) == '.gz'
+    end
+  end
+end

data/lib/redshift_connector/data_file_bundle_params.rb

@@ -0,0 +1,25 @@
+require 'redshift_connector/logger'
+
+module RedshiftConnector
+  class DataFileBundleParams
+    def initialize(
+      bucket: nil,
+      schema:,
+      table:,
+      txn_id: nil,
+      logger: RedshiftConnector.logger
+    )
+      @bucket = bucket
+      @schema = schema
+      @table = table
+      @txn_id = txn_id
+      @logger = logger
+    end
+
+    attr_reader :bucket
+    attr_reader :schema
+    attr_reader :table
+    attr_reader :txn_id
+    attr_reader :logger
+  end
+end

data/lib/redshift_connector/data_file_bundle_reader.rb

@@ -0,0 +1,72 @@
+require 'redshift_connector/logger'
+require 'forwardable'
+
+module RedshiftConnector
+  class DataFileBundleReader
+    extend Forwardable
+
+    DEFAULT_BATCH_SIZE = 1000
+
+    def initialize(bundle, filter: nil, batch_size: DEFAULT_BATCH_SIZE, logger: RedshiftConnector.logger)
+      @bundle = bundle
+      @filter = filter || lambda {|*row| row }
+      @batch_size = batch_size || 1000
+      @logger = logger
+    end
+
+    attr_reader :bundle
+    attr_reader :batch_size
+    attr_reader :logger
+
+    def_delegators '@bundle', :url, :bucket, :key
+
+    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
+      @bundle.data_files.select {|obj| obj.data_object? }
+    end
+
+    REPORT_SIZE = 10_0000
+
+    def each_batch(report: true)
+      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
+  end
+end

data/lib/redshift_connector/exception.rb

@@ -0,0 +1,5 @@
+module RedshiftConnector
+  class Error < ::StandardError; end
+  class ExportError < Error; end
+  class ImportError < Error; end
+end

data/lib/redshift_connector/exporter.rb

@@ -0,0 +1,40 @@
+module RedshiftConnector
+  module Exporter
+    @default_data_source = nil
+
+    def Exporter.default_data_source=(ds)
+      @default_data_source = ds
+    end
+
+    def Exporter.default_data_source
+      @default_data_source or raise ArgumentError, "RedshiftConnector::Exporter.default_data_source was not set"
+    end
+
+    def Exporter.builder
+      default_data_source.exporter_builder
+    end
+
+    def Exporter.for_table_delta(**params)
+      builder.build_for_table_delta(**params)
+    end
+
+    def Exporter.for_table(**params)
+      builder.build_for_table(**params)
+    end
+
+    def Exporter.for_query(**params)
+      builder.build_for_query(**params)
+    end
+
+    def Exporter.foreach(**params, &block)
+      exporter = for_query(**params)
+      bundle = exporter.execute
+      r = DataFileBundleReader.new(bundle, logger: bundle.logger)
+      begin
+        r.each_row(&block)
+      ensure
+        bundle.clear if bundle.respond_to?(:clear)
+      end
+    end
+  end
+end

data/lib/redshift_connector/exporter_builder.rb

@@ -0,0 +1,49 @@
+require 'redshift_connector/query'
+require 'redshift_connector/logger'
+
+module RedshiftConnector
+  class ExporterBuilder
+    def initialize(ds:, exporter_class:)
+      @ds = ds
+      @exporter_class = exporter_class
+    end
+
+    def build_for_table_delta(schema:, table:, condition:, columns:, bundle_params:, logger: RedshiftConnector.logger)
+      query = DeltaQuery.new(schema: schema, table: table, columns: columns, condition: condition)
+      @exporter_class.new(ds: @ds, query: query, bundle_params: bundle_params, logger: logger)
+    end
+
+    def build_for_table(schema:, table:, columns:, bundle_params:, logger: RedshiftConnector.logger)
+      query = SelectAllQuery.new(schema: schema, table: table, columns: columns)
+      @exporter_class.new(ds: @ds, query: query, bundle_params: bundle_params, logger: logger)
+    end
+
+    def build_for_query(
+        schema:,
+        table:,
+        bucket: nil,
+        query:,
+        txn_id: "#{Time.now.strftime('%Y%m%d_%H%M%S')}_#{$$}",
+        enable_sort: false,
+        logger: RedshiftConnector.logger,
+        quiet: false
+    )
+      logger = NullLogger.new if quiet
+      bundle_params = DataFileBundleParams.new(
+        bucket: bucket,
+        schema: schema,
+        table: table,
+        txn_id: txn_id,
+        logger: logger
+      )
+      @exporter_class.new(
+        ds: @ds,
+        query: ArbitraryQuery.new(query),
+        bundle_params: bundle_params,
+        enable_sort: enable_sort,
+        logger: logger
+      )
+    end
+
+  end
+end

data/lib/redshift_connector/immediate_exporter.rb

@@ -0,0 +1,19 @@
+require 'redshift_connector/s3_data_file_bundle'
+require 'redshift_connector/logger'
+
+module RedshiftConnector
+  class ImmediateExporter
+    def initialize(bundle:, logger: RedshiftConnector.logger)
+      @bundle = bundle
+      @logger = logger
+    end
+
+    attr_reader :bundle
+    attr_reader :logger
+
+    def execute
+      @logger.info "USE #{@bundle.url}*"
+      @bundle
+    end
+  end
+end

data/lib/redshift_connector/importer.rb

@@ -0,0 +1,58 @@
+# create module
+module RedshiftConnector
+  module Importer
+  end
+end
+
+require 'redshift_connector/importer/upsert'
+require 'redshift_connector/importer/insert_delta'
+require 'redshift_connector/importer/rebuild_rename'
+require 'redshift_connector/importer/rebuild_truncate'
+require 'redshift_connector/logger'
+
+module RedshiftConnector
+  module Importer
+    def Importer.for_delta_upsert(table:, columns:, delete_cond: nil, upsert_columns: nil, logger: RedshiftConnector.logger)
+      if delete_cond and upsert_columns
+        raise ArgumentError, "delete_cond and upsert_columns are exclusive"
+      end
+      importer =
+        if delete_cond
+          Importer::InsertDelta.new(
+            dao: table.classify.constantize,
+            columns: columns,
+            delete_cond: delete_cond,
+            logger: logger
+          )
+        elsif upsert_columns
+          Importer::Upsert.new(
+            dao: table.classify.constantize,
+            columns: columns,
+            upsert_columns: upsert_columns,
+            logger: logger
+          )
+        else
+          raise ArgumentError, "either of delete_cond or upsert_columns is required for delta import"
+        end
+      importer
+    end
+
+    def Importer.for_rebuild(strategy: 'rename', table:, columns:, logger: RedshiftConnector.logger)
+      c = get_rebuild_class(strategy)
+      c.new(
+        dao: table.classify.constantize,
+        columns: columns,
+        logger: logger
+      )
+    end
+
+    def Importer.get_rebuild_class(strategy)
+      case strategy.to_s
+      when 'rename' then RebuildRename
+      when 'truncate' then RebuildTruncate
+      else
+        raise ArgumentError, "unsupported rebuild strategy: #{strategy.inspect}"
+      end
+    end
+  end
+end

data/lib/redshift_connector/importer/activerecord-import.rb

@@ -0,0 +1,2 @@
+require 'activerecord-import'
+require 'activerecord-import/base'

data/lib/redshift_connector/importer/insert_delta.rb

@@ -0,0 +1,31 @@
+require 'redshift_connector/importer/activerecord-import'
+require 'redshift_connector/logger'
+
+module RedshiftConnector
+  class Importer::InsertDelta
+    def initialize(dao:, columns:, delete_cond:, logger: RedshiftConnector.logger)
+      @dao = dao
+      @columns = columns
+      @delete_cond = delete_cond
+      @logger = logger
+    end
+
+    def execute(bundle)
+      delete_rows(@delete_cond)
+      import(bundle)
+    end
+
+    def delete_rows(cond_expr)
+      @logger.info "DELETE #{@dao.table_name} where (#{cond_expr})"
+      @dao.connection.execute("delete from #{@dao.table_name} where #{cond_expr}")
+      @logger.info "deleted."
+    end
+
+    def import(bundle)
+      @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/rebuild_rename.rb

@@ -0,0 +1,58 @@
+require 'redshift_connector/importer/activerecord-import'
+require 'redshift_connector/logger'
+require 'thread'
+
+module RedshiftConnector
+  class Importer::RebuildRename
+    def initialize(dao:, columns:, logger: RedshiftConnector.logger)
+      @dao = dao
+      @columns = columns
+      @logger = logger
+    end
+
+    def execute(bundle)
+      dest_table = @dao.table_name
+      tmp_table = "#{dest_table}_new"
+      old_table = "#{dest_table}_old"
+
+      tmp_dao = self.class.make_temporary_dao(@dao)
+      tmp_dao.table_name = tmp_table
+
+      exec_update "drop table if exists #{tmp_table}"
+      exec_update "create table #{tmp_table} like #{dest_table}"
+      import(tmp_dao, bundle)
+      exec_update "drop table if exists #{old_table}"
+      # Atomic table exchange
+      exec_update "rename table #{dest_table} to #{old_table}, #{tmp_table} to #{dest_table}"
+    end
+
+    # Duplicates DAO (ActiveRecord class) and names it.
+    # Newer activerecord-import requires a class name (not a table name),
+    # we must prepare some name for temporary DAO class.
+    def self.make_temporary_dao(orig)
+      tmp = orig.dup
+      const_set("TemporaryDAO_#{get_unique_sequence}", tmp)
+      tmp.name   # fix class name
+      tmp
+    end
+
+    @dao_seq = 0
+    @dao_seq_lock = Mutex.new
+
+    def self.get_unique_sequence
+      @dao_seq_lock.synchronize { @dao_seq += 1 }
+    end
+
+    def exec_update(query)
+      @logger.info query
+      @dao.connection.execute(query)
+    end
+
+    def import(dao, bundle)
+      @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/rebuild_truncate.rb

@@ -0,0 +1,30 @@
+require 'redshift_connector/importer/activerecord-import'
+require 'redshift_connector/logger'
+
+module RedshiftConnector
+  class Importer::RebuildTruncate
+    def initialize(dao:, columns:, logger: RedshiftConnector.logger)
+      @dao = dao
+      @columns = columns
+      @logger = logger
+    end
+
+    def execute(bundle)
+      truncate_table(@dao.table_name)
+      import(bundle)
+    end
+
+    def truncate_table(table_name)
+      @logger.info "TRUNCATE #{table_name}"
+      @dao.connection.execute("truncate #{table_name}")
+      @logger.info "truncated."
+    end
+
+    def import(bundle)
+      @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,24 @@
+require 'redshift_connector/importer/activerecord-import'
+require 'redshift_connector/logger'
+
+module RedshiftConnector
+  class Importer::Upsert
+    def initialize(dao:, columns:, upsert_columns:, logger: RedshiftConnector.logger)
+      @dao = dao
+      @columns = columns
+      @upsert_columns = upsert_columns
+      @logger = logger
+    end
+
+    def execute(bundle)
+      import(bundle)
+    end
+
+    def import(bundle)
+      @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