Rails rake tasks for dump & restore of PostgreSQL databases

pg_dump.rake

# lib/tasks/pg_dump.rake
# Public fork/gist at: https://gist.github.com/e12e/e0c7d2cc1d30d18c8050b309a43450ac
#
# Original source: https://gist.github.com/hopsoft/56ba6f55fe48ad7f8b90
# Merged with: https://gist.github.com/kofronpi/37130f5ed670465b1fe2d170f754f8c6
#
# This is the code from the comment:
# https://gist.github.com/hopsoft/56ba6f55fe48ad7f8b90#gistcomment-2275324
# Made into a fork/separate gist for easier download, with redundant shell-script removed.
#
# Updates:
# 30072024
# - fixes for rails 7
# - move helptext into rake task (db:dump:help)
# - default to posgresql format if none given (was: sql)
#
# install:
# cd lib/tasks
# wget https://gist.github.com/e12e/e0c7d2cc1d30d18c8050b309a43450ac

namespace :db do
  desc 'Dumps the database to backups'
  task dump: :environment do
    dump_fmt   = ensure_format(ENV.fetch('format', nil))
    dump_sfx   = suffix_for_format(dump_fmt)
    backup_dir = backup_directory(Rails.env, create: true)
    full_path  = nil
    cmd        = nil

    with_config do |app, host, port, db, user|
      full_path = "#{backup_dir}/#{Time.now.strftime('%Y%m%d%H%M%S')}_#{db}.#{dump_sfx}"
      cmd       = build_cmd(dump_fmt:, user:, host:, port:, db:, full_path:)
    end

    puts cmd
    system cmd
    puts ''
    puts "Dumped to file: #{full_path}"
    puts ''
  end

  namespace :dump do
    desc 'Helptext for pg_dump/db:dump and restore tasks'
    task :help do
      puts <<~EOF
        Usage:

        # dump the development db
        rake db:dump

        # dump the db in a specific format
        rake db:dump format=sql

        # dump a table (e.g. users table)
        rake db:dump:table table=users

        # dump a table in a specific format
        rake db:dump:table table=users format=sql

        # list dumps
        rake db:dumps

        # dump the production db
        RAILS_ENV=production rake db:dump

        # restore db based on a backup file pattern (e.g. timestamp)
        rake db:restore pattern=20170101

        # note: config/database.yml is used for database configuration,
        #       but you will be prompted for the database user's password
      EOF
    end

    desc 'Dumps a specific table to backups'
    task table: :environment do
      table_name = ENV.fetch('table', nil)

      if table_name.present?
        dump_fmt   = ensure_format(ENV.fetch('format', nil))
        dump_sfx   = suffix_for_format(dump_fmt)
        backup_dir = backup_directory(Rails.env, create: true)
        full_path  = nil
        cmd        = nil

        with_config do |app, host, port, db, user|
          full_path = "#{backup_dir}/#{Time.now.strftime('%Y%m%d%H%M%S')}_#{db}.#{table_name.parameterize.underscore}.#{dump_sfx}"
          cmd       = build_cmd(dump_fmt:, user:, host:, port:, db:, table_name:, full_path:)
        end

        puts cmd
        system cmd
        puts ''
        puts "Dumped to file: #{full_path}"
        puts ''
      else
        puts 'Please specify a table name'
      end
    end
  end

  desc 'Show the existing database backups'
  task dumps: :environment do
    backup_dir = backup_directory
    puts "#{backup_dir}"
    system "/bin/ls -lt #{backup_dir}"
  end

  desc 'Restores the database from a backup using PATTERN'
  task restore: :environment do
    pattern = ENV.fetch('pattern', nil)

    if pattern.present?
      file = nil
      cmd  = nil

      with_config do |app, host, port, db, user|
        backup_dir = backup_directory
        files      = Dir.glob("#{backup_dir}/**/*#{pattern}*")

        case files.size
        when 0
          puts "No backups found for the pattern '#{pattern}'"
        when 1
          file = files.first
          fmt  = format_for_file file
          host_port_str = build_host_port_params(host:, port:)

          case fmt
          when nil
            puts "No recognized dump file suffix: #{file}"
          when 'p'
            cmd = "psql -U '#{user}' #{host_port_str} -d '#{db}' -f '#{file}'"
          when 'c'
            # Database to create/restore specified in dump - with -C/--create:
            cmd = "pg_restore -F #{fmt} -v -c -C -j 8 -U '#{user}' #{host_port_str} -d 'template1' '#{file}'"
          else
            cmd = "pg_restore -F #{fmt} -v -c -C -U '#{user}' #{host_port_str} -d 'template1' '#{file}'"
          end
        else
          puts "Too many files match the pattern '#{pattern}':"
          puts ' ' + files.join("\n ")
          puts ''
          puts 'Try a more specific pattern'
          puts ''
        end
      end
      unless cmd.nil?
        # Avoid drop/create db, as that requires db creat privilieges, and
        # is redundant with pb_restore -c (clean) and -C (create).
        # Rake::Task["db:drop"].invoke
        # Rake::Task["db:create"].invoke
        puts cmd
        system cmd
        puts ''
        puts "Restored from file: #{file}"
        puts ''
      end
    else
      puts 'Please specify a file pattern for the backup to restore (e.g. timestamp)'
    end
  end

  private

  def ensure_format(format)
    return format if %w[c p t d].include?(format)

    case format
    when 'dump' then 'c'
    when 'sql' then 'p'
    when 'tar' then 't'
    when 'dir' then 'd'
    else 'c'
    end
  end

  def suffix_for_format(suffix)
    case suffix
    when 'c' then 'dump'
    when 'p' then 'sql'
    when 't' then 'tar'
    when 'd' then 'dir'
    else nil
    end
  end

  def format_for_file(file)
    case file
    when /\.dump$/ then 'c'
    when /\.sql$/  then 'p'
    when /\.dir$/  then 'd'
    when /\.tar$/  then 't'
    else nil
    end
  end

  def build_cmd(dump_fmt:, user:, host:, db:, full_path:, port: nil, table_name: nil)
    cmd_str = "pg_dump -F #{dump_fmt} -v -O -U '#{user}'"
    cmd_str << build_host_port_params(host:, port:)
    cmd_str << " -d '#{db}'"
    cmd_str << " -t '#{table_name}'" if table_name
    cmd_str << " -f '#{full_path}'"
    cmd_str
  end

  def build_host_port_params(host:, port: nil)
    conn_sub_str = " -h '#{host}'"
    conn_sub_str << " -p '#{port}'" if port # Allow fallback to socket auth
    conn_sub_str
  end

  def backup_directory(suffix = nil, create: false)
    backup_dir = Rails.root.join('db/backups')

    if create and !Dir.exist?(backup_dir)
      puts "Creating #{backup_dir} .."
      FileUtils.mkdir_p(backup_dir)
    end

    backup_dir
  end

  def with_config
    yield Rails.application.class.module_parent_name.underscore,
          ActiveRecord::Base.connection_db_config.configuration_hash[:host] || 'localhost',
          ActiveRecord::Base.connection_db_config.configuration_hash[:port],
          ActiveRecord::Base.connection_db_config.configuration_hash[:database],
          ActiveRecord::Base.connection_db_config.configuration_hash[:username]
  end
end