Add latest changes from gitlab-org/gitlab@master

1 files changed, 74 insertions, 185 deletions

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index 75738a62305..1fb1faf6db5 100644
--- a/lib/gitlab/database.rb
+++ b/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
-    end
+    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,93 +78,90 @@ 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?
 
-      Kernel.warn ERB.new(Rainbow.new.wrap(<<~EOS).red).result
-
-                  ██     ██  █████  ██████  ███    ██ ██ ███    ██  ██████ 
-                  ██     ██ ██   ██ ██   ██ ████   ██ ██ ████   ██ ██      
-                  ██  █  ██ ███████ ██████  ██ ██  ██ ██ ██ ██  ██ ██   ███ 
-                  ██ ███ ██ ██   ██ ██   ██ ██  ██ ██ ██ ██  ██ ██ ██    ██ 
-                   ███ ███  ██   ██ ██   ██ ██   ████ ██ ██   ████  ██████  
-
-        ******************************************************************************
-          You are using PostgreSQL <%= Gitlab::Database.version %>, 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:
-            https://gitlab.com/gitlab-org/gitlab-development-kit/-/blob/main/doc/howto/postgresql.md#upgrade-postgresql
-          <% end %>
-          Please upgrade your environment to a supported PostgreSQL version, see
-          https://docs.gitlab.com/ee/install/requirements.html#database for details.
-        ******************************************************************************
-      EOS
-    rescue ActiveRecord::ActiveRecordError, PG::Error
-      # ignore - happens when Rake tasks yet have to create a database, e.g. for testing
+      DATABASES.each do |name, connection|
+        next if connection.postgresql_minimum_supported_version?
+
+        Kernel.warn ERB.new(Rainbow.new.wrap(<<~EOS).red).result
+
+                    ██     ██  █████  ██████  ███    ██ ██ ███    ██  ██████ 
+                    ██     ██ ██   ██ ██   ██ ████   ██ ██ ████   ██ ██      
+                    ██  █  ██ ███████ ██████  ██ ██  ██ ██ ██ ██  ██ ██   ███ 
+                    ██ ███ ██ ██   ██ ██   ██ ██  ██ ██ ██ ██  ██ ██ ██    ██ 
+                     ███ ███  ██   ██ ██   ██ ██   ████ ██ ██   ████  ██████  
+
+          ******************************************************************************
+            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:
+              https://gitlab.com/gitlab-org/gitlab-development-kit/-/blob/main/doc/howto/postgresql.md#upgrade-postgresql
+            <% end %>
+            Please upgrade your environment to a supported PostgreSQL version, see
+            https://docs.gitlab.com/ee/install/requirements.html#database for details.
+          ******************************************************************************
+        EOS
+      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)
@@ -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']
+      main.system_id
     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
-            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,31 +267,17 @@ 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
-    end
-
-    def self.open_transactions_baseline
-      @open_transactions_baseline ||= 0
+      main.reset_open_transactions_baseline
     end
-    private_class_method :open_transactions_baseline
 
     # Monkeypatch rails with upgraded database observability
     def self.install_monkey_patches
@@ -409,5 +300,3 @@ module Gitlab
     end
   end
 end
-
-Gitlab::Database.prepend_mod_with('Gitlab::Database')