Refactor Gitlab::Database to support multiple DBs

This refactors the module Gitlab::Database so it supports multiple
databases. Logic tied to a connection is now scoped to instances of
Gitlab::Dababase::Connection. For every database, an instance of this
class is created and stored in Gitlab::Database::DATABASES.

The module Gitlab::Database still exposes various methods that forward
the call to a corresponding instance of Database::Connection. These
wrappers and their use will be changed in separate commits, as otherwise
the diffs become far too large for anybody to make sense of.

See https://gitlab.com/gitlab-org/gitlab/-/issues/331776 for more
information.

config/initializers/database_config.rb

@@ -23,8 +23,9 @@ end
 db_config = Gitlab::Database.config ||
             Rails.application.config.database_configuration[Rails.env]
 
-db_config['pool'] = Gitlab::Database.default_pool_size
-ActiveRecord::Base.establish_connection(db_config)
+ActiveRecord::Base.establish_connection(
+  db_config.merge(pool: Gitlab::Database.default_pool_size)
+)
 
 Gitlab.ee do
   if Gitlab::Runtime.sidekiq? && Gitlab::Geo.geo_database_configured?

ee/app/workers/geo/prune_event_log_worker.rb

@@ -18,7 +18,7 @@ module Geo
 
     def perform
       return if Gitlab::Database.read_only?
-      return unless Gitlab::Database.healthy?
+      return unless Gitlab::Database.main.healthy?
 
       unless ::GeoNode.secondary_nodes.any?
         Geo::PruneEventLogService.new(:all).execute

ee/app/workers/geo/scheduler/scheduler_worker.rb

@@ -182,7 +182,7 @@ module Geo
 
       def update_pending_resources
         if reload_queue?
-          @pending_resources = Gitlab::Database.geo_uncached_queries { load_pending_resources }
+          @pending_resources = Gitlab::Database.main.geo_uncached_queries { load_pending_resources }
           set_backoff_time! if should_apply_backoff?
         end
       end

ee/lib/ee/gitlab/database.rb → ee/lib/ee/gitlab/database/connection.rb

@@ -3,9 +3,8 @@
 module EE
   module Gitlab
     module Database
+      module Connection
         extend ActiveSupport::Concern
-
-      class_methods do
         extend ::Gitlab::Utils::Override
 
         override :read_only?
@@ -20,7 +19,7 @@ module EE
         def geo_uncached_queries(&block)
           raise 'No block given' unless block_given?
 
-          ActiveRecord::Base.uncached do
+          scope.uncached do
             if ::Gitlab::Geo.secondary?
               Geo::TrackingBase.uncached(&block)
             else

ee/spec/lib/ee/gitlab/database_spec.rb → ee/spec/lib/ee/gitlab/database/connection_spec.rb

@@ -2,10 +2,12 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database do
+RSpec.describe Gitlab::Database::Connection do
   include ::EE::GeoHelpers
 
-  describe '.read_only?' do
+  let(:connection) { described_class.new }
+
+  describe '#read_only?' do
     context 'with Geo enabled' do
       before do
         allow(Gitlab::Geo).to receive(:enabled?) { true }
@@ -16,7 +18,7 @@ RSpec.describe Gitlab::Database do
         let(:geo_node) { create(:geo_node) }
 
         it 'returns true' do
-          expect(described_class.read_only?).to be_truthy
+          expect(connection.read_only?).to be_truthy
         end
       end
 
@@ -24,14 +26,14 @@ RSpec.describe Gitlab::Database do
         let(:geo_node) { create(:geo_node, :primary) }
 
         it 'returns false when is Geo primary node' do
-          expect(described_class.read_only?).to be_falsey
+          expect(connection.read_only?).to be_falsey
         end
       end
     end
 
     context 'with Geo disabled' do
       it 'returns false' do
-        expect(described_class.read_only?).to be_falsey
+        expect(connection.read_only?).to be_falsey
       end
     end
 
@@ -41,30 +43,30 @@ RSpec.describe Gitlab::Database do
       end
 
       it 'returns true' do
-        expect(described_class.read_only?).to be_truthy
+        expect(connection.read_only?).to be_truthy
       end
     end
   end
 
-  describe '.healthy?' do
+  describe '#healthy?' do
     it 'returns true when replication lag is not too great' do
       allow(Postgresql::ReplicationSlot).to receive(:lag_too_great?).and_return(false)
 
-      expect(described_class.healthy?).to be_truthy
+      expect(connection.healthy?).to be_truthy
     end
 
     it 'returns false when replication lag is too great' do
       allow(Postgresql::ReplicationSlot).to receive(:lag_too_great?).and_return(true)
 
-      expect(described_class.healthy?).to be_falsey
+      expect(connection.healthy?).to be_falsey
     end
   end
 
-  describe '.geo_uncached_queries' do
+  describe '#geo_uncached_queries' do
     context 'when no block is given' do
       it 'raises error' do
         expect do
-          described_class.geo_uncached_queries
+          connection.geo_uncached_queries
         end.to raise_error('No block given')
       end
     end
@@ -79,7 +81,7 @@ RSpec.describe Gitlab::Database do
         expect(ActiveRecord::Base).to receive(:uncached).and_call_original
 
         expect do |b|
-          described_class.geo_uncached_queries(&b)
+          connection.geo_uncached_queries(&b)
         end.to yield_control
       end
     end
@@ -95,7 +97,7 @@ RSpec.describe Gitlab::Database do
         expect(ActiveRecord::Base).to receive(:uncached).and_call_original
 
         expect do |b|
-          described_class.geo_uncached_queries(&b)
+          connection.geo_uncached_queries(&b)
         end.to yield_control
       end
     end
@@ -106,7 +108,7 @@ RSpec.describe Gitlab::Database do
         expect(ActiveRecord::Base).to receive(:uncached).and_call_original
 
         expect do |b|
-          described_class.geo_uncached_queries(&b)
+          connection.geo_uncached_queries(&b)
         end.to yield_control
       end
     end

ee/spec/models/ee/group_spec.rb

@@ -840,7 +840,7 @@ RSpec.describe Group do
 
       context 'in read-only mode' do
         before do
-          allow(Gitlab::Database).to receive(:read_only?).and_return(true)
+          allow(Gitlab::Database.main).to receive(:read_only?).and_return(true)
           allow(group).to receive(:create_or_update).and_raise(ActiveRecord::ReadOnlyRecord)
         end
 

ee/spec/models/ee/user_spec.rb

@@ -287,7 +287,7 @@ RSpec.describe User do
     end
 
     it 'does not clear remember_created_at when in a GitLab read-only instance' do
-      allow(Gitlab::Database).to receive(:read_only?) { true }
+      allow(Gitlab::Database.main).to receive(:read_only?) { true }
 
       expect { subject.forget_me! }.not_to change(subject, :remember_created_at)
     end
@@ -303,7 +303,7 @@ RSpec.describe User do
     end
 
     it 'does not update remember_created_at when in a Geo read-only instance' do
-      allow(Gitlab::Database).to receive(:read_only?) { true }
+      allow(Gitlab::Database.main).to receive(:read_only?) { true }
 
       expect { subject.remember_me! }.not_to change(subject, :remember_created_at)
     end

ee/spec/services/audit_event_service_spec.rb

@@ -336,7 +336,7 @@ RSpec.describe AuditEventService, :request_store do
 
     context 'on a read-only instance' do
       before do
-        allow(Gitlab::Database).to receive(:read_only?).and_return(true)
+        allow(Gitlab::Database.main).to receive(:read_only?).and_return(true)
       end
 
       it 'does not create an event record in the database' do

ee/spec/workers/geo/prune_event_log_worker_spec.rb

@@ -29,7 +29,7 @@ RSpec.describe Geo::PruneEventLogWorker, :geo do
       end
 
       it 'does nothing when database is not feeling healthy' do
-        allow(EE::Gitlab::Database).to receive(:healthy?).and_return(false)
+        allow(Gitlab::Database.main).to receive(:healthy?).and_return(false)
 
         expect(Geo::PruneEventLogService).not_to receive(:new)
 

lib/gitlab/database.rb

@@ -45,27 +45,18 @@ module Gitlab
     # It does not include the default public schema
     EXTRA_SCHEMAS = [DYNAMIC_PARTITIONS_SCHEMA, STATIC_PARTITIONS_SCHEMA].freeze
 
-    DEFAULT_POOL_HEADROOM = 10
-
-    # We configure the database connection pool size automatically based on the
-    # configured concurrency. We also add some headroom, to make sure we don't run
-    # out of connections when more threads besides the 'user-facing' ones are
-    # running.
-    #
-    # Read more about this in doc/development/database/client_side_connection_pool.md
-    def self.default_pool_size
-      headroom = (ENV["DB_POOL_HEADROOM"].presence || DEFAULT_POOL_HEADROOM).to_i
-
-      Gitlab::Runtime.max_threads + headroom
+    DATABASES = ActiveRecord::Base
+      .connection_handler
+      .connection_pools
+      .each_with_object({}) do |pool, hash|
+        hash[pool.db_config.name.to_sym] = Connection.new(pool.connection_klass)
       end
+      .freeze
 
-    def self.config
-      default_config_hash = ActiveRecord::Base.configurations.find_db_config(Rails.env)&.configuration_hash || {}
+    PRIMARY_DATABASE_NAME = ActiveRecord::Base.connection_db_config.name.to_sym
 
-      default_config_hash.with_indifferent_access.tap do |hash|
-        # Match config/initializers/database_config.rb
-        hash[:pool] ||= default_pool_size
-      end
+    def self.main
+      DATABASES[PRIMARY_DATABASE_NAME]
     end
 
     def self.has_config?(database_name)
@@ -87,72 +78,68 @@ module Gitlab
       name.to_s == CI_DATABASE_NAME
     end
 
+    def self.default_pool_size
+      main.default_pool_size
+    end
+
+    def self.config
+      main.config
+    end
+
     def self.username
-      config['username'] || ENV['USER']
+      main.username
     end
 
     def self.database_name
-      config['database']
+      main.database_name
     end
 
     def self.adapter_name
-      config['adapter']
+      main.adapter_name
     end
 
     def self.human_adapter_name
-      if postgresql?
-        'PostgreSQL'
-      else
-        'Unknown'
-      end
+      main.human_adapter_name
     end
 
-    # Disables prepared statements for the current database connection.
     def self.disable_prepared_statements
-      ActiveRecord::Base.establish_connection(config.merge(prepared_statements: false))
+      main.disable_prepared_statements
     end
 
-    # @deprecated
     def self.postgresql?
-      adapter_name.casecmp('postgresql') == 0
+      main.postgresql?
     end
 
     def self.read_only?
-      false
+      main.read_only?
     end
 
     def self.read_write?
-      !self.read_only?
+      main.read_write?
     end
 
-    # Check whether the underlying database is in read-only mode
     def self.db_read_only?
-      pg_is_in_recovery =
-        ActiveRecord::Base
-          .connection
-          .execute('SELECT pg_is_in_recovery()')
-          .first
-          .fetch('pg_is_in_recovery')
-
-      Gitlab::Utils.to_boolean(pg_is_in_recovery)
+      main.db_read_only?
     end
 
     def self.db_read_write?
-      !self.db_read_only?
+      main.db_read_write?
     end
 
     def self.version
-      @version ||= database_version.match(/\A(?:PostgreSQL |)([^\s]+).*\z/)[1]
+      main.version
     end
 
     def self.postgresql_minimum_supported_version?
-      version.to_f >= MINIMUM_POSTGRES_VERSION
+      main.postgresql_minimum_supported_version?
     end
 
     def self.check_postgres_version_and_print_warning
-      return if Gitlab::Database.postgresql_minimum_supported_version?
       return if Gitlab::Runtime.rails_runner?
 
+      DATABASES.each do |name, connection|
+        next if connection.postgresql_minimum_supported_version?
+
         Kernel.warn ERB.new(Rainbow.new.wrap(<<~EOS).red).result
 
                     ██     ██  █████  ██████  ███    ██ ██ ███    ██  ██████ 
@@ -162,7 +149,7 @@ module Gitlab
                      ███ ███  ██   ██ ██   ██ ██   ████ ██ ██   ████  ██████  
 
           ******************************************************************************
-          You are using PostgreSQL <%= Gitlab::Database.version %>, but PostgreSQL >= <%= Gitlab::Database::MINIMUM_POSTGRES_VERSION %>
+            You are using PostgreSQL <%= Gitlab::Database.version %> for the #{name} database, but PostgreSQL >= <%= Gitlab::Database::MINIMUM_POSTGRES_VERSION %>
             is required for this version of GitLab.
             <% if Rails.env.development? || Rails.env.test? %>
             If using gitlab-development-kit, please find the relevant steps here:
@@ -175,6 +162,7 @@ module Gitlab
       rescue ActiveRecord::ActiveRecordError, PG::Error
         # ignore - happens when Rake tasks yet have to create a database, e.g. for testing
       end
+    end
 
     def self.nulls_order(field, direction = :asc, nulls_order = :nulls_last)
       raise ArgumentError unless [:nulls_last, :nulls_first].include?(nulls_order)
@@ -206,132 +194,49 @@ module Gitlab
       "'f'"
     end
 
-    def self.with_connection_pool(pool_size)
-      pool = create_connection_pool(pool_size)
-
-      begin
-        yield(pool)
-      ensure
-        pool.disconnect!
-      end
+    def self.with_connection_pool(...)
+      main.with_connection_pool(...)
     end
 
-    # Bulk inserts a number of rows into a table, optionally returning their
-    # IDs.
-    #
-    # table - The name of the table to insert the rows into.
-    # rows - An Array of Hash instances, each mapping the columns to their
-    #        values.
-    # return_ids - When set to true the return value will be an Array of IDs of
-    #              the inserted rows
-    # disable_quote - A key or an Array of keys to exclude from quoting (You
-    #                 become responsible for protection from SQL injection for
-    #                 these keys!)
-    # on_conflict - Defines an upsert. Values can be: :disabled (default) or
-    #               :do_nothing
-    def self.bulk_insert(table, rows, return_ids: false, disable_quote: [], on_conflict: nil)
-      return if rows.empty?
-
-      keys = rows.first.keys
-      columns = keys.map { |key| connection.quote_column_name(key) }
-
-      disable_quote = Array(disable_quote).to_set
-      tuples = rows.map do |row|
-        keys.map do |k|
-          disable_quote.include?(k) ? row[k] : connection.quote(row[k])
-        end
-      end
-
-      sql = <<-EOF
-        INSERT INTO #{table} (#{columns.join(', ')})
-        VALUES #{tuples.map { |tuple| "(#{tuple.join(', ')})" }.join(', ')}
-      EOF
-
-      sql = "#{sql} ON CONFLICT DO NOTHING" if on_conflict == :do_nothing
-
-      sql = "#{sql} RETURNING id" if return_ids
-
-      result = connection.execute(sql)
-
-      if return_ids
-        result.values.map { |tuple| tuple[0].to_i }
-      else
-        []
-      end
+    def self.bulk_insert(...)
+      main.bulk_insert(...)
     end
 
     def self.sanitize_timestamp(timestamp)
       MAX_TIMESTAMP_VALUE > timestamp ? timestamp : MAX_TIMESTAMP_VALUE.dup
     end
 
-    # pool_size - The size of the DB pool.
-    # host - An optional host name to use instead of the default one.
-    def self.create_connection_pool(pool_size, host = nil, port = nil)
-      original_config = Gitlab::Database.config
-
-      env_config = original_config.merge(pool: pool_size)
-      env_config[:host] = host if host
-      env_config[:port] = port if port
-
-      ActiveRecord::ConnectionAdapters::ConnectionHandler.new.establish_connection(env_config)
+    def self.create_connection_pool(...)
+      main.create_connection_pool(...)
     end
 
     def self.connection
-      ActiveRecord::Base.connection
+      main.connection
     end
     private_class_method :connection
 
-    def self.cached_column_exists?(table_name, column_name)
-      connection.schema_cache.columns_hash(table_name).has_key?(column_name.to_s)
+    def self.cached_column_exists?(...)
+      main.cached_column_exists?(...)
     end
 
-    def self.cached_table_exists?(table_name)
-      exists? && connection.schema_cache.data_source_exists?(table_name)
+    def self.cached_table_exists?(...)
+      main.cached_table_exists?(...)
     end
 
     def self.database_version
-      row = connection.execute("SELECT VERSION()").first
-
-      row['version']
+      main.database_version
     end
 
     def self.exists?
-      connection
-
-      true
-    rescue StandardError
-      false
+      main.exists?
     end
 
     def self.system_id
-      row = connection.execute('SELECT system_identifier FROM pg_control_system()').first
-
-      row['system_identifier']
-    end
-
-    # @param [ActiveRecord::Connection] ar_connection
-    # @return [String]
-    def self.get_write_location(ar_connection)
-      use_new_load_balancer_query = Gitlab::Utils.to_boolean(ENV['USE_NEW_LOAD_BALANCER_QUERY'], default: true)
-
-      sql = if use_new_load_balancer_query
-              <<~NEWSQL
-                SELECT CASE
-                    WHEN pg_is_in_recovery() = true AND EXISTS (SELECT 1 FROM pg_stat_get_wal_senders())
-                      THEN pg_last_wal_replay_lsn()::text
-                    WHEN pg_is_in_recovery() = false
-                      THEN pg_current_wal_insert_lsn()::text
-                      ELSE NULL
-                    END AS location;
-              NEWSQL
-            else
-              <<~SQL
-                SELECT pg_current_wal_insert_lsn()::text AS location
-              SQL
+      main.system_id
     end
 
-      row = ar_connection.select_all(sql).first
-      row['location'] if row
+    def self.get_write_location(...)
+      main.get_write_location(...)
     end
 
     private_class_method :database_version
@@ -362,32 +267,18 @@ module Gitlab
       'unknown'
     end
 
-    # inside_transaction? will return true if the caller is running within a transaction. Handles special cases
-    # when running inside a test environment, where tests may be wrapped in transactions
     def self.inside_transaction?
-      if Rails.env.test?
-        ActiveRecord::Base.connection.open_transactions > open_transactions_baseline
-      else
-        ActiveRecord::Base.connection.open_transactions > 0
-      end
+      main.inside_transaction?
     end
 
-    # These methods that access @open_transactions_baseline are not thread-safe.
-    # These are fine though because we only call these in RSpec's main thread. If we decide to run
-    # specs multi-threaded, we would need to use something like ThreadGroup to keep track of this value
     def self.set_open_transactions_baseline
-      @open_transactions_baseline = ActiveRecord::Base.connection.open_transactions
+      main.set_open_transactions_baseline
     end
 
     def self.reset_open_transactions_baseline
-      @open_transactions_baseline = 0
+      main.reset_open_transactions_baseline
     end
 
-    def self.open_transactions_baseline
-      @open_transactions_baseline ||= 0
-    end
-    private_class_method :open_transactions_baseline
-
     # Monkeypatch rails with upgraded database observability
     def self.install_monkey_patches
       ActiveRecord::Base.prepend(ActiveRecordBaseTransactionMetrics)
@@ -409,5 +300,3 @@ module Gitlab
     end
   end
 end
-
-Gitlab::Database.prepend_mod_with('Gitlab::Database')

lib/gitlab/database/connection.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    # Configuration settings and methods for interacting with a PostgreSQL
+    # database, with support for multiple databases.
+    class Connection
+      DEFAULT_POOL_HEADROOM = 10
+
+      attr_reader :scope
+
+      # Initializes a new `Database`.
+      #
+      # The `scope` argument must be an object (such as `ActiveRecord::Base`)
+      # that supports retrieving connections and connection pools.
+      def initialize(scope = ActiveRecord::Base)
+        @config = nil
+        @scope = scope
+        @version = nil
+        @open_transactions_baseline = 0
+      end
+
+      # We configure the database connection pool size automatically based on
+      # the configured concurrency. We also add some headroom, to make sure we
+      # don't run out of connections when more threads besides the 'user-facing'
+      # ones are running.
+      #
+      # Read more about this in
+      # doc/development/database/client_side_connection_pool.md
+      def default_pool_size
+        headroom =
+          (ENV["DB_POOL_HEADROOM"].presence || DEFAULT_POOL_HEADROOM).to_i
+
+        Gitlab::Runtime.max_threads + headroom
+      end
+
+      def config
+        @config ||=
+          scope.connection_db_config.configuration_hash.with_indifferent_access
+      end
+
+      def pool_size
+        config[:pool] || default_pool_size
+      end
+
+      def username
+        config[:username] || ENV['USER']
+      end
+
+      def database_name
+        config[:database]
+      end
+
+      def adapter_name
+        config[:adapter]
+      end
+
+      def human_adapter_name
+        if postgresql?
+          'PostgreSQL'
+        else
+          'Unknown'
+        end
+      end
+
+      def postgresql?
+        adapter_name.casecmp('postgresql') == 0
+      end
+
+      # Disables prepared statements for the current database connection.
+      def disable_prepared_statements
+        scope.establish_connection(config.merge(prepared_statements: false))
+      end
+
+      def read_only?
+        false
+      end
+
+      def read_write?
+        !read_only?
+      end
+
+      # Check whether the underlying database is in read-only mode
+      def db_read_only?
+        pg_is_in_recovery =
+          scope
+            .connection
+            .execute('SELECT pg_is_in_recovery()')
+            .first
+            .fetch('pg_is_in_recovery')
+
+        Gitlab::Utils.to_boolean(pg_is_in_recovery)
+      end
+
+      def db_read_write?
+        !db_read_only?
+      end
+
+      def version
+        @version ||= database_version.match(/\A(?:PostgreSQL |)([^\s]+).*\z/)[1]
+      end
+
+      def database_version
+        connection.execute("SELECT VERSION()").first['version']
+      end
+
+      def postgresql_minimum_supported_version?
+        version.to_f >= MINIMUM_POSTGRES_VERSION
+      end
+
+      # Bulk inserts a number of rows into a table, optionally returning their
+      # IDs.
+      #
+      # table - The name of the table to insert the rows into.
+      # rows - An Array of Hash instances, each mapping the columns to their
+      #        values.
+      # return_ids - When set to true the return value will be an Array of IDs of
+      #              the inserted rows
+      # disable_quote - A key or an Array of keys to exclude from quoting (You
+      #                 become responsible for protection from SQL injection for
+      #                 these keys!)
+      # on_conflict - Defines an upsert. Values can be: :disabled (default) or
+      #               :do_nothing
+      def bulk_insert(table, rows, return_ids: false, disable_quote: [], on_conflict: nil)
+        return if rows.empty?
+
+        keys = rows.first.keys
+        columns = keys.map { |key| connection.quote_column_name(key) }
+
+        disable_quote = Array(disable_quote).to_set
+        tuples = rows.map do |row|
+          keys.map do |k|
+            disable_quote.include?(k) ? row[k] : connection.quote(row[k])
+          end
+        end
+
+        sql = <<-EOF
+          INSERT INTO #{table} (#{columns.join(', ')})
+          VALUES #{tuples.map { |tuple| "(#{tuple.join(', ')})" }.join(', ')}
+        EOF
+
+        sql = "#{sql} ON CONFLICT DO NOTHING" if on_conflict == :do_nothing
+
+        sql = "#{sql} RETURNING id" if return_ids
+
+        result = connection.execute(sql)
+
+        if return_ids
+          result.values.map { |tuple| tuple[0].to_i }
+        else
+          []
+        end
+      end
+
+      # pool_size - The size of the DB pool.
+      # host - An optional host name to use instead of the default one.
+      # port - An optional port to connect to.
+      def create_connection_pool(pool_size, host = nil, port = nil)
+        original_config = config
+        env_config = original_config.merge(pool: pool_size)
+
+        env_config[:host] = host if host
+        env_config[:port] = port if port
+
+        ActiveRecord::ConnectionAdapters::ConnectionHandler
+          .new.establish_connection(env_config)
+      end
+
+      def with_connection_pool(pool_size)
+        pool = create_connection_pool(pool_size)
+
+        begin
+          yield(pool)
+        ensure
+          pool.disconnect!
+        end
+      end
+
+      def cached_column_exists?(table_name, column_name)
+        connection
+          .schema_cache.columns_hash(table_name)
+          .has_key?(column_name.to_s)
+      end
+
+      def cached_table_exists?(table_name)
+        exists? && connection.schema_cache.data_source_exists?(table_name)
+      end
+
+      def exists?
+        connection
+
+        true
+      rescue StandardError
+        false
+      end
+
+      def system_id
+        row = connection
+          .execute('SELECT system_identifier FROM pg_control_system()')
+          .first
+
+        row['system_identifier']
+      end
+
+      # @param [ActiveRecord::Connection] ar_connection
+      # @return [String]
+      def get_write_location(ar_connection)
+        use_new_load_balancer_query = Gitlab::Utils
+          .to_boolean(ENV['USE_NEW_LOAD_BALANCER_QUERY'], default: true)
+
+        sql =
+          if use_new_load_balancer_query
+            <<~NEWSQL
+              SELECT CASE
+                  WHEN pg_is_in_recovery() = true AND EXISTS (SELECT 1 FROM pg_stat_get_wal_senders())
+                    THEN pg_last_wal_replay_lsn()::text
+                  WHEN pg_is_in_recovery() = false
+                    THEN pg_current_wal_insert_lsn()::text
+                    ELSE NULL
+                  END AS location;
+            NEWSQL
+          else
+            <<~SQL
+              SELECT pg_current_wal_insert_lsn()::text AS location
+            SQL
+          end
+
+        row = ar_connection.select_all(sql).first
+        row['location'] if row
+      end
+
+      # inside_transaction? will return true if the caller is running within a
+      # transaction. Handles special cases when running inside a test
+      # environment, where tests may be wrapped in transactions
+      def inside_transaction?
+        base = Rails.env.test? ? @open_transactions_baseline : 0
+
+        scope.connection.open_transactions > base
+      end
+
+      # These methods that access @open_transactions_baseline are not
+      # thread-safe.  These are fine though because we only call these in
+      # RSpec's main thread. If we decide to run specs multi-threaded, we would
+      # need to use something like ThreadGroup to keep track of this value
+      def set_open_transactions_baseline
+        @open_transactions_baseline = scope.connection.open_transactions
+      end
+
+      def reset_open_transactions_baseline
+        @open_transactions_baseline = 0
+      end
+
+      def connection
+        scope.connection
+      end
+    end
+  end
+end
+
+Gitlab::Database::Connection.prepend_mod_with('Gitlab::Database::Connection')

lib/gitlab/database/load_balancing.rb

@@ -79,7 +79,7 @@ module Gitlab
       end
 
       def self.pool_size
-        Gitlab::Database.config[:pool]
+        Gitlab::Database.main.pool_size
       end
 
       # Returns true if load balancing is to be enabled.

scripts/rspec_helpers.sh

@@ -84,10 +84,6 @@ function rspec_simple_job() {
 function rspec_db_library_code() {
   local db_files="spec/lib/gitlab/database/ spec/support/helpers/database/"
 
-  if [[ -d "ee/" ]]; then
-    db_files="${db_files} ee/spec/lib/ee/gitlab/database_spec.rb"
-  fi
-
   rspec_simple_job "-- ${db_files}"
 }
 

spec/initializers/database_config_spec.rb

@@ -21,37 +21,23 @@ RSpec.describe 'Database config initializer' do
 
   let(:max_threads) { 8 }
 
-  context "no existing pool size is set" do
-    before do
-      stub_database_config(pool_size: nil)
-    end
-
-    it "sets it based on the max number of worker threads" do
-      expect { subject }.to change { Gitlab::Database.config['pool'] }.from(nil).to(18)
-
-      expect(ActiveRecord::Base.connection_db_config.pool).to eq(18)
-    end
-  end
-
-  context "the existing pool size is smaller than the max number of worker threads" do
-    before do
-      stub_database_config(pool_size: 1)
-    end
+  context 'when no custom headroom is specified' do
+    it 'sets the pool size based on the number of worker threads' do
+      old = ActiveRecord::Base.connection_db_config.pool
 
-    it "sets it based on the max number of worker threads" do
-      expect { subject }.to change { Gitlab::Database.config['pool'] }.from(1).to(18)
+      expect(old).not_to eq(18)
 
-      expect(ActiveRecord::Base.connection_db_config.pool).to eq(18)
-    end
+      expect { subject }
+        .to change { ActiveRecord::Base.connection_db_config.pool }
+        .from(old)
+        .to(18)
     end
 
-  context "and the existing pool size is larger than the max number of worker threads" do
-    before do
-      stub_database_config(pool_size: 100)
-    end
+    it 'overwrites custom pool settings' do
+      config = Gitlab::Database.config.merge(pool: 42)
 
-    it "sets it based on the max number of worker threads" do
-      expect { subject }.to change { Gitlab::Database.config['pool'] }.from(100).to(18)
+      allow(Gitlab::Database.main).to receive(:config).and_return(config)
+      subject
 
       expect(ActiveRecord::Base.connection_db_config.pool).to eq(18)
     end
@@ -61,25 +47,16 @@ RSpec.describe 'Database config initializer' do
     let(:headroom) { 15 }
 
     before do
-      stub_database_config(pool_size: 1)
       stub_env("DB_POOL_HEADROOM", headroom)
     end
 
     it "adds headroom on top of the calculated size" do
-      expect { subject }.to change { Gitlab::Database.config['pool'] }
-                              .from(1)
-                              .to(max_threads + headroom)
+      old = ActiveRecord::Base.connection_db_config.pool
 
-      expect(ActiveRecord::Base.connection_db_config.pool).to eq(max_threads + headroom)
+      expect { subject }
+        .to change { ActiveRecord::Base.connection_db_config.pool }
+        .from(old)
+        .to(23)
     end
   end
-
-  def stub_database_config(pool_size:)
-    original_config = Gitlab::Database.config
-
-    config = original_config.dup
-    config['pool'] = pool_size
-
-    allow(Gitlab::Database).to receive(:config).and_return(config)
-  end
 end

spec/lib/gitlab/auth_spec.rb

@@ -844,7 +844,7 @@ RSpec.describe Gitlab::Auth, :use_clean_rails_memory_store_caching do
 
       context 'when the database is read-only' do
         before do
-          allow(Gitlab::Database).to receive(:read_only?).and_return(true)
+          allow(Gitlab::Database.main).to receive(:read_only?).and_return(true)
         end
 
         it 'does not increment failed_attempts when true and password is incorrect' do

spec/lib/gitlab/database/connection_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Connection do
+  let(:connection) { described_class.new }
+
+  describe '#default_pool_size' do
+    before do
+      allow(Gitlab::Runtime).to receive(:max_threads).and_return(7)
+    end
+
+    it 'returns the max thread size plus a fixed headroom of 10' do
+      expect(connection.default_pool_size).to eq(17)
+    end
+
+    it 'returns the max thread size plus a DB_POOL_HEADROOM if this env var is present' do
+      stub_env('DB_POOL_HEADROOM', '7')
+
+      expect(connection.default_pool_size).to eq(14)
+    end
+  end
+
+  describe '#config' do
+    it 'returns a HashWithIndifferentAccess' do
+      expect(connection.config).to be_an_instance_of(HashWithIndifferentAccess)
+    end
+
+    it 'returns a default pool size' do
+      expect(connection.config).to include(pool: connection.default_pool_size)
+    end
+  end
+
+  describe '#pool_size' do
+    context 'when no explicit size is configured' do
+      it 'returns the default pool size' do
+        expect(connection.config).to receive(:[]).with(:pool).and_return(nil)
+
+        expect(connection.pool_size).to eq(connection.default_pool_size)
+      end
+    end
+
+    context 'when an explicit pool size is set' do
+      it 'returns the pool size' do
+        expect(connection.config).to receive(:[]).with(:pool).and_return(4)
+
+        expect(connection.pool_size).to eq(4)
+      end
+    end
+  end
+
+  describe '#username' do
+    context 'when a username is set' do
+      it 'returns the username' do
+        allow(connection).to receive(:config).and_return(username: 'bob')
+
+        expect(connection.username).to eq('bob')
+      end
+    end
+
+    context 'when a username is not set' do
+      it 'returns the value of the USER environment variable' do
+        allow(connection).to receive(:config).and_return(username: nil)
+        allow(ENV).to receive(:[]).with('USER').and_return('bob')
+
+        expect(connection.username).to eq('bob')
+      end
+    end
+  end
+
+  describe '#database_name' do
+    it 'returns the name of the database' do
+      allow(connection).to receive(:config).and_return(database: 'test')
+
+      expect(connection.database_name).to eq('test')
+    end
+  end
+
+  describe '#adapter_name' do
+    it 'returns the database adapter name' do
+      allow(connection).to receive(:config).and_return(adapter: 'test')
+
+      expect(connection.adapter_name).to eq('test')
+    end
+  end
+
+  describe '#human_adapter_name' do
+    context 'when the adapter is PostgreSQL' do
+      it 'returns PostgreSQL' do
+        allow(connection).to receive(:config).and_return(adapter: 'postgresql')
+
+        expect(connection.human_adapter_name).to eq('PostgreSQL')
+      end
+    end
+
+    context 'when the adapter is not PostgreSQL' do
+      it 'returns Unknown' do
+        allow(connection).to receive(:config).and_return(adapter: 'kittens')
+
+        expect(connection.human_adapter_name).to eq('Unknown')
+      end
+    end
+  end
+
+  describe '#postgresql?' do
+    context 'when using PostgreSQL' do
+      it 'returns true' do
+        allow(connection).to receive(:adapter_name).and_return('PostgreSQL')
+
+        expect(connection.postgresql?).to eq(true)
+      end
+    end
+
+    context 'when not using PostgreSQL' do
+      it 'returns false' do
+        allow(connection).to receive(:adapter_name).and_return('MySQL')
+
+        expect(connection.postgresql?).to eq(false)
+      end
+    end
+  end
+
+  describe '#disable_prepared_statements' do
+    around do |example|
+      original_config = ::Gitlab::Database.config
+
+      example.run
+
+      connection.scope.establish_connection(original_config)
+    end
+
+    it 'disables prepared statements' do
+      connection.scope.establish_connection(
+        ::Gitlab::Database.config.merge(prepared_statements: true)
+      )
+
+      expect(connection.scope.connection.prepared_statements).to eq(true)
+
+      connection.disable_prepared_statements
+
+      expect(connection.scope.connection.prepared_statements).to eq(false)
+    end
+  end
+
+  describe '#read_only?' do
+    it 'returns false' do
+      expect(connection.read_only?).to eq(false)
+    end
+  end
+
+  describe '#read_write' do
+    it 'returns true' do
+      expect(connection.read_write?).to eq(true)
+    end
+  end
+
+  describe '#db_read_only?' do
+    it 'detects a read-only database' do
+      allow(connection.scope.connection)
+        .to receive(:execute)
+        .with('SELECT pg_is_in_recovery()')
+        .and_return([{ "pg_is_in_recovery" => "t" }])
+
+      expect(connection.db_read_only?).to be_truthy
+    end
+
+    it 'detects a read-only database' do
+      allow(connection.scope.connection)
+        .to receive(:execute)
+        .with('SELECT pg_is_in_recovery()')
+        .and_return([{ "pg_is_in_recovery" => true }])
+
+      expect(connection.db_read_only?).to be_truthy
+    end
+
+    it 'detects a read-write database' do
+      allow(connection.scope.connection)
+        .to receive(:execute)
+        .with('SELECT pg_is_in_recovery()')
+        .and_return([{ "pg_is_in_recovery" => "f" }])
+
+      expect(connection.db_read_only?).to be_falsey
+    end
+
+    it 'detects a read-write database' do
+      allow(connection.scope.connection)
+        .to receive(:execute)
+        .with('SELECT pg_is_in_recovery()')
+        .and_return([{ "pg_is_in_recovery" => false }])
+
+      expect(connection.db_read_only?).to be_falsey
+    end
+  end
+
+  describe '#db_read_write?' do
+    it 'detects a read-only database' do
+      allow(connection.scope.connection)
+        .to receive(:execute)
+        .with('SELECT pg_is_in_recovery()')
+        .and_return([{ "pg_is_in_recovery" => "t" }])
+
+      expect(connection.db_read_write?).to eq(false)
+    end
+
+    it 'detects a read-only database' do
+      allow(connection.scope.connection)
+        .to receive(:execute)
+        .with('SELECT pg_is_in_recovery()')
+        .and_return([{ "pg_is_in_recovery" => true }])
+
+      expect(connection.db_read_write?).to eq(false)
+    end
+
+    it 'detects a read-write database' do
+      allow(connection.scope.connection)
+        .to receive(:execute)
+        .with('SELECT pg_is_in_recovery()')
+        .and_return([{ "pg_is_in_recovery" => "f" }])
+
+      expect(connection.db_read_write?).to eq(true)
+    end
+
+    it 'detects a read-write database' do
+      allow(connection.scope.connection)
+        .to receive(:execute)
+        .with('SELECT pg_is_in_recovery()')
+        .and_return([{ "pg_is_in_recovery" => false }])
+
+      expect(connection.db_read_write?).to eq(true)
+    end
+  end
+
+  describe '#version' do
+    around do |example|
+      connection.instance_variable_set(:@version, nil)
+      example.run
+      connection.instance_variable_set(:@version, nil)
+    end
+
+    context "on postgresql" do
+      it "extracts the version number" do
+        allow(connection)
+          .to receive(:database_version)
+          .and_return("PostgreSQL 9.4.4 on x86_64-apple-darwin14.3.0")
+
+        expect(connection.version).to eq '9.4.4'
+      end
+    end
+
+    it 'memoizes the result' do
+      count = ActiveRecord::QueryRecorder
+        .new { 2.times { connection.version } }
+        .count
+
+      expect(count).to eq(1)
+    end
+  end
+
+  describe '#postgresql_minimum_supported_version?' do
+    it 'returns false when using PostgreSQL 10' do
+      allow(connection).to receive(:version).and_return('10')
+
+      expect(connection.postgresql_minimum_supported_version?).to eq(false)
+    end
+
+    it 'returns false when using PostgreSQL 11' do
+      allow(connection).to receive(:version).and_return('11')
+
+      expect(connection.postgresql_minimum_supported_version?).to eq(false)
+    end
+
+    it 'returns true when using PostgreSQL 12' do
+      allow(connection).to receive(:version).and_return('12')
+
+      expect(connection.postgresql_minimum_supported_version?).to eq(true)
+    end
+  end
+
+  describe '#bulk_insert' do
+    before do
+      allow(connection).to receive(:connection).and_return(dummy_connection)
+      allow(dummy_connection).to receive(:quote_column_name, &:itself)
+      allow(dummy_connection).to receive(:quote, &:itself)
+      allow(dummy_connection).to receive(:execute)
+    end
+
+    let(:dummy_connection) { double(:connection) }
+
+    let(:rows) do
+      [
+        { a: 1, b: 2, c: 3 },
+        { c: 6, a: 4, b: 5 }
+      ]
+    end
+
+    it 'does nothing with empty rows' do
+      expect(dummy_connection).not_to receive(:execute)
+
+      connection.bulk_insert('test', [])
+    end
+
+    it 'uses the ordering from the first row' do
+      expect(dummy_connection).to receive(:execute) do |sql|
+        expect(sql).to include('(1, 2, 3)')
+        expect(sql).to include('(4, 5, 6)')
+      end
+
+      connection.bulk_insert('test', rows)
+    end
+
+    it 'quotes column names' do
+      expect(dummy_connection).to receive(:quote_column_name).with(:a)
+      expect(dummy_connection).to receive(:quote_column_name).with(:b)
+      expect(dummy_connection).to receive(:quote_column_name).with(:c)
+
+      connection.bulk_insert('test', rows)
+    end
+
+    it 'quotes values' do
+      1.upto(6) do |i|
+        expect(dummy_connection).to receive(:quote).with(i)
+      end
+
+      connection.bulk_insert('test', rows)
+    end
+
+    it 'does not quote values of a column in the disable_quote option' do
+      [1, 2, 4, 5].each do |i|
+        expect(dummy_connection).to receive(:quote).with(i)
+      end
+
+      connection.bulk_insert('test', rows, disable_quote: :c)
+    end
+
+    it 'does not quote values of columns in the disable_quote option' do
+      [2, 5].each do |i|
+        expect(dummy_connection).to receive(:quote).with(i)
+      end
+
+      connection.bulk_insert('test', rows, disable_quote: [:a, :c])
+    end
+
+    it 'handles non-UTF-8 data' do
+      expect { connection.bulk_insert('test', [{ a: "\255" }]) }.not_to raise_error
+    end
+
+    context 'when using PostgreSQL' do
+      it 'allows the returning of the IDs of the inserted rows' do
+        result = double(:result, values: [['10']])
+
+        expect(dummy_connection)
+          .to receive(:execute)
+          .with(/RETURNING id/)
+          .and_return(result)
+
+        ids = connection
+          .bulk_insert('test', [{ number: 10 }], return_ids: true)
+
+        expect(ids).to eq([10])
+      end
+
+      it 'allows setting the upsert to do nothing' do
+        expect(dummy_connection)
+          .to receive(:execute)
+          .with(/ON CONFLICT DO NOTHING/)
+
+        connection
+          .bulk_insert('test', [{ number: 10 }], on_conflict: :do_nothing)
+      end
+    end
+  end
+
+  describe '#create_connection_pool' do
+    it 'creates a new connection pool with specific pool size' do
+      pool = connection.create_connection_pool(5)
+
+      begin
+        expect(pool)
+          .to be_kind_of(ActiveRecord::ConnectionAdapters::ConnectionPool)
+
+        expect(pool.db_config.pool).to eq(5)
+      ensure
+        pool.disconnect!
+      end
+    end
+
+    it 'allows setting of a custom hostname' do
+      pool = connection.create_connection_pool(5, '127.0.0.1')
+
+      begin
+        expect(pool.db_config.host).to eq('127.0.0.1')
+      ensure
+        pool.disconnect!
+      end
+    end
+
+    it 'allows setting of a custom hostname and port' do
+      pool = connection.create_connection_pool(5, '127.0.0.1', 5432)
+
+      begin
+        expect(pool.db_config.host).to eq('127.0.0.1')
+        expect(pool.db_config.configuration_hash[:port]).to eq(5432)
+      ensure
+        pool.disconnect!
+      end
+    end
+  end
+
+  describe '#with_connection_pool' do
+    it 'creates a new connection pool and disconnect it after used' do
+      closed_pool = nil
+
+      connection.with_connection_pool(1) do |pool|
+        pool.with_connection do |connection|
+          connection.execute('SELECT 1 AS value')
+        end
+
+        expect(pool).to be_connected
+
+        closed_pool = pool
+      end
+
+      expect(closed_pool).not_to be_connected
+    end
+
+    it 'disconnects the pool even an exception was raised' do
+      error = Class.new(RuntimeError)
+      closed_pool = nil
+
+      begin
+        connection.with_connection_pool(1) do |pool|
+          pool.with_connection do |connection|
+            connection.execute('SELECT 1 AS value')
+          end
+
+          closed_pool = pool
+
+          raise error, 'boom'
+        end
+      rescue error
+      end
+
+      expect(closed_pool).not_to be_connected
+    end
+  end
+
+  describe '#cached_column_exists?' do
+    it 'only retrieves data once' do
+      expect(connection.scope.connection)
+        .to receive(:columns)
+        .once.and_call_original
+
+      2.times do
+        expect(connection.cached_column_exists?(:projects, :id)).to be_truthy
+        expect(connection.cached_column_exists?(:projects, :bogus_column)).to be_falsey
+      end
+    end
+  end
+
+  describe '#cached_table_exists?' do
+    it 'only retrieves data once per table' do
+      expect(connection.scope.connection)
+        .to receive(:data_source_exists?)
+        .with(:projects)
+        .once.and_call_original
+
+      expect(connection.scope.connection)
+        .to receive(:data_source_exists?)
+        .with(:bogus_table_name)
+        .once.and_call_original
+
+      2.times do
+        expect(connection.cached_table_exists?(:projects)).to be_truthy
+        expect(connection.cached_table_exists?(:bogus_table_name)).to be_falsey
+      end
+    end
+
+    it 'returns false when database does not exist' do
+      expect(connection.scope).to receive(:connection) do
+        raise ActiveRecord::NoDatabaseError, 'broken'
+      end
+
+      expect(connection.cached_table_exists?(:projects)).to be(false)
+    end
+  end
+
+  describe '#exists?' do
+    it 'returns true if `ActiveRecord::Base.connection` succeeds' do
+      expect(connection.scope).to receive(:connection)
+
+      expect(connection.exists?).to be(true)
+    end
+
+    it 'returns false if `ActiveRecord::Base.connection` fails' do
+      expect(connection.scope).to receive(:connection) do
+        raise ActiveRecord::NoDatabaseError, 'broken'
+      end
+
+      expect(connection.exists?).to be(false)
+    end
+  end
+
+  describe '#system_id' do
+    it 'returns the PostgreSQL system identifier' do
+      expect(connection.system_id).to be_an_instance_of(Integer)
+    end
+  end
+
+  describe '#get_write_location' do
+    it 'returns a string' do
+      expect(connection.get_write_location(connection.scope.connection))
+        .to be_a(String)
+    end
+
+    it 'returns nil if there are no results' do
+      expect(connection.get_write_location(double(select_all: []))).to be_nil
+    end
+  end
+end

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -42,7 +42,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
 
       original_db_config = Gitlab::Database.config
       modified_db_config = original_db_config.merge(load_balancing: lb_config)
-      expect(Gitlab::Database).to receive(:config).and_return(modified_db_config)
+      expect(Gitlab::Database.main).to receive(:config).and_return(modified_db_config)
 
       expect(described_class.configuration).to eq(lb_config)
     end
@@ -401,7 +401,7 @@ RSpec.describe Gitlab::Database::LoadBalancing do
 
         original_db_config = Gitlab::Database.config
         modified_db_config = original_db_config.merge(load_balancing: { hosts: hosts })
-        allow(Gitlab::Database).to receive(:config).and_return(modified_db_config)
+        allow(Gitlab::Database.main).to receive(:config).and_return(modified_db_config)
 
         ::Gitlab::Database::LoadBalancing::Session.clear_session
       end

spec/lib/gitlab/database_spec.rb

@@ -15,32 +15,6 @@ RSpec.describe Gitlab::Database do
     end
   end
 
-  describe '.default_pool_size' do
-    before do
-      allow(Gitlab::Runtime).to receive(:max_threads).and_return(7)
-    end
-
-    it 'returns the max thread size plus a fixed headroom of 10' do
-      expect(described_class.default_pool_size).to eq(17)
-    end
-
-    it 'returns the max thread size plus a DB_POOL_HEADROOM if this env var is present' do
-      stub_env('DB_POOL_HEADROOM', '7')
-
-      expect(described_class.default_pool_size).to eq(14)
-    end
-  end
-
-  describe '.config' do
-    it 'returns a HashWithIndifferentAccess' do
-      expect(described_class.config).to be_an_instance_of(HashWithIndifferentAccess)
-    end
-
-    it 'returns a default pool size' do
-      expect(described_class.config).to include(pool: described_class.default_pool_size)
-    end
-  end
-
   describe '.has_config?' do
     context 'two tier database config' do
       before do
@@ -114,108 +88,11 @@ RSpec.describe Gitlab::Database do
     end
   end
 
-  describe '.adapter_name' do
-    it 'returns the name of the adapter' do
-      expect(described_class.adapter_name).to be_an_instance_of(String)
-    end
-
-    it 'returns Unknown when using anything else' do
-      allow(described_class).to receive(:postgresql?).and_return(false)
-
-      expect(described_class.human_adapter_name).to eq('Unknown')
-    end
-  end
-
-  describe '.human_adapter_name' do
-    it 'returns PostgreSQL when using PostgreSQL' do
-      expect(described_class.human_adapter_name).to eq('PostgreSQL')
-    end
-  end
-
-  describe '.system_id' do
-    it 'returns the PostgreSQL system identifier' do
-      expect(described_class.system_id).to be_an_instance_of(Integer)
-    end
-  end
-
-  describe '.disable_prepared_statements' do
-    around do |example|
-      original_config = ::Gitlab::Database.config
-
-      example.run
-
-      ActiveRecord::Base.establish_connection(original_config)
-    end
-
-    it 'disables prepared statements' do
-      ActiveRecord::Base.establish_connection(::Gitlab::Database.config.merge(prepared_statements: true))
-      expect(ActiveRecord::Base.connection.prepared_statements).to eq(true)
-
-      expect(ActiveRecord::Base).to receive(:establish_connection)
-        .with(a_hash_including({ 'prepared_statements' => false })).and_call_original
-
-      described_class.disable_prepared_statements
-
-      expect(ActiveRecord::Base.connection.prepared_statements).to eq(false)
-    end
-  end
-
-  describe '.postgresql?' do
-    subject { described_class.postgresql? }
-
-    it { is_expected.to satisfy { |val| val == true || val == false } }
-  end
-
-  describe '.version' do
-    around do |example|
-      described_class.instance_variable_set(:@version, nil)
-      example.run
-      described_class.instance_variable_set(:@version, nil)
-    end
-
-    context "on postgresql" do
-      it "extracts the version number" do
-        allow(described_class).to receive(:database_version)
-          .and_return("PostgreSQL 9.4.4 on x86_64-apple-darwin14.3.0")
-
-        expect(described_class.version).to eq '9.4.4'
-      end
-    end
-
-    it 'memoizes the result' do
-      count = ActiveRecord::QueryRecorder
-        .new { 2.times { described_class.version } }
-        .count
-
-      expect(count).to eq(1)
-    end
-  end
-
-  describe '.postgresql_minimum_supported_version?' do
-    it 'returns false when using PostgreSQL 10' do
-      allow(described_class).to receive(:version).and_return('10')
-
-      expect(described_class.postgresql_minimum_supported_version?).to eq(false)
-    end
-
-    it 'returns false when using PostgreSQL 11' do
-      allow(described_class).to receive(:version).and_return('11')
-
-      expect(described_class.postgresql_minimum_supported_version?).to eq(false)
-    end
-
-    it 'returns true when using PostgreSQL 12' do
-      allow(described_class).to receive(:version).and_return('12')
-
-      expect(described_class.postgresql_minimum_supported_version?).to eq(true)
-    end
-  end
-
   describe '.check_postgres_version_and_print_warning' do
     subject { described_class.check_postgres_version_and_print_warning }
 
     it 'prints a warning if not compliant with minimum postgres version' do
-      allow(described_class).to receive(:postgresql_minimum_supported_version?).and_return(false)
+      allow(described_class.main).to receive(:postgresql_minimum_supported_version?).and_return(false)
 
       expect(Kernel).to receive(:warn).with(/You are using PostgreSQL/)
 
@@ -223,7 +100,7 @@ RSpec.describe Gitlab::Database do
     end
 
     it 'doesnt print a warning if compliant with minimum postgres version' do
-      allow(described_class).to receive(:postgresql_minimum_supported_version?).and_return(true)
+      allow(described_class.main).to receive(:postgresql_minimum_supported_version?).and_return(true)
 
       expect(Kernel).not_to receive(:warn).with(/You are using PostgreSQL/)
 
@@ -231,7 +108,7 @@ RSpec.describe Gitlab::Database do
     end
 
     it 'doesnt print a warning in Rails runner environment' do
-      allow(described_class).to receive(:postgresql_minimum_supported_version?).and_return(false)
+      allow(described_class.main).to receive(:postgresql_minimum_supported_version?).and_return(false)
       allow(Gitlab::Runtime).to receive(:rails_runner?).and_return(true)
 
       expect(Kernel).not_to receive(:warn).with(/You are using PostgreSQL/)
@@ -240,13 +117,13 @@ RSpec.describe Gitlab::Database do
     end
 
     it 'ignores ActiveRecord errors' do
-      allow(described_class).to receive(:postgresql_minimum_supported_version?).and_raise(ActiveRecord::ActiveRecordError)
+      allow(described_class.main).to receive(:postgresql_minimum_supported_version?).and_raise(ActiveRecord::ActiveRecordError)
 
       expect { subject }.not_to raise_error
     end
 
     it 'ignores Postgres errors' do
-      allow(described_class).to receive(:postgresql_minimum_supported_version?).and_raise(PG::Error)
+      allow(described_class.main).to receive(:postgresql_minimum_supported_version?).and_raise(PG::Error)
 
       expect { subject }.not_to raise_error
     end
@@ -262,231 +139,6 @@ RSpec.describe Gitlab::Database do
     it { expect(described_class.nulls_first_order('column', 'DESC')).to eq 'column DESC NULLS FIRST'}
   end
 
-  describe '.with_connection_pool' do
-    it 'creates a new connection pool and disconnect it after used' do
-      closed_pool = nil
-
-      described_class.with_connection_pool(1) do |pool|
-        pool.with_connection do |connection|
-          connection.execute('SELECT 1 AS value')
-        end
-
-        expect(pool).to be_connected
-
-        closed_pool = pool
-      end
-
-      expect(closed_pool).not_to be_connected
-    end
-
-    it 'disconnects the pool even an exception was raised' do
-      error = Class.new(RuntimeError)
-      closed_pool = nil
-
-      begin
-        described_class.with_connection_pool(1) do |pool|
-          pool.with_connection do |connection|
-            connection.execute('SELECT 1 AS value')
-          end
-
-          closed_pool = pool
-
-          raise error, 'boom'
-        end
-      rescue error
-      end
-
-      expect(closed_pool).not_to be_connected
-    end
-  end
-
-  describe '.bulk_insert' do
-    before do
-      allow(described_class).to receive(:connection).and_return(connection)
-      allow(connection).to receive(:quote_column_name, &:itself)
-      allow(connection).to receive(:quote, &:itself)
-      allow(connection).to receive(:execute)
-    end
-
-    let(:connection) { double(:connection) }
-
-    let(:rows) do
-      [
-        { a: 1, b: 2, c: 3 },
-        { c: 6, a: 4, b: 5 }
-      ]
-    end
-
-    it 'does nothing with empty rows' do
-      expect(connection).not_to receive(:execute)
-
-      described_class.bulk_insert('test', [])
-    end
-
-    it 'uses the ordering from the first row' do
-      expect(connection).to receive(:execute) do |sql|
-        expect(sql).to include('(1, 2, 3)')
-        expect(sql).to include('(4, 5, 6)')
-      end
-
-      described_class.bulk_insert('test', rows)
-    end
-
-    it 'quotes column names' do
-      expect(connection).to receive(:quote_column_name).with(:a)
-      expect(connection).to receive(:quote_column_name).with(:b)
-      expect(connection).to receive(:quote_column_name).with(:c)
-
-      described_class.bulk_insert('test', rows)
-    end
-
-    it 'quotes values' do
-      1.upto(6) do |i|
-        expect(connection).to receive(:quote).with(i)
-      end
-
-      described_class.bulk_insert('test', rows)
-    end
-
-    it 'does not quote values of a column in the disable_quote option' do
-      [1, 2, 4, 5].each do |i|
-        expect(connection).to receive(:quote).with(i)
-      end
-
-      described_class.bulk_insert('test', rows, disable_quote: :c)
-    end
-
-    it 'does not quote values of columns in the disable_quote option' do
-      [2, 5].each do |i|
-        expect(connection).to receive(:quote).with(i)
-      end
-
-      described_class.bulk_insert('test', rows, disable_quote: [:a, :c])
-    end
-
-    it 'handles non-UTF-8 data' do
-      expect { described_class.bulk_insert('test', [{ a: "\255" }]) }.not_to raise_error
-    end
-
-    context 'when using PostgreSQL' do
-      it 'allows the returning of the IDs of the inserted rows' do
-        result = double(:result, values: [['10']])
-
-        expect(connection)
-          .to receive(:execute)
-          .with(/RETURNING id/)
-          .and_return(result)
-
-        ids = described_class
-          .bulk_insert('test', [{ number: 10 }], return_ids: true)
-
-        expect(ids).to eq([10])
-      end
-
-      it 'allows setting the upsert to do nothing' do
-        expect(connection)
-          .to receive(:execute)
-          .with(/ON CONFLICT DO NOTHING/)
-
-        described_class
-          .bulk_insert('test', [{ number: 10 }], on_conflict: :do_nothing)
-      end
-    end
-  end
-
-  describe '.create_connection_pool' do
-    it 'creates a new connection pool with specific pool size' do
-      pool = described_class.create_connection_pool(5)
-
-      begin
-        expect(pool)
-          .to be_kind_of(ActiveRecord::ConnectionAdapters::ConnectionPool)
-
-        expect(pool.db_config.pool).to eq(5)
-      ensure
-        pool.disconnect!
-      end
-    end
-
-    it 'allows setting of a custom hostname' do
-      pool = described_class.create_connection_pool(5, '127.0.0.1')
-
-      begin
-        expect(pool.db_config.host).to eq('127.0.0.1')
-      ensure
-        pool.disconnect!
-      end
-    end
-
-    it 'allows setting of a custom hostname and port' do
-      pool = described_class.create_connection_pool(5, '127.0.0.1', 5432)
-
-      begin
-        expect(pool.db_config.host).to eq('127.0.0.1')
-        expect(pool.db_config.configuration_hash[:port]).to eq(5432)
-      ensure
-        pool.disconnect!
-      end
-    end
-  end
-
-  describe '.cached_column_exists?' do
-    it 'only retrieves data once' do
-      expect(ActiveRecord::Base.connection).to receive(:columns).once.and_call_original
-
-      2.times do
-        expect(described_class.cached_column_exists?(:projects, :id)).to be_truthy
-        expect(described_class.cached_column_exists?(:projects, :bogus_column)).to be_falsey
-      end
-    end
-  end
-
-  describe '.cached_table_exists?' do
-    it 'only retrieves data once per table' do
-      expect(ActiveRecord::Base.connection).to receive(:data_source_exists?).with(:projects).once.and_call_original
-      expect(ActiveRecord::Base.connection).to receive(:data_source_exists?).with(:bogus_table_name).once.and_call_original
-
-      2.times do
-        expect(described_class.cached_table_exists?(:projects)).to be_truthy
-        expect(described_class.cached_table_exists?(:bogus_table_name)).to be_falsey
-      end
-    end
-
-    it 'returns false when database does not exist' do
-      expect(ActiveRecord::Base).to receive(:connection) { raise ActiveRecord::NoDatabaseError, 'broken' }
-
-      expect(described_class.cached_table_exists?(:projects)).to be(false)
-    end
-  end
-
-  describe '.exists?' do
-    it 'returns true if `ActiveRecord::Base.connection` succeeds' do
-      expect(ActiveRecord::Base).to receive(:connection)
-
-      expect(described_class.exists?).to be(true)
-    end
-
-    it 'returns false if `ActiveRecord::Base.connection` fails' do
-      expect(ActiveRecord::Base).to receive(:connection) { raise ActiveRecord::NoDatabaseError, 'broken' }
-
-      expect(described_class.exists?).to be(false)
-    end
-  end
-
-  describe '.get_write_location' do
-    it 'returns a string' do
-      connection = ActiveRecord::Base.connection
-
-      expect(described_class.get_write_location(connection)).to be_a(String)
-    end
-
-    it 'returns nil if there are no results' do
-      connection = double(select_all: [])
-
-      expect(described_class.get_write_location(connection)).to be_nil
-    end
-  end
-
   describe '.db_config_name' do
     it 'returns the db_config name for the connection' do
       connection = ActiveRecord::Base.connection
@@ -516,42 +168,6 @@ RSpec.describe Gitlab::Database do
     end
   end
 
-  describe '.read_only?' do
-    it 'returns false' do
-      expect(described_class.read_only?).to be_falsey
-    end
-  end
-
-  describe '.db_read_only?' do
-    before do
-      allow(ActiveRecord::Base.connection).to receive(:execute).and_call_original
-    end
-
-    it 'detects a read-only database' do
-      allow(ActiveRecord::Base.connection).to receive(:execute).with('SELECT pg_is_in_recovery()').and_return([{ "pg_is_in_recovery" => "t" }])
-
-      expect(described_class.db_read_only?).to be_truthy
-    end
-
-    it 'detects a read-only database' do
-      allow(ActiveRecord::Base.connection).to receive(:execute).with('SELECT pg_is_in_recovery()').and_return([{ "pg_is_in_recovery" => true }])
-
-      expect(described_class.db_read_only?).to be_truthy
-    end
-
-    it 'detects a read-write database' do
-      allow(ActiveRecord::Base.connection).to receive(:execute).with('SELECT pg_is_in_recovery()').and_return([{ "pg_is_in_recovery" => "f" }])
-
-      expect(described_class.db_read_only?).to be_falsey
-    end
-
-    it 'detects a read-write database' do
-      allow(ActiveRecord::Base.connection).to receive(:execute).with('SELECT pg_is_in_recovery()').and_return([{ "pg_is_in_recovery" => false }])
-
-      expect(described_class.db_read_only?).to be_falsey
-    end
-  end
-
   describe '#sanitize_timestamp' do
     let(:max_timestamp) { Time.at((1 << 31) - 1) }
 

spec/services/packages/create_event_service_spec.rb

@@ -46,7 +46,7 @@ RSpec.describe Packages::CreateEventService do
 
         context 'on a read-only instance' do
           before do
-            allow(Gitlab::Database).to receive(:read_only?).and_return(true)
+            allow(Gitlab::Database.main).to receive(:read_only?).and_return(true)
           end
 
           it 'does not create an event' do

spec/support/shared_examples/services/boards/create_service_shared_examples.rb

@@ -12,7 +12,7 @@ RSpec.shared_examples 'boards recent visit create service' do
   end
 
   it 'returns nil when database is read only' do
-    allow(Gitlab::Database).to receive(:read_only?) { true }
+    allow(Gitlab::Database.main).to receive(:read_only?) { true }
 
     expect(service.execute(board)).to be_nil
   end