postgresql.rake

require 'capistrano/postgresql/helper_methods'
require 'capistrano/postgresql/password_helpers'
require 'capistrano/postgresql/psql_helpers'
include Capistrano::Postgresql::HelperMethods
include Capistrano::Postgresql::PasswordHelpers
include Capistrano::Postgresql::PsqlHelpers

namespace :load do
  task :defaults do
    # Options necessary for database.yml creation (pg_template|helper_methods.rb)
    set :pg_env, -> { fetch(:rails_env) || fetch(:stage) }
    set :pg_encoding, 'unicode'
    set :pg_database, -> { "#{fetch(:application)}_#{fetch(:stage)}" }
    set :pg_pool, 5
    set :pg_username, -> { fetch(:pg_database) }
    set :pg_generate_random_password, nil
    set :pg_ask_for_password, nil
    set :pg_password, -> { pg_password_generate }
    set :pg_socket, ''
    set :pg_host, -> do # for multiple release nodes automatically use server hostname (IP?) in the database.yml
      release_roles(:all).count == 1 && release_roles(:all).first == primary(:db) ? 'localhost' : primary(:db).hostname
    end
    set :pg_port, 5432
    set :pg_timeout, 5000 # 5 seconds (rails default)
    # General settings
    set :pg_without_sudo, false # issues/22 | Contributed by snake66
    set :pg_system_user, 'postgres'
    set :pg_system_db, 'postgres'
    set :pg_use_hstore, false
    set :pg_extensions, []
    set :pg_templates_path, 'config/deploy/templates'
  end
end

namespace :postgresql do

  desc 'Steps to upgrade the gem to version 4.0'
  task :upgrade4 do
    on roles :db do
      execute :mkdir, '-pv', File.dirname(archetype_database_yml_file)
      execute :cp, database_yml_file, archetype_database_yml_file
    end
  end

  # undocumented, for a reason: drops database. Use with care!
  task :remove_all do
    on release_roles :all do
      execute :rm, database_yml_file if test "[ -e #{database_yml_file} ]"
    end
    on primary :db do
      execute :rm, archetype_database_yml_file if test "[ -e #{archetype_database_yml_file} ]"
    end
    on roles :db do
      psql'execute', fetch(:pg_system_db), "-p #{fetch(:pg_port)} -c", %Q{"DROP database \\"#{fetch(:pg_database)}\\";"} if database_exists?
      psql 'execute', fetch(:pg_system_db),"-p #{fetch(:pg_port)} -c", %Q{"DROP user \\"#{fetch(:pg_username)}\\";"}if database_user_exists?
      remove_extensions
    end
    puts 'Removed database.yml from all hosts, Database, Database User, and Removed Extensions'
  end

  task :remove_app_database_yml_files do
    # We should never delete archetype files. The generate_database_yml_archetype task will handle updates
    on release_roles :app do
        execute :rm, database_yml_file if test "[ -e #{database_yml_file} ]"
    end
  end

  desc 'Remove pg_extension from postgresql db'
  task :remove_extensions do
    remove_extensions
  end

  desc 'Add pg_extension to postgresql db'
  task :add_extensions do
    on roles :db do
      if Array( fetch(:pg_extensions) ).any?
        Array( fetch(:pg_extensions) ).each do |ext|
          next if [nil, false, ''].include?(ext)
          psql 'execute', fetch(:pg_database), "-p #{fetch(:pg_port)} -c", %Q{"CREATE EXTENSION IF NOT EXISTS #{ext};"}unless extension_exists?(ext)
        end
      end
    end
  end

  desc 'Create or update pg_username in database'
  task :create_database_user do
    on roles :db do
      unless database_user_exists?
        # If you use CREATE USER instead of CREATE ROLE the LOGIN right is granted automatically; otherwise you must specify it in the WITH clause of the CREATE statement.
        psql 'execute', fetch(:pg_system_db), "-p #{fetch(:pg_port)} -c", %Q{"CREATE USER \\"#{fetch(:pg_username)}\\" PASSWORD}, redact("'#{fetch(:pg_password)}'"), %Q{;"}
      end
      if database_user_password_different?
        # Ensure updating the password in your deploy/ENV.rb files updates the user, server side
        psql 'execute', fetch(:pg_system_db), "-p #{fetch(:pg_port)} -c", %Q{"ALTER USER \\"#{fetch(:pg_username)}\\" WITH PASSWORD}, redact("'#{fetch(:pg_password)}'"), %Q{;"}
      end
    end
  end

  desc 'Create database'
  task :create_database do
    on roles :db do
      unless database_exists?
        psql 'execute', fetch(:pg_system_db), "-p #{fetch(:pg_port)} -c", %Q{"CREATE DATABASE \\"#{fetch(:pg_database)}\\" OWNER \\"#{fetch(:pg_username)}\\";"}
      end
    end
  end

  # This task creates the archetype database.yml file on the primary db server. This is done once when a new DB user is created.
  desc 'Generate database.yml archetype'
  task :generate_database_yml_archetype do
    on primary :db do
      if test "[ -e #{archetype_database_yml_file} ]" # Archetype already exists. Just update values that changed. Make sure we don't overwrite it to protect generated passwords.
        upload!(
            StringIO.new(pg_template(true, download!(archetype_database_yml_file))),
            archetype_database_yml_file
        )
        # Net::SCP.upload!(
        #     self.host.hostname,
        #     self.host.user,
        #     StringIO.new(pg_template(true, download!(archetype_database_yml_file))),
        #     archetype_database_yml_file,
        #     ssh: { port: self.host.port }
        # )
      else
        execute :mkdir, '-pv', File.dirname(archetype_database_yml_file)
        upload!(
            StringIO.new(pg_template),
            archetype_database_yml_file
        )
        # Net::SCP.upload!(
        #     self.host.hostname,
        #     self.host.user,
        #     StringIO.new(pg_template),
        #     archetype_database_yml_file,
        #     ssh: { port: self.host.port }
        # )
      end
    end
  end

  # This task copies the archetype database file on the primary db server to all clients. This is done on every setup, to ensure new servers get a copy as well.
  desc 'Copy archetype database.yml from primary db server to clients'
  task :generate_database_yml do
    database_yml_contents = nil
    on primary :db do
      database_yml_contents = download! archetype_database_yml_file
    end
    on release_roles :all do
      execute :mkdir, '-pv', File.dirname(database_yml_file)
      upload!(
          StringIO.new(database_yml_contents),
          database_yml_file
      )
      # Net::SCP.upload!(
      #     self.host.hostname,
      #     self.host.user,
      #     StringIO.new(database_yml_contents),
      #     database_yml_file,
      #     ssh: { port: self.host.port }
      # )
    end
  end

  task :database_yml_symlink do
    set :linked_files, fetch(:linked_files, []).push('config/database.yml')
  end

  after 'deploy:started', 'postgresql:database_yml_symlink'

  desc 'Postgresql setup tasks'
  task :setup do
    puts "* ===== Postgresql Setup ===== *\n"
    puts " All psql commands will be run #{fetch(:pg_without_sudo) ? 'without sudo' : 'with sudo'}\n You can modify this in your app/config/deploy/#{fetch(:rails_env)}.rb by setting the pg_without_sudo boolean.\n"
    if release_roles(:app).empty?
      warn " WARNING: There are no servers in your app/config/deploy/#{fetch(:rails_env)}.rb with a :app role... Skipping Postgresql setup."
    else
      if release_roles(:db).empty? # Test to be sure we have a :db role host
        warn " WARNING: There is no server in your app/config/deploy/#{fetch(:rails_env)}.rb with a :db role... Skipping Postgresql setup."
      elsif !fetch(:pg_password) && !fetch(:pg_generate_random_password) && !fetch(:pg_ask_for_password)
        warn " WARNING: There is no :pg_password set in your app/config/deploy/#{fetch(:rails_env)}.rb.\n  If you don't wish to set it, 'set :pg_generate_random_password, true' or 'set :pg_ask_for_password, true' are available!"
      elsif fetch(:pg_generate_random_password) && fetch(:pg_ask_for_password)
        warn " WARNING: You cannot have both :pg_generate_random_password and :pg_ask_for_password enabled in app/config/deploy/#{fetch(:rails_env)}.rb."
      else
        invoke 'postgresql:remove_app_database_yml_files' # Deletes old yml files from all servers. Allows you to avoid having to manually delete the files on your app servers to get a new pool size for example. Don't touch the archetype file to avoid deleting generated passwords.
        invoke 'postgresql:create_database_user'
        invoke 'postgresql:create_database'
        invoke 'postgresql:add_extensions'
        invoke 'postgresql:generate_database_yml_archetype'
        invoke 'postgresql:generate_database_yml'
      end
    end
    puts "* ============================= *"
  end
end

desc 'Server setup tasks'
task :setup do
  invoke 'postgresql:setup'
end