Add latest changes from gitlab-org/gitlab@16-4-stable-eev16.4.0-rc42

4 files changed, 147 insertions, 65 deletions

diff --git a/lib/backup/database.rb b/lib/backup/database.rb
index 12656cb3702..f70a7e41862 100644
--- a/lib/backup/database.rb
+++ b/lib/backup/database.rb
@@ -27,18 +27,20 @@ module Backup
     def dump(destination_dir, backup_id)
       FileUtils.mkdir_p(destination_dir)
 
-      each_database_snapshot_id do |database_name, snapshot_id|
-        base_model = base_models_for_backup[database_name]
+      each_database(destination_dir) do |database_name, current_db|
+        model = current_db[:model]
+        snapshot_id = current_db[:snapshot_id]
 
-        config = base_model.connection_db_config.configuration_hash
+        pg_env = model.config[:pg_env]
+        connection = model.connection
+        active_record_config = model.config[:activerecord]
+        pg_database = active_record_config[:database]
 
         db_file_name = file_name(destination_dir, database_name)
         FileUtils.rm_f(db_file_name)
 
-        pg_database = config[:database]
-
         progress.print "Dumping PostgreSQL database #{pg_database} ... "
-        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_id}" if snapshot_id
@@ -53,11 +55,13 @@ module Backup
           end
         end
 
-        success = Backup::Dump::Postgres.new.dump(pg_database, db_file_name, pgsql_args)
+        success = with_transient_pg_env(pg_env) do
+          Backup::Dump::Postgres.new.dump(pg_database, db_file_name, pgsql_args)
+        end
 
-        base_model.connection.rollback_transaction if snapshot_id
+        connection.rollback_transaction if snapshot_id
 
-        raise DatabaseBackupError.new(config, db_file_name) unless success
+        raise DatabaseBackupError.new(active_record_config, db_file_name) unless success
 
         report_success(success)
         progress.flush
@@ -72,8 +76,10 @@ module Backup
 
     override :restore
     def restore(destination_dir)
-      base_models_for_backup.each do |database_name, base_model|
-        config = base_model.connection_db_config.configuration_hash
+      base_models_for_backup.each do |database_name, _base_model|
+        backup_model = Backup::DatabaseModel.new(database_name)
+
+        config = backup_model.config[:activerecord]
 
         db_file_name = file_name(destination_dir, database_name)
         database = config[:database]
@@ -94,21 +100,23 @@ module Backup
         # hanging out from a failed upgrade
         drop_tables(database_name)
 
-        decompress_rd, decompress_wr = IO.pipe
-        decompress_pid = spawn(*%w(gzip -cd), out: decompress_wr, in: db_file_name)
-        decompress_wr.close
-
-        status, @errors =
-          case config[:adapter]
-          when "postgresql" then
-            progress.print "Restoring PostgreSQL database #{database} ... "
-            pg_env(config)
-            execute_and_track_errors(pg_restore_cmd(database), decompress_rd)
-          end
-        decompress_rd.close
-
-        Process.waitpid(decompress_pid)
-        success = $?.success? && status.success?
+        pg_env = backup_model.config[:pg_env]
+        success = with_transient_pg_env(pg_env) do
+          decompress_rd, decompress_wr = IO.pipe
+          decompress_pid = spawn(*%w[gzip -cd], out: decompress_wr, in: db_file_name)
+          decompress_wr.close
+
+          status, @errors =
+            case config[:adapter]
+            when "postgresql" then
+              progress.print "Restoring PostgreSQL database #{database} ... "
+              execute_and_track_errors(pg_restore_cmd(database), decompress_rd)
+            end
+          decompress_rd.close
+
+          Process.waitpid(decompress_pid)
+          $?.success? && status.success?
+        end
 
         if @errors.present?
           progress.print "------ BEGIN ERRORS -----\n".color(:yellow)
@@ -204,30 +212,6 @@ module Backup
       end
     end
 
-    def pg_env(config)
-      args = {
-        username: 'PGUSER',
-        host: 'PGHOST',
-        port: 'PGPORT',
-        password: 'PGPASSWORD',
-        # SSL
-        sslmode: 'PGSSLMODE',
-        sslkey: 'PGSSLKEY',
-        sslcert: 'PGSSLCERT',
-        sslrootcert: 'PGSSLROOTCERT',
-        sslcrl: 'PGSSLCRL',
-        sslcompression: 'PGSSLCOMPRESSION'
-      }
-      args.each do |opt, arg|
-        # This enables the use of different PostgreSQL settings in
-        # case PgBouncer is used. PgBouncer clears the search path,
-        # which wreaks havoc on Rails if connections are reused.
-        override = "GITLAB_BACKUP_#{arg}"
-        val = ENV[override].presence || config[opt].to_s.presence
-        ENV[arg] = val if val
-      end
-    end
-
     def report_success(success)
       if success
         progress.puts '[DONE]'.color(:green)
@@ -251,30 +235,45 @@ module Backup
       puts_time 'done'.color(:green)
     end
 
+    def with_transient_pg_env(extended_env)
+      ENV.merge!(extended_env)
+      result = yield
+      ENV.reject! { |k, _| extended_env.key?(k) }
+
+      result
+    end
+
     def pg_restore_cmd(database)
       ['psql', database]
     end
 
-    def each_database_snapshot_id(&block)
-      @database_to_snapshot_id = {}
+    def each_database(destination_dir, &block)
+      databases = {}
+      ::Gitlab::Database::EachDatabase.each_connection(
+        only: base_models_for_backup.keys, include_shared: false
+      ) do |_connection, name|
+        next if databases[name]
+
+        backup_model = Backup::DatabaseModel.new(name)
 
-      if @database_to_snapshot_id.empty?
-        ::Gitlab::Database::EachDatabase.each_connection(
-          only: base_models_for_backup.keys, include_shared: false
-        ) do |connection, database_name|
-          @database_to_snapshot_id[database_name] = nil
+        databases[name] = {
+          model: backup_model
+        }
 
-          next unless Gitlab::Database.database_mode == Gitlab::Database::MODE_MULTIPLE_DATABASES
+        next unless Gitlab::Database.database_mode == Gitlab::Database::MODE_MULTIPLE_DATABASES
 
-          Gitlab::Database::TransactionTimeoutSettings.new(connection).disable_timeouts
+        connection = backup_model.connection
 
+        begin
+          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()")
+          databases[name][:snapshot_id] = connection.select_value("SELECT pg_export_snapshot()")
+        rescue ActiveRecord::ConnectionNotEstablished
+          raise Backup::DatabaseBackupError.new(backup_model.config[:activerecord], file_name(destination_dir, name))
         end
       end
 
-      @database_to_snapshot_id.each(&block)
+      databases.each(&block)
     end
   end
 end
diff --git a/lib/backup/database_model.rb b/lib/backup/database_model.rb
new file mode 100644
index 00000000000..6129a3ce891
--- /dev/null
+++ b/lib/backup/database_model.rb
@@ -0,0 +1,80 @@
+# frozen_string_literal: true
+
+module Backup
+  class DatabaseModel
+    SUPPORTED_OVERRIDES = {
+      username: 'PGUSER',
+      host: 'PGHOST',
+      port: 'PGPORT',
+      password: 'PGPASSWORD',
+      # SSL
+      sslmode: 'PGSSLMODE',
+      sslkey: 'PGSSLKEY',
+      sslcert: 'PGSSLCERT',
+      sslrootcert: 'PGSSLROOTCERT',
+      sslcrl: 'PGSSLCRL',
+      sslcompression: 'PGSSLCOMPRESSION'
+    }.freeze
+
+    attr_reader :config
+
+    def initialize(name)
+      configure_model(name)
+    end
+
+    def connection
+      @model.connection
+    end
+
+    private
+
+    def configure_model(name)
+      source_model = Gitlab::Database.database_base_models_with_gitlab_shared[name]
+
+      @model = backup_model_for(name)
+
+      original_config = source_model.connection_db_config.configuration_hash.dup
+
+      @config = config_for_backup(original_config)
+
+      @model.establish_connection(
+        ActiveRecord::DatabaseConfigurations::HashConfig.new(
+          source_model.connection_db_config.env_name,
+          name.to_s,
+          original_config.merge(@config[:activerecord])
+        )
+      )
+
+      Gitlab::Database::LoadBalancing::Setup.new(@model).setup
+    end
+
+    def backup_model_for(name)
+      klass_name = name.camelize
+
+      return "#{self.class.name}::#{klass_name}".constantize if self.class.const_defined?(klass_name.to_sym, false)
+
+      self.class.const_set(klass_name, Class.new(ApplicationRecord))
+    end
+
+    def config_for_backup(config)
+      db_config = {
+        activerecord: config,
+        pg_env: {}
+      }
+      SUPPORTED_OVERRIDES.each do |opt, arg|
+        # This enables the use of different PostgreSQL settings in
+        # case PgBouncer is used. PgBouncer clears the search path,
+        # which wreaks havoc on Rails if connections are reused.
+        override = "GITLAB_BACKUP_#{arg}"
+        val = ENV[override].presence || config[opt].to_s.presence
+
+        next unless val
+
+        db_config[:pg_env][arg] = val
+        db_config[:activerecord][opt] = val
+      end
+
+      db_config
+    end
+  end
+end
diff --git a/lib/backup/manager.rb b/lib/backup/manager.rb
index 60239781926..2cded4a55bb 100644
--- a/lib/backup/manager.rb
+++ b/lib/backup/manager.rb
@@ -484,7 +484,7 @@ module Backup
 
         puts_time 'Unpacking backup ... '.color(:blue)
 
-        if Kernel.system(*%W(tar -xf #{tar_file}))
+        if Kernel.system(*%W[tar -xf #{tar_file}])
           puts_time 'Unpacking backup ... '.color(:blue) + 'done'.color(:green)
         else
           puts_time 'Unpacking backup failed'.color(:red)
@@ -494,7 +494,7 @@ module Backup
     end
 
     def tar_version
-      tar_version, _ = Gitlab::Popen.popen(%w(tar --version))
+      tar_version, _ = Gitlab::Popen.popen(%w[tar --version])
       tar_version.dup.force_encoding('locale').split("\n").first
     end
 
diff --git a/lib/backup/repositories.rb b/lib/backup/repositories.rb
index 199da8821d9..3b1547148d8 100644
--- a/lib/backup/repositories.rb
+++ b/lib/backup/repositories.rb
@@ -73,7 +73,10 @@ module Backup
     def enqueue_project(project)
       strategy.enqueue(project, Gitlab::GlRepository::PROJECT)
       strategy.enqueue(project, Gitlab::GlRepository::WIKI)
-      strategy.enqueue(project, Gitlab::GlRepository::DESIGN)
+
+      return unless project.design_management_repository
+
+      strategy.enqueue(project.design_management_repository, Gitlab::GlRepository::DESIGN)
     end
 
     def enqueue_snippet(snippet)