1 files changed, 59 insertions, 0 deletions
diff --git a/lib/backup/database_connection.rb b/lib/backup/database_connection.rb
new file mode 100644
index 00000000000..f3f0a5dfcb5
--- /dev/null
+++ b/lib/backup/database_connection.rb
@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+module Backup
+  class DatabaseConnection
+    attr_reader :database_configuration, :snapshot_id
+
+    delegate :connection_name, to: :database_configuration
+    delegate :connection, to: :@backup_model
+
+    # Initializes a database connection
+    #
+    # @param [String] connection_name the key from `database.yml` for multi-database connection configuration
+    def initialize(connection_name)
+      @database_configuration = Backup::DatabaseConfiguration.new(connection_name)
+      @backup_model = backup_model
+      @snapshot_id = nil
+
+      configure_backup_model
+    end
+
+    # Start a new transaction and run pg_export_snapshot()
+    # Returns the snapshot identifier
+    #
+    # @return [String] snapshot identifier
+    def export_snapshot!
+      Gitlab::Database::TransactionTimeoutSettings.new(connection).disable_timeouts
+
+      connection.begin_transaction(isolation: :repeatable_read)
+      @snapshot_id = connection.select_value("SELECT pg_export_snapshot()")
+    end
+
+    # Rollback the transaction to release the effects of pg_export_snapshot()
+    def release_snapshot!
+      return unless snapshot_id
+
+      connection.rollback_transaction
+      @snapshot_id = nil
+    end
+
+    private
+
+    delegate :activerecord_configuration, to: :database_configuration, private: true
+
+    def configure_backup_model
+      @backup_model.establish_connection(activerecord_configuration)
+
+      Gitlab::Database::LoadBalancing::Setup.new(@backup_model).setup
+    end
+
+    # Creates a disposable model to be used to host the Backup connection only
+    def backup_model
+      klass_name = connection_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
+  end
+end