module EY
  module Backup
    class Postgresql < Engine
      register 'postgresql'

      def dump(database_name, basename)
        file = basename + '.dump'

        command = "PGPASSWORD='#{password}' pg_dump -h #{host} --create --format=c -Upostgres #{database_name} 2> /tmp/eybackup.$$.dumperr"

        if gpg?
          command << " | " << GPGEncryptor.command_for(key_id)
          file << GPGEncryptor.extension
        end

        command << " > #{file}"

        run(command, database_name)

        file
      end

      def load(database_name, file)
        if file =~ /.gpz$/ # GPG?
          abort "Cannot restore a GPG backup directly; decrypt the file (#{file}) using your key and then load with pg_restore."
        end

        cycle_database(database_name)

        command = "cat #{file}"

        command << " | PGPASSWORD='#{password}' pg_restore -h #{host} --clean --format=c -Upostgres -d #{database_name} 2> /tmp/eybackup.$$.dumperr"

        run(command, database_name)
      end

      def check_connections(database_name)
        stdout = StringIO.new()
        active_connections = spawn(%Q{PGPASSWORD='#{password}' psql -U postgres -h #{host} -t -c "select count(*) from pg_stat_activity where datname='#{database_name}';"}, stdout)
        
        if stdout.string.to_i > 0
          res = ''
          unless force
            puts "There are currently #{stdout.string.to_i} connections on database: '#{database_name}'; can I kill these to continue (Y/n):"
            Timeout::timeout(30){
              res = gets.strip
            }
          end
          
          if res.upcase == 'Y' or force
            cancel_connections(database_name)
          else
            EY::Backup.logger.fatal(%Q{ERROR: Target database has active connections. For more information, see "Restore or load a database" in docs.engineyard.com})
          end
        end
      end
      
      def cancel_connections(database_name)
        spawn(%Q{psql -U postgres -h #{host} -c"SELECT pg_terminate_backend(pg_stat_activity.pid)
          FROM pg_stat_activity
          WHERE pg_stat_activity.datname = '#{database_name}';"})
      end
      
      def drop_database(database_name)
        spawn("PGPASSWORD='#{password}' dropdb -h #{host} -Upostgres #{database_name}")
      end

      def create_database(database_name)
        spawn("PGPASSWORD='#{password}' createdb -U#{username} -h #{host} #{database_name}")
      end
      
      def check_if_replica
        stdout = StringIO.new()
        spawn(%Q{PGPASSWORD='#{password}' psql -U postgres -h #{host} -t -c "select pg_is_in_recovery()"| head -n 1}, stdout)
        unless stdout.string.chomp =~ /^\W*f$/
          EY::Backup.logger.fatal(%Q{ERROR: Target host: '#{host}' is currently a replica in recovery mode; restore operations need to be processed against the master.})
        end
      end
      
      def create_command(database_name)
        stdout = StringIO.new()
        
        spawn(%Q{PGPASSWORD='#{password}' psql -U postgres -h #{host} -t -c "SELECT 'CREATE DATABASE ' || datname ||
          ' WITH OWNER ' || pg_user.usename ||
          CASE (select pg_encoding_to_char(encoding) from pg_database where datname='template1') 
          WHEN pg_encoding_to_char(encoding) 
            THEN ''
            ELSE ' template=template0'
          END ||
          ' ENCODING ''' || pg_encoding_to_char(encoding) ||
          ''' LC_COLLATE ''' || datcollate||
          ''' LC_CTYPE ''' || datctype || 
          ''' CONNECTION LIMIT ' || datconnlimit || ';'
        FROM pg_database
        INNER JOIN pg_user
          ON pg_user.usesysid = pg_database.datdba
        WHERE datname = '#{database_name}'"
        }, stdout)
        stdout.string
      end
      
      def cycle_database(database_name)
        create_cmd = create_command(database_name).chomp
        if create_cmd == ''
          create_database(database_name)
        else
          check_connections(database_name)
          drop_database(database_name)
          spawn(%Q{PGPASSWORD='#{password}' psql -U postgres -h #{host} -t -c "#{create_cmd}"})
        end
      end

      def suffix
        /\.(dump|gpz)$/
      end
    end
  end
end