Add latest changes from gitlab-org/gitlab@16-0-stable-eev16.0.0-rc42

4 files changed, 71 insertions, 28 deletions

diff --git a/lib/backup/database.rb b/lib/backup/database.rb
index bd3832c7327..28bc78a3932 100644
--- a/lib/backup/database.rb
+++ b/lib/backup/database.rb
@@ -27,7 +27,7 @@ module Backup
     def dump(destination_dir, backup_id)
       FileUtils.mkdir_p(destination_dir)
 
-      snapshot_ids.each do |database_name, snapshot_id|
+      each_database_snapshot_id do |database_name, snapshot_id|
         base_model = base_models_for_backup[database_name]
 
         config = base_model.connection_db_config.configuration_hash
@@ -41,7 +41,7 @@ module Backup
         pg_env(config)
         pgsql_args = ["--clean"] # Pass '--clean' to include 'DROP TABLE' statements in the DB dump.
         pgsql_args << '--if-exists'
-        pgsql_args << "--snapshot=#{snapshot_ids[database_name]}"
+        pgsql_args << "--snapshot=#{snapshot_id}" if snapshot_id
 
         if Gitlab.config.backup.pg_schema
           pgsql_args << '-n'
@@ -55,7 +55,7 @@ module Backup
 
         success = Backup::Dump::Postgres.new.dump(pg_database, db_file_name, pgsql_args)
 
-        base_model.connection.rollback_transaction
+        base_model.connection.rollback_transaction if snapshot_id
 
         raise DatabaseBackupError.new(config, db_file_name) unless success
 
@@ -63,8 +63,10 @@ module Backup
         progress.flush
       end
     ensure
-      base_models_for_backup.each do |_database_name, base_model|
-        Gitlab::Database::TransactionTimeoutSettings.new(base_model.connection).restore_timeouts
+      ::Gitlab::Database::EachDatabase.each_database_connection(
+        only: base_models_for_backup.keys, include_shared: false
+      ) do |connection, _|
+        Gitlab::Database::TransactionTimeoutSettings.new(connection).restore_timeouts
       end
     end
 
@@ -237,31 +239,42 @@ module Backup
     private
 
     def drop_tables(database_name)
+      puts_time 'Cleaning the database ... '.color(:blue)
+
       if Rake::Task.task_defined? "gitlab:db:drop_tables:#{database_name}"
-        puts_time 'Cleaning the database ... '.color(:blue)
         Rake::Task["gitlab:db:drop_tables:#{database_name}"].invoke
-        puts_time 'done'.color(:green)
-      elsif Gitlab::Database.database_base_models.one?
-        # In single database, we do not have rake tasks per database
-        puts_time 'Cleaning the database ... '.color(:blue)
+      else
+        # In single database (single or two connections)
         Rake::Task["gitlab:db:drop_tables"].invoke
-        puts_time 'done'.color(:green)
       end
+
+      puts_time 'done'.color(:green)
     end
 
     def pg_restore_cmd(database)
       ['psql', database]
     end
 
-    def snapshot_ids
-      @snapshot_ids ||= base_models_for_backup.each_with_object({}) do |(database_name, base_model), snapshot_ids|
-        Gitlab::Database::TransactionTimeoutSettings.new(base_model.connection).disable_timeouts
+    def each_database_snapshot_id(&block)
+      @database_to_snapshot_id = {}
+
+      if @database_to_snapshot_id.empty?
+        ::Gitlab::Database::EachDatabase.each_database_connection(
+          only: base_models_for_backup.keys, include_shared: false
+        ) do |connection, database_name|
+          @database_to_snapshot_id[database_name] = nil
 
-        base_model.connection.begin_transaction(isolation: :repeatable_read)
+          next unless Gitlab::Database.database_mode == Gitlab::Database::MODE_MULTIPLE_DATABASES
 
-        snapshot_ids[database_name] =
-          base_model.connection.execute("SELECT pg_export_snapshot() as snapshot_id;").first['snapshot_id']
+          Gitlab::Database::TransactionTimeoutSettings.new(connection).disable_timeouts
+
+          connection.begin_transaction(isolation: :repeatable_read)
+
+          @database_to_snapshot_id[database_name] = connection.select_value("SELECT pg_export_snapshot()")
+        end
       end
+
+      @database_to_snapshot_id.each(&block)
     end
   end
 end
diff --git a/lib/backup/gitaly_backup.rb b/lib/backup/gitaly_backup.rb
index 57dd74c7950..53c998efd71 100644
--- a/lib/backup/gitaly_backup.rb
+++ b/lib/backup/gitaly_backup.rb
@@ -9,14 +9,15 @@ module Backup
     # @param [StringIO] progress IO interface to output progress
     # @param [Integer] max_parallelism max parallelism when running backups
     # @param [Integer] storage_parallelism max parallelism per storage (is affected by max_parallelism)
-    def initialize(progress, max_parallelism: nil, storage_parallelism: nil, incremental: false, backup_id: nil)
+    # @param [Boolean] incremental if incremental backups should be created.
+    def initialize(progress, max_parallelism: nil, storage_parallelism: nil, incremental: false)
       @progress = progress
       @max_parallelism = max_parallelism
       @storage_parallelism = storage_parallelism
       @incremental = incremental
     end
 
-    def start(type, backup_repos_path, backup_id: nil)
+    def start(type, backup_repos_path, backup_id: nil, remove_all_repositories: nil)
       raise Error, 'already started' if started?
 
       if type == :create && !incremental?
@@ -35,9 +36,13 @@ module Backup
       args = ['-layout', 'pointer']
       args += ['-parallel', @max_parallelism.to_s] if @max_parallelism
       args += ['-parallel-storage', @storage_parallelism.to_s] if @storage_parallelism
-      if type == :create
+
+      case type
+      when :create
         args += ['-incremental'] if incremental?
         args += ['-id', backup_id] if backup_id
+      when :restore
+        args += ['-remove-all-repositories', remove_all_repositories.join(',')] if remove_all_repositories
       end
 
       @input_stream, stdout, @thread = Open3.popen2(build_env, bin_path, command, '-path', backup_repos_path, *args)
@@ -77,11 +82,7 @@ module Backup
     #
     # @see https://gitlab.com/gitlab-org/gitaly/-/blob/master/doc/gitaly-backup.md
     def schedule_backup_job(repository, always_create:)
-      connection_params = Gitlab::GitalyClient.connection_data(repository.storage)
-
       json_job = {
-        address: connection_params['address'],
-        token: connection_params['token'],
         storage_name: repository.storage,
         relative_path: repository.relative_path,
         gl_project_path: repository.gl_project_path,
@@ -91,10 +92,21 @@ module Backup
       @input_stream.puts(json_job)
     end
 
+    def gitaly_servers
+      Gitlab.config.repositories.storages.keys.index_with do |storage_name|
+        Gitlab::GitalyClient.connection_data(storage_name)
+      end
+    end
+
+    def gitaly_servers_encoded
+      Base64.strict_encode64(Gitlab::Json.dump(gitaly_servers))
+    end
+
     def build_env
       {
         'SSL_CERT_FILE' => Gitlab::X509::Certificate.default_cert_file,
-        'SSL_CERT_DIR' => Gitlab::X509::Certificate.default_cert_dir
+        'SSL_CERT_DIR' => Gitlab::X509::Certificate.default_cert_dir,
+        'GITALY_SERVERS' => gitaly_servers_encoded
       }.merge(ENV)
     end
 
diff --git a/lib/backup/manager.rb b/lib/backup/manager.rb
index ba4a26ba714..b5e1634004a 100644
--- a/lib/backup/manager.rb
+++ b/lib/backup/manager.rb
@@ -15,6 +15,10 @@ module Backup
       repositories_paths: 'REPOSITORIES_PATHS'
     }.freeze
 
+    YAML_PERMITTED_CLASSES = [
+      ActiveSupport::TimeWithZone, ActiveSupport::TimeZone, Symbol, Time
+    ].freeze
+
     TaskDefinition = Struct.new(
       :enabled, # `true` if the task can be used. Treated as `true` when not specified.
       :human_name, # Name of the task used for logging.
@@ -247,7 +251,9 @@ module Backup
     end
 
     def read_backup_information
-      @backup_information ||= YAML.load_file(File.join(backup_path, MANIFEST_NAME))
+      @backup_information ||= YAML.safe_load_file(
+        File.join(backup_path, MANIFEST_NAME),
+        permitted_classes: YAML_PERMITTED_CLASSES)
     end
 
     def write_backup_information
@@ -416,6 +422,12 @@ module Backup
       end
     end
 
+    def puts_available_timestamps
+      available_timestamps.each do |available_timestamp|
+        puts_time " " + available_timestamp
+      end
+    end
+
     def unpack(source_backup_id)
       if source_backup_id.blank? && non_tarred_backup?
         puts_time "Non tarred backup found in #{backup_path}, using that"
@@ -431,7 +443,7 @@ module Backup
         elsif backup_file_list.many? && source_backup_id.nil?
           puts_time 'Found more than one backup:'
           # print list of available backups
-          puts_time " " + available_timestamps.join("\n ")
+          puts_available_timestamps
 
           if incremental?
             puts_time 'Please specify which one you want to create an incremental backup for:'
diff --git a/lib/backup/repositories.rb b/lib/backup/repositories.rb
index 4f4a098f374..218df3fcb6c 100644
--- a/lib/backup/repositories.rb
+++ b/lib/backup/repositories.rb
@@ -30,7 +30,7 @@ module Backup
 
     override :restore
     def restore(destination_path)
-      strategy.start(:restore, destination_path)
+      strategy.start(:restore, destination_path, remove_all_repositories: remove_all_repositories)
       enqueue_consecutive
 
     ensure
@@ -44,6 +44,12 @@ module Backup
 
     attr_reader :strategy, :storages, :paths
 
+    def remove_all_repositories
+      return if paths.present?
+
+      storages.presence || Gitlab.config.repositories.storages.keys
+    end
+
     def enqueue_consecutive
       enqueue_consecutive_projects
       enqueue_consecutive_snippets