Class: Tapsoob::Push

Inherits:

Operation

• Object
• Operation
• Tapsoob::Push

Defined in:

lib/tapsoob/operation.rb

Instance Attribute Summary

Attributes inherited from Operation

#database_url, #dump_path, #opts

Instance Method Summary

• #fetch_local_tables_info ⇒ Object
• #file_prefix ⇒ Object
• #local_tables_info ⇒ Object
• #push_data ⇒ Object
• #push_data_from_file(stream, progress) ⇒ Object
• #push_data_parallel ⇒ Object
• #push_data_serial ⇒ Object
• #push_indexes ⇒ Object
• #push_partial_data ⇒ Object
• #push_reset_sequences ⇒ Object
• #push_schema ⇒ Object
• #record_count ⇒ Object
• #run ⇒ Object
• #tables ⇒ Object
• #to_hash ⇒ Object

Methods inherited from Operation

#add_completed_table, #apply_table_filter, #catch_errors, #completed_tables, #completed_tables_mutex, #data?, #db, #default_chunksize, #exclude_tables, #exiting?, factory, #format_number, #indexes_first?, #initialize, #log, #parallel?, #parallel_workers, #resuming?, #schema?, #setup_signal_trap, #store_session, #stream_state, #stream_state=, #table_filter

Constructor Details

This class inherits a constructor from Tapsoob::Operation

Instance Method Details

#fetch_local_tables_info ⇒ Object

# File 'lib/tapsoob/operation.rb', line 650

def fetch_local_tables_info
  tables_with_counts = {}
  tbls = Dir.glob(File.join(dump_path, "schemas", "*")).map { |path| File.basename(path, ".rb") }
  tbls.each do |table|
    if File.exist?(File.join(dump_path, "data", "#{table}.json"))
      # Read NDJSON format - each line is a separate JSON chunk
      total_rows = 0
      File.readlines(File.join(dump_path, "data", "#{table}.json")).each do |line|
        chunk = JSON.parse(line.strip)
        total_rows += chunk["data"].size if chunk["data"]
      end
      tables_with_counts[table] = total_rows
    else
      tables_with_counts[table] = 0
    end
  end
  apply_table_filter(tables_with_counts)
end

#file_prefix ⇒ Object

# File 'lib/tapsoob/operation.rb', line 419

def file_prefix
  "push"
end

#local_tables_info ⇒ Object

# File 'lib/tapsoob/operation.rb', line 633

def local_tables_info
  opts[:local_tables_info] ||= fetch_local_tables_info
end

#push_data ⇒ Object

# File 'lib/tapsoob/operation.rb', line 500

def push_data
  log.info "Sending data"

  log.info "#{tables.size} tables, #{format_number(record_count)} records"

  if parallel?
    push_data_parallel
  else
    push_data_serial
  end
end

#push_data_from_file(stream, progress) ⇒ Object

# File 'lib/tapsoob/operation.rb', line 576

def push_data_from_file(stream, progress)
  loop do
    if exiting?
      store_session
      exit 0
    end

    row_size = 0
    chunksize = stream.state[:chunksize]

    begin
      chunksize = Tapsoob::Utils.calculate_chunksize(chunksize) do |c|
        stream.state[:chunksize] = c.to_i
        encoded_data, row_size, elapsed_time = nil
        d1 = c.time_delta do
          encoded_data, row_size, elapsed_time = stream.fetch({ :type => "file", :source => dump_path })
        end

        data = nil
        d2 = c.time_delta do
          data = {
            :state        => stream.to_hash,
            :checksum     => Tapsoob::Utils.checksum(encoded_data).to_s,
            :encoded_data => encoded_data
          }
        end

        stream.fetch_data_to_database(data)
        log.debug "row size: #{row_size}"
        self.stream_state = stream.to_hash

        c.idle_secs = (d1 + d2)

        elapsed_time
      end
    rescue Tapsoob::CorruptedData => e
      # retry the same data, it got corrupted somehow.
      next
    rescue Tapsoob::DuplicatePrimaryKeyError => e
      # verify the stream and retry it
      stream.verify_stream
      stream = JSON.generate({ :state => stream.to_hash })
      next
    end
    stream.state[:chunksize] = chunksize

    # Update progress bar by 1 chunk
    progress.inc(1) if progress

    break if stream.complete?
  end

  progress.finish if progress
  add_completed_table(stream.table_name)
  self.stream_state = {}
end

#push_data_parallel ⇒ Object

# File 'lib/tapsoob/operation.rb', line 533

def push_data_parallel
  log.info "Using #{parallel_workers} parallel workers"

  multi_progress = opts[:progress] ? MultiProgressBar.new(parallel_workers) : nil
  table_queue = Queue.new

  tables.each do |table_name, count|
    data_file = File.join(dump_path, "data", "#{table_name}.json")
    next unless File.exist?(data_file) && count > 0
    table_queue << [table_name, count]
  end

  workers = (1..parallel_workers).map do
    Thread.new do
      loop do
        break if table_queue.empty?

        table_name, count = table_queue.pop(true) rescue break

        # Each thread gets its own connection from the pool
        db[table_name.to_sym].truncate if @opts[:purge]
        stream = Tapsoob::DataStream.factory(db, {
          :table_name => table_name,
          :chunksize => default_chunksize
        }, {
          :"skip-duplicates" => opts[:"skip-duplicates"] || false,
          :"discard-identity" => opts[:"discard-identity"] || false,
          :purge => opts[:purge] || false,
          :debug => opts[:debug]
        })

        estimated_chunks = [(count.to_f / default_chunksize).ceil, 1].max
        progress = multi_progress ? multi_progress.create_bar(table_name.to_s, estimated_chunks) : nil

        push_data_from_file(stream, progress)
      end
    end
  end

  workers.each(&:join)
  multi_progress.stop if multi_progress
end

#push_data_serial ⇒ Object

# File 'lib/tapsoob/operation.rb', line 512

def push_data_serial
  tables.each do |table_name, count|
    # Skip if data file doesn't exist or has no data
    data_file = File.join(dump_path, "data", "#{table_name}.json")
    next unless File.exist?(data_file) && count > 0
    db[table_name.to_sym].truncate if @opts[:purge]
    stream = Tapsoob::DataStream.factory(db, {
      :table_name => table_name,
      :chunksize => default_chunksize
    }, {
      :"skip-duplicates" => opts[:"skip-duplicates"] || false,
      :"discard-identity" => opts[:"discard-identity"] || false,
      :purge => opts[:purge] || false,
      :debug => opts[:debug]
    })
    estimated_chunks = [(count.to_f / default_chunksize).ceil, 1].max
    progress = (opts[:progress] ? ProgressBar.new(table_name.to_s, estimated_chunks) : nil)
    push_data_from_file(stream, progress)
  end
end

#push_indexes ⇒ Object

# File 'lib/tapsoob/operation.rb', line 441

def push_indexes
  idxs = {}
  table_idxs = Dir.glob(File.join(dump_path, "indexes", "*.json")).map { |path| File.basename(path, '.json') }
  table_idxs.each do |table_idx|
    # Read NDJSON format - each line is a separate index
    index_file = File.join(dump_path, "indexes", "#{table_idx}.json")
    idxs[table_idx] = File.readlines(index_file).map { |line| JSON.parse(line.strip) }
  end

  return unless idxs.size > 0

  log.info "Sending indexes"

  # Calculate max title width for consistent alignment
  filtered_idxs = apply_table_filter(idxs).select { |table, indexes| indexes.size > 0 }
  max_title_width = filtered_idxs.keys.map { |table| "#{table} indexes".length }.max || 14

  filtered_idxs.each do |table, indexes|
    progress = ProgressBar.new("#{table} indexes", indexes.size, STDOUT, max_title_width)
    indexes.each do |idx|
      Tapsoob::Utils.load_indexes(database_url, idx)
      progress.inc(1)
    end
    progress.finish
  end
end

#push_partial_data ⇒ Object

# File 'lib/tapsoob/operation.rb', line 487

def push_partial_data
  return if stream_state == {}

  table_name = stream_state[:table_name]
  record_count = tables[table_name.to_s]
  log.info "Resuming #{table_name}, #{format_number(record_count)} records"
  stream = Tapsoob::DataStream.factory(db, stream_state)
  chunksize = stream_state[:chunksize] || default_chunksize
  estimated_chunks = [(record_count.to_f / chunksize).ceil, 1].max
  progress = (opts[:progress] ? ProgressBar.new(table_name.to_s, estimated_chunks) : nil)
  push_data_from_file(stream, progress)
end

#push_reset_sequences ⇒ Object

# File 'lib/tapsoob/operation.rb', line 481

def push_reset_sequences
  log.info "Resetting sequences"

  Tapsoob::Utils.schema_bin(:reset_db_sequences, database_url)
end

#push_schema ⇒ Object

# File 'lib/tapsoob/operation.rb', line 468

def push_schema
  log.info "Sending schema"

  progress = ProgressBar.new('Schema', tables.size)
  tables.each do |table, count|
    log.debug "Loading '#{table}' schema\n"
    # Reuse existing db connection for better performance
    Tapsoob::Utils.load_schema(dump_path, db, table)
    progress.inc(1)
  end
  progress.finish
end

#record_count ⇒ Object

# File 'lib/tapsoob/operation.rb', line 646

def record_count
  @record_count ||= local_tables_info.values.inject(0) { |a,c| a += c }
end

#run ⇒ Object

# File 'lib/tapsoob/operation.rb', line 427

def run
  catch_errors do
    unless resuming?
      push_schema if schema?
      push_indexes if indexes_first? && schema?
    end
    setup_signal_trap
    push_partial_data if data? && resuming?
    push_data if data?
    push_indexes if !indexes_first? && schema?
    push_reset_sequences
  end
end

#tables ⇒ Object

# File 'lib/tapsoob/operation.rb', line 637

def tables
  h = {}
  local_tables_info.each do |table_name, count|
    next if completed_tables.include?(table_name.to_s)
    h[table_name.to_s] = count
  end
  h
end

#to_hash ⇒ Object

# File 'lib/tapsoob/operation.rb', line 423

def to_hash
  super.merge(:local_tables_info => local_tables_info)
end