logical_replication_dumpr.rb

logical_replication_dumpr.rb

require 'pg'
require 'timeout'
require 'time'
require 'json'

THE_QUERY = <<SQL
SELECT
    sum(table_size) AS table_size,
    sum(indexes_size) AS indexes_size,
    sum(total_size) AS total_size
FROM (
    SELECT
        table_name,
        pg_table_size(table_name) AS table_size,
        pg_indexes_size(table_name) AS indexes_size,
        pg_total_relation_size(table_name) AS total_size
    FROM (
        SELECT ('"' || table_schema || '"."' || table_name || '"') AS table_name
        FROM information_schema.tables
    ) AS all_tables
    ORDER BY total_size DESC
) AS pretty_sizes;
SQL
    
PROGRESS_QUERY = 'select sr.srsubstate as st, count(*) as c from pg_subscription_rel sr join pg_class cls on (cls.oid = sr.srrelid) group by sr.srsubstate;'

def run_queries(*queries)
  conn_url = ENV['CONN_URL']
  conn = nil
  error = false
  query_results = nil

  begin
    Timeout.timeout 10 do
      conn = PG::Connection.new(conn_url)
    end
    
    query_results = queries.map do |query_sql|
      conn.query(query_sql).to_a
    end

  rescue Timeout::Error
    error = :timeout
  rescue PG::Error, SocketError
    error = :pg_or_socket
  end

  begin
    yield(*query_results) if block_given? && !error
  rescue StandardError => e
    puts "ERROR: #{e}"
  end

  conn&.close rescue nil

  error
end

loop do
  time = Time.now

  error = run_queries(THE_QUERY, PROGRESS_QUERY) do |res, prg|
    data = ["[", Time.now.utc.to_s, "]"]
    %w[table_size indexes_size total_size].each do |name|
      value = res[0][name]
      value_gb = "%.2f" % (res[0][name].to_i / 1073741824.0)
      data << "#{name}=#{value} #{name}_gb=#{value_gb}"
    end

    total = 0
    finished = 0
    prg.each do |row|
      total += row['c'].to_i
      finished += row['c'].to_i if row['st'] == 'r'
    end
    data << ("finished=%.2f%%" % [finished / total.to_f * 100.0])

    File.open('./prd-output.txt', 'a') do |fp|
      fp.puts(data.join(" "))
    end
    puts(data.join(" "))
  end
  puts(error) if error != nil

  wait = (Time.now.to_f - time.to_f)
  sleep(10.0 - wait) if wait < 10.0
end