Quickly copy a postgres table

copy_table.rb

require 'pg'

# Copy a table, e.g. when altering a column to a wider numeric type.
# Minimal usage when db server is localhost:
# SRC_TABLE=table DST_TABLE=table_modified DBNAME=mydb ruby copy_table.rb
class CopyTable
  BUFFER_CAPACITY = 2000

  def initialize(src_table:, dst_table:, dbname:, host:, user:, password:)
    @src_table = src_table
    @dst_table = dst_table
    @dbname = dbname
    @host = host
    @user = user
    @password = password
  end

  def copy!
    dst_conn.copy_data("COPY #{@dst_table} FROM STDIN") do
      src_conn.copy_data("COPY #{@src_table} TO STDOUT") do
        buffered_copy
      end
    end
  end

  # Instead of copying each row into the remote database as soon as we have
  # it, buffer the copy rows and flush periodically for better performance.
  # NB: maintainting an array of strings and joining on flush rather than
  # repeatedly concatenating to a string is the approach recommended by
  # https://bugs.ruby-lang.org/issues/12024#note-16 since we can't
  # preallocate big strings.
  def buffered_copy
    buffer = []
    count = 0
    while row = src_conn.get_copy_data
      buffer << row
      if buffer.size >= BUFFER_CAPACITY
        count += BUFFER_CAPACITY
        print "\r#{count}"
        dst_conn.put_copy_data(buffer.join)
        buffer = []
      end
    end
    dst_conn.put_copy_data(buffer.join) if buffer.any?
  end

  def dst_conn
    @dst_conn ||= PG::Connection.open(conn_details)
  end

  def src_conn
    @src_conn ||= PG::Connection.open(conn_details)
  end

  def conn_details
    {
      dbname: @dbname,
      host: @host,
      user: @user,
      password: @password
    }
  end
end

if __FILE__ == $0
  CopyTable.new(
    src_table: ENV.fetch('SRC_TABLE'),
    dst_table: ENV.fetch('DST_TABLE'),
    dbname: ENV.fetch('DBNAME'),
    host: ENV['DB_HOST'],
    user: ENV['DB_USER'],
    password: ENV['DB_PASSWORD']
  ).copy!
end