1 files changed, 64 insertions, 8 deletions

diff --git a/lib/backup/database.rb b/lib/backup/database.rb
index d4c1ce260e4..851445f703d 100644
--- a/lib/backup/database.rb
+++ b/lib/backup/database.rb
@@ -8,10 +8,18 @@ module Backup
     attr_reader :progress
     attr_reader :config, :db_file_name
 
-    def initialize(progress)
+    IGNORED_ERRORS = [
+      # Ignore the DROP errors; recent database dumps will use --if-exists with pg_dump
+      /does not exist$/,
+      # User may not have permissions to drop extensions or schemas
+      /must be owner of/
+    ].freeze
+    IGNORED_ERRORS_REGEXP = Regexp.union(IGNORED_ERRORS).freeze
+
+    def initialize(progress, filename: nil)
       @progress = progress
       @config = YAML.load_file(File.join(Rails.root, 'config', 'database.yml'))[Rails.env]
-      @db_file_name = File.join(Gitlab.config.backup.path, 'db', 'database.sql.gz')
+      @db_file_name = filename || File.join(Gitlab.config.backup.path, 'db', 'database.sql.gz')
     end
 
     def dump
@@ -27,6 +35,7 @@ module Backup
           progress.print "Dumping PostgreSQL database #{config['database']} ... "
           pg_env
           pgsql_args = ["--clean"] # Pass '--clean' to include 'DROP TABLE' statements in the DB dump.
+          pgsql_args << '--if-exists'
 
           if Gitlab.config.backup.pg_schema
             pgsql_args << '-n'
@@ -48,6 +57,8 @@ module Backup
       end
 
       report_success(success)
+      progress.flush
+
       raise Backup::Error, 'Backup failed' unless success
     end
 
@@ -56,26 +67,65 @@ module Backup
       decompress_pid = spawn(*%w(gzip -cd), out: decompress_wr, in: db_file_name)
       decompress_wr.close
 
-      restore_pid =
+      status, errors =
         case config["adapter"]
         when "postgresql" then
           progress.print "Restoring PostgreSQL database #{config['database']} ... "
           pg_env
-          spawn('psql', config['database'], in: decompress_rd)
+          execute_and_track_errors(pg_restore_cmd, decompress_rd)
         end
       decompress_rd.close
 
-      success = [decompress_pid, restore_pid].all? do |pid|
-        Process.waitpid(pid)
-        $?.success?
+      Process.waitpid(decompress_pid)
+      success = $?.success? && status.success?
+
+      if errors.present?
+        progress.print "------ BEGIN ERRORS -----\n".color(:yellow)
+        progress.print errors.join.color(:yellow)
+        progress.print "------ END ERRORS -------\n".color(:yellow)
       end
 
       report_success(success)
-      abort Backup::Error, 'Restore failed' unless success
+      raise Backup::Error, 'Restore failed' unless success
+
+      errors
     end
 
     protected
 
+    def ignore_error?(line)
+      IGNORED_ERRORS_REGEXP.match?(line)
+    end
+
+    def execute_and_track_errors(cmd, decompress_rd)
+      errors = []
+
+      Open3.popen3(ENV, *cmd) do |stdin, stdout, stderr, thread|
+        stdin.binmode
+
+        out_reader = Thread.new do
+          data = stdout.read
+          $stdout.write(data)
+        end
+
+        err_reader = Thread.new do
+          until (raw_line = stderr.gets).nil?
+            warn(raw_line)
+            errors << raw_line unless ignore_error?(raw_line)
+          end
+        end
+
+        begin
+          IO.copy_stream(decompress_rd, stdin)
+        rescue Errno::EPIPE
+        end
+
+        stdin.close
+        [thread, out_reader, err_reader].each(&:join)
+        [thread.value, errors]
+      end
+    end
+
     def pg_env
       args = {
         'username'  => 'PGUSER',
@@ -100,5 +150,11 @@ module Backup
         progress.puts '[FAILED]'.color(:red)
       end
     end
+
+    private
+
+    def pg_restore_cmd
+      ['psql', config['database']]
+    end
   end
 end