Add latest changes from gitlab-org/gitlab@master

gems/gitlab-backup-cli/lib/gitlab/backup/cli/errors.rb

@@ -5,6 +5,7 @@ module Gitlab
     module Cli
       module Errors
         autoload :DatabaseBackupError, 'gitlab/backup/cli/errors/database_backup_error'
+        autoload :DatabaseCleanupError, 'gitlab/backup/cli/errors/database_cleanup_error'
         autoload :DatabaseConfigMissingError, 'gitlab/backup/cli/errors/database_config_missing_error'
         autoload :DatabaseMissingConnectionError, 'gitlab/backup/cli/errors/database_missing_connection_error'
         autoload :FileBackupError, 'gitlab/backup/cli/errors/file_backup_error'

gems/gitlab-backup-cli/lib/gitlab/backup/cli/errors/database_cleanup_error.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Backup
+    module Cli
+      module Errors
+        class DatabaseCleanupError < StandardError
+          attr_reader :task, :path, :error
+
+          def initialize(task:, path:, error:)
+            @task = task
+            @path = path
+            @error = error
+
+            super(build_message)
+          end
+
+          private
+
+          def build_message
+            "Failed to cleanup GitLab databases \n" \
+              "Running the following rake task: '#{task}' (from: #{path}) failed:\n" \
+              "#{error}"
+          end
+        end
+      end
+    end
+  end
+end

gems/gitlab-backup-cli/lib/gitlab/backup/cli/targets/database.rb

@@ -18,6 +18,10 @@ module Gitlab
           ].freeze
           IGNORED_ERRORS_REGEXP = Regexp.union(IGNORED_ERRORS).freeze
 
+          # Rake task used to drop all tables from GitLab databases
+          # This task is executed before restoring data
+          DROP_TABLES_TASK = "gitlab:db:drop_tables"
+
           attr_reader :errors
 
           def initialize(context)
@@ -66,6 +70,10 @@ module Gitlab
           def restore(source)
             databases = Gitlab::Backup::Cli::Services::Postgres.new(context)
 
+            # Drop all tables Load the schema to ensure we don't have any newer tables
+            # hanging out from a failed upgrade
+            drop_tables!
+
             databases.each do |db|
               database_name = db.configuration.name
               pg_database_name = db.configuration.database
@@ -89,10 +97,6 @@ module Gitlab
                 next
               end
 
-              # Drop all tables Load the schema to ensure we don't have any newer tables
-              # hanging out from a failed upgrade
-              drop_tables(db)
-
               Gitlab::Backup::Cli::Output.info "Restoring PostgreSQL database #{pg_database_name} ... "
 
               status = restore_tables(database: db, filepath: db_file_name)
@@ -151,18 +155,22 @@ module Gitlab
             Gitlab::Backup::Cli::Output.print_tag(status ? :success : :failure)
           end
 
-          def drop_tables(database)
-            pg_database_name = database.configuration.database
-            Gitlab::Backup::Cli::Output.print_info "Cleaning the '#{pg_database_name}' database ... "
+          def drop_tables!
+            Gitlab::Backup::Cli::Output.print_info "Cleaning existing databases ... "
 
-            if Rake::Task.task_defined? "gitlab:db:drop_tables:#{database.configuration.name}"
-              Rake::Task["gitlab:db:drop_tables:#{database.configuration.name}"].invoke
-            else
-              # In single database (single or two connections)
-              Rake::Task["gitlab:db:drop_tables"].invoke
+            gitlab_path = context.gitlab_basepath
+
+            # Drop existing tables from configured databases before restoring from a backup
+            rake = Utils::Rake.new(DROP_TABLES_TASK, chdir: gitlab_path).execute
+
+            unless rake.success?
+              Gitlab::Backup::Cli::Output.print_tag(:failure)
+
+              raise Errors::DatabaseCleanupError.new(task: DROP_TABLES_TASK, path: gitlab_path, error: rake.stderr)
             end
 
             Gitlab::Backup::Cli::Output.print_tag(:success)
+            Gitlab::Backup::Cli::Output.info(rake.output) unless rake.output.empty?
           end
 
           def restore_tables(database:, filepath:)

gems/gitlab-backup-cli/lib/gitlab/backup/cli/utils.rb

@@ -6,6 +6,7 @@ module Gitlab
       module Utils
         autoload :Compression, 'gitlab/backup/cli/utils/compression'
         autoload :PgDump, 'gitlab/backup/cli/utils/pg_dump'
+        autoload :Rake, 'gitlab/backup/cli/utils/rake'
         autoload :Tar, 'gitlab/backup/cli/utils/tar'
       end
     end

gems/gitlab-backup-cli/lib/gitlab/backup/cli/utils/rake.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Backup
+    module Cli
+      module Utils
+        class Rake
+          # @return [Array<String>] a list of tasks to be executed
+          attr_reader :tasks
+
+          # @return [String|Pathname] a path where rake tasks are run from
+          attr_reader :chdir
+
+          # @param [Array<String>] *tasks a list of tasks to be executed
+          # @param [String|Pathname] chdir a path where rake tasks are run from
+          def initialize(*tasks, chdir: Gitlab::Backup::Cli.root)
+            @tasks = tasks
+            @chdir = chdir
+          end
+
+          # @return [self]
+          def execute
+            Bundler.with_original_env do
+              @result = Shell::Command.new(*rake_command, chdir: chdir).capture
+            end
+
+            self
+          end
+
+          # Return whether the execution was a success or not
+          #
+          # @return [Boolean] whether the execution was a success
+          def success?
+            @result&.status&.success? || false
+          end
+
+          # Return the captured rake output
+          #
+          # @return [String] stdout content
+          def output
+            @result&.stdout || ''
+          end
+
+          # Return the captured error content
+          #
+          # @return [String] stdout content
+          def stderr
+            @result&.stderr || ''
+          end
+
+          # Return the captured execution duration
+          #
+          # @return [Float] execution duration
+          def duration
+            @result&.duration || 0.0
+          end
+
+          private
+
+          # Return a list of commands necessary to execute `rake`
+          #
+          # @return [Array<String (frozen)>] array of commands to be used by Shellout
+          def rake_command
+            %w[bundle exec rake] + tasks
+          end
+        end
+      end
+    end
+  end
+end