atlas_recipe :postgresql do

during :provision, %w(
  create_user
  create_database
  database_yml
  pgpass
  logrotate_backup
)

end

namespace :atlas do

namespace :postgresql do
  desc "Create user if it doesn't already exist"
  task :create_user do
    privileged_on primary(:db) do
      user = fetch(:atlas_postgresql_user)

      unless test("sudo -Hiu postgres psql -c '\\du' | grep -q #{user}")
        passwd = fetch(:atlas_postgresql_password)
        md5 = Digest::MD5.hexdigest(passwd + user)
        execute "sudo -Hiu postgres psql -c " +
                %Q["CREATE USER #{user} PASSWORD 'md5#{md5}';"]
      end
    end
  end

  desc "Create database if it doesn't already exist"
  task :create_database do
    privileged_on primary(:db) do
      user = fetch(:atlas_postgresql_user)
      db = fetch(:atlas_postgresql_database)

      unless test("sudo -Hiu postgres psql -l | grep -w -q #{db}")
        execute "sudo -Hiu postgres createdb -O #{user} #{db}"
      end
    end
  end

  desc "Generate database.yml"
  task :database_yml do
    yaml = {
      fetch(:rails_env).to_s => {
        "adapter" => "postgresql",
        "encoding" => "unicode",
        "database" => fetch(:atlas_postgresql_database).to_s,
        "pool" => fetch(:atlas_postgresql_pool_size).to_i,
        "username" => fetch(:atlas_postgresql_user).to_s,
        "password" => fetch(:atlas_postgresql_password).to_s,
        "host" => fetch(:atlas_postgresql_host).to_s
      }
    }
    fetch(:atlas_postgresql_password)
    on release_roles(:all) do
      put YAML.dump(yaml),
          "#{shared_path}/config/database.yml",
          :mode => "600"
    end
  end

  desc "Generate pgpass file (needed by backup scripts)"
  task :pgpass do
    fetch(:atlas_postgresql_password)
    on release_roles(:all) do
      template "pgpass.erb",
               fetch(:atlas_postgresql_pgpass_path),
               :mode => "600"
    end
  end

  desc "Configure logrotate to back up the database daily"
  task :logrotate_backup do
    on roles(:backup) do
      backup_path = fetch(:atlas_postgresql_backup_path)
      execute :mkdir, "-p", File.dirname(backup_path)
      execute :touch, backup_path
    end

    privileged_on roles(:backup) do |host, user|
      template\
        "postgresql-backup-logrotate.erb",
        "/etc/logrotate.d/postgresql-backup-#{application_basename}",
        :owner => "root:root",
        :mode => "644",
        :binding => binding,
        :sudo => true
    end
  end

  desc "Dump the database to FILE"
  task :dump do
    on primary(:db) do
      with_pgpassfile do
        execute :pg_dump,
          "-Fc -Z9 -O",
          "-x", fetch(:atlas_postgresql_dump_options),
          "-f", remote_dump_file,
          connection_flags,
          fetch(:atlas_postgresql_database)
      end

      download!(remote_dump_file, local_dump_file)

      info(
        "Exported #{fetch(:atlas_postgresql_database)} "\
        "to #{local_dump_file}."
        )
    end
  end

  desc "Restore database from FILE"
  task :restore do
    on primary(:db) do
      exit 1 unless agree(
        "\nErase existing #{fetch(:rails_env)} database "\
        "and restore from local file: #{local_dump_file}? "
        )

      upload!(local_dump_file, remote_dump_file)

      with_pgpassfile do
        execute :pg_restore,
          "-O -c",
          connection_flags,
          "-d", fetch(:atlas_postgresql_database),
          remote_dump_file
      end
    end
  end

  def local_dump_file
    ENV.fetch("FILE", "#{fetch(:atlas_postgresql_database)}.dmp")
  end

  def remote_dump_file
    "/tmp/#{fetch(:atlas_postgresql_database)}.dmp"
  end

  def connection_flags
    [
      "-U", fetch(:atlas_postgresql_user),
      "-h", fetch(:atlas_postgresql_host)
    ].join(" ")
  end

  def with_pgpassfile(&block)
    with(:pgpassfile => fetch(:atlas_postgresql_pgpass_path), &block)
  end
end

end