<commit_before>require 'sequel'
require 'yaml'
require 'ostruct'
require 'logger'
require 'fastandand'

Sequel.extension :migration

# TODO possibly use Gem::Package::TarWriter to write tar files
# TODO when restoring, could use a SizeQueue to make sure the db is kept busy

# TODO need to version the dumps, or something like that.
# So the slowest-changing variables are the db, the io stream
# and the page size.
# table will change every call. Will IO stream change between
# table changes? No. So a currying type approach will work.
# Somebody must have done this before.
# But table and io are often related (ie table going to one file)
# TODO This really should be Wyrm::Hole. Or maybe Wyrm::Hole should
# be the codec that connects two DbPumps, for direct transfer?
class DbPump
  # some codecs might ignore io, eg if a dbpump is talking to another dbpump
  def initialize( db, table_name, io: STDOUT, codec: :marshal, page_size: 10000, dry_run: false )
    self.codec = codec
    self.db = db
    self.table_name = table_name
    self.io = io
    self.page_size = page_size
    self.dry_run = dry_run
    yield self if block_given?
  end

  attr_accessor :io, :page_size, :dry_run

  # These affect cached values
  attr_reader :db, :table_name

  def table_name=( name_sym )
    @primary_keys = nil
    @table_dataset = nil
    @table_name = name_sym
  end

  def db=( other_db )
    @primary_keys = nil
    @table_dataset = nil
    @db = other_db
    @db.extension :pagination
  end

  def dry_run?; dry_run; end

  class RespondsTo
    def initialize( *methods )
      @methods = methods
    end

    def ===( instance )
      @methods.all?{|m| instance.respond_to? m}
    end
  end

  def codec=( codec_thing )
    @codec =
    case codec_thing
    when :yaml; YamlCodec.new
    when :marshal; MarshalCodec.new
    when Class
      codec.new
    when RespondsTo.new( :encode, :decode )
      codec
    else
      raise "unknown codec #{codec}"
    end
  end

  attr_reader :codec

  # TODO could use msgpack as serialization here, but its API is unpleasant.

  class MarshalCodec
    def encode( obj, io )
      Marshal.dump obj, io
    end

    def decode( io, &block )
      obj = Marshal.load(io)
      yield obj if block_given?
      obj
    end
  end

  class MsgPackCodec
    def encode( obj, io )
      Marshal.dump obj, io
    end

    def decode( io, &block )
      obj = Marshal.load(io)
      yield obj if block_given?
      obj
    end
  end

  class YamlCodec
    def encode( obj, io )
      YAML.dump obj, io
    end

    def decode( io, &block )
      obj = YAML.load(io)
      yield obj if block_given?
      obj
    end
  end

  def logger
    @logger ||= Logger.new STDERR
  end

  def primary_keys
    @primary_keys ||= db.schema(table_name).select{|df| df.last[:primary_key]}.map{|df| df.first}
  end

  def table_dataset
    @table_dataset ||= db[table_name.to_sym]
  end

  # TODO possibly use select from outer / inner join to
  # http://www.numerati.com/2012/06/26/reading-large-result-sets-with-hibernate-and-mysql/
  # because mysql is useless
  def paginated_dump
    table_dataset.order(*primary_keys).each_page(page_size) do |page|
      logger.info page.sql
      page.each do |row|
        unless dry_run?
          codec.encode row.values, io
        end
      end
    end
  end

  # have to use this for non-integer pks
  # The idea is that large offsets are expensive in the db because the db server has to read
  # through the data set to reach the required offset. So make that only ids, and then
  # do the main select from the limited id list.
  # TODO could speed this up by have a query thread which runs the next page-query while
  # the current one is being written/compressed.
  # select * from massive as full
  #   inner join (select id from massive order by whatever limit m, n) limit
  #   on full.id = limit.id
  # order by full.whatever
  def inner_dump
    # could possibly overrride Dataset#paginate(page_no, page_size, record_count=nil)
    0.step(table_dataset.count, page_size).each do |offset|
      limit_dataset = table_dataset.select( *primary_keys ).limit( page_size, offset ).order( *primary_keys )
      page = table_dataset.join( limit_dataset, Hash[ primary_keys.map{|f| [f,f]} ] ).order( *primary_keys ).qualify_to(table_name)
      logger.info page.sql
      page.each do |row|
        unless dry_run?
          codec.encode row.values, io
        end
      end
    end
  end

  # TODO need to also dump a first row containing useful stuff:
  # - source table name<commit_msg>qualify_to is deprecated<commit_after>require 'sequel'
require 'yaml'
require 'ostruct'
require 'logger'
require 'fastandand'

Sequel.extension :migration

# TODO possibly use Gem::Package::TarWriter to write tar files
# TODO when restoring, could use a SizeQueue to make sure the db is kept busy

# TODO need to version the dumps, or something like that.
# So the slowest-changing variables are the db, the io stream
# and the page size.
# table will change every call. Will IO stream change between
# table changes? No. So a currying type approach will work.
# Somebody must have done this before.
# But table and io are often related (ie table going to one file)
# TODO This really should be Wyrm::Hole. Or maybe Wyrm::Hole should
# be the codec that connects two DbPumps, for direct transfer?
class DbPump
  # some codecs might ignore io, eg if a dbpump is talking to another dbpump
  def initialize( db, table_name, io: STDOUT, codec: :marshal, page_size: 10000, dry_run: false )
    self.codec = codec
    self.db = db
    self.table_name = table_name
    self.io = io
    self.page_size = page_size
    self.dry_run = dry_run
    yield self if block_given?
  end

  attr_accessor :io, :page_size, :dry_run

  # These affect cached values
  attr_reader :db, :table_name

  def table_name=( name_sym )
    @primary_keys = nil
    @table_dataset = nil
    @table_name = name_sym
  end

  def db=( other_db )
    @primary_keys = nil
    @table_dataset = nil
    @db = other_db
    @db.extension :pagination
  end

  def dry_run?; dry_run; end

  class RespondsTo
    def initialize( *methods )
      @methods = methods
    end

    def ===( instance )
      @methods.all?{|m| instance.respond_to? m}
    end
  end

  def codec=( codec_thing )
    @codec =
    case codec_thing
    when :yaml; YamlCodec.new
    when :marshal; MarshalCodec.new
    when Class
      codec.new
    when RespondsTo.new( :encode, :decode )
      codec
    else
      raise "unknown codec #{codec}"
    end
  end

  attr_reader :codec

  # TODO could use msgpack as serialization here, but its API is unpleasant.

  class MarshalCodec
    def encode( obj, io )
      Marshal.dump obj, io
    end

    def decode( io, &block )
      obj = Marshal.load(io)
      yield obj if block_given?
      obj
    end
  end

  class MsgPackCodec
    def encode( obj, io )
      Marshal.dump obj, io
    end

    def decode( io, &block )
      obj = Marshal.load(io)
      yield obj if block_given?
      obj
    end
  end

  class YamlCodec
    def encode( obj, io )
      YAML.dump obj, io
    end

    def decode( io, &block )
      obj = YAML.load(io)
      yield obj if block_given?
      obj
    end
  end

  def logger
    @logger ||= Logger.new STDERR
  end

  def primary_keys
    @primary_keys ||= db.schema(table_name).select{|df| df.last[:primary_key]}.map{|df| df.first}
  end

  def table_dataset
    @table_dataset ||= db[table_name.to_sym]
  end

  # TODO possibly use select from outer / inner join to
  # http://www.numerati.com/2012/06/26/reading-large-result-sets-with-hibernate-and-mysql/
  # because mysql is useless
  def paginated_dump
    table_dataset.order(*primary_keys).each_page(page_size) do |page|
      logger.info page.sql
      page.each do |row|
        unless dry_run?
          codec.encode row.values, io
        end
      end
    end
  end

  # have to use this for non-integer pks
  # The idea is that large offsets are expensive in the db because the db server has to read
  # through the data set to reach the required offset. So make that only ids, and then
  # do the main select from the limited id list.
  # TODO could speed this up by have a query thread which runs the next page-query while
  # the current one is being written/compressed.
  # select * from massive as full
  #   inner join (select id from massive order by whatever limit m, n) limit
  #   on full.id = limit.id
  # order by full.whatever
  def inner_dump
    # could possibly overrride Dataset#paginate(page_no, page_size, record_count=nil)
    0.step(table_dataset.count, page_size).each do |offset|
      limit_dataset = table_dataset.select( *primary_keys ).limit( page_size, offset ).order( *primary_keys )
      page = table_dataset.join( limit_dataset, Hash[ primary_keys.map{|f| [f,f]} ] ).order( *primary_keys ).qualify(table_name)
      logger.info page.sql
      page.each do |row|
        unless dry_run?
          codec.encode row.values, io
        end
      end
    end
  end

  # TODO need to also dump a first row containing useful stuff:
  # - source table name
{'fraction_non_alphanumeric': 0.07877989905639675, 'fraction_numerical': 0.0030721966205837174, 'mean_word_length': 2.805845511482255, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2805475', 'n_tokens_mistral': 2965, 'n_tokens_neox': 2837, 'n_words': 1059}