Add `main_database?` helper

This ensures that we use `main:`
as an indicator for a main unsharded
database.

This adds a helper method, and additionally
updated `database.yml.postgresql` to properly
support that for `rspec`.

config/database.yml.postgresql

@@ -2,6 +2,7 @@
 # PRODUCTION
 #
 production:
+  main:
     adapter: postgresql
     encoding: unicode
     database: gitlabhq_production
@@ -22,6 +23,7 @@ production:
 # Development specific
 #
 development:
+  main:
     adapter: postgresql
     encoding: unicode
     database: gitlabhq_development
@@ -35,6 +37,7 @@ development:
 # Staging specific
 #
 staging:
+  main:
     adapter: postgresql
     encoding: unicode
     database: gitlabhq_staging
@@ -46,6 +49,7 @@ staging:
 # re-generated from your development database when you run "rake".
 # Do not set this db to the same as development or production.
 test: &test
+  main:
     adapter: postgresql
     encoding: unicode
     database: gitlabhq_test

config/initializers/active_record_force_reconnects.rb

 # frozen_string_literal: true
 
 Gitlab::Database::ConnectionTimer.configure do |config|
-  config.interval = Rails.application.config_for(:database)[:force_reconnect_interval]
+  configuration_hash = ActiveRecord::Base.configurations.find_db_config(Rails.env).configuration_hash
+  config.interval = configuration_hash[:force_reconnect_interval]
 end
 
 ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(Gitlab::Database::PostgresqlAdapter::ForceDisconnectableMixin)

lib/backup/database.rb

@@ -18,7 +18,7 @@ module Backup
 
     def initialize(progress, filename: nil)
       @progress = progress
-      @config = YAML.load_file(File.join(Rails.root, 'config', 'database.yml'))[Rails.env]
+      @config = ActiveRecord::Base.configurations.find_db_config(Rails.env).configuration_hash
       @db_file_name = filename || File.join(Gitlab.config.backup.path, 'db', 'database.sql.gz')
     end
 
@@ -30,9 +30,9 @@ module Backup
       compress_rd.close
 
       dump_pid =
-        case config["adapter"]
+        case config[:adapter]
         when "postgresql" then
-          progress.print "Dumping PostgreSQL database #{config['database']} ... "
+          progress.print "Dumping PostgreSQL database #{database} ... "
           pg_env
           pgsql_args = ["--clean"] # Pass '--clean' to include 'DROP TABLE' statements in the DB dump.
           pgsql_args << '--if-exists'
@@ -47,7 +47,7 @@ module Backup
             end
           end
 
-          Process.spawn('pg_dump', *pgsql_args, config['database'], out: compress_wr)
+          Process.spawn('pg_dump', *pgsql_args, database, out: compress_wr)
         end
       compress_wr.close
 
@@ -68,9 +68,9 @@ module Backup
       decompress_wr.close
 
       status, errors =
-        case config["adapter"]
+        case config[:adapter]
         when "postgresql" then
-          progress.print "Restoring PostgreSQL database #{config['database']} ... "
+          progress.print "Restoring PostgreSQL database #{database} ... "
           pg_env
           execute_and_track_errors(pg_restore_cmd, decompress_rd)
         end
@@ -93,6 +93,10 @@ module Backup
 
     protected
 
+    def database
+      @config[:database]
+    end
+
     def ignore_error?(line)
       IGNORED_ERRORS_REGEXP.match?(line)
     end
@@ -128,17 +132,17 @@ module Backup
 
     def pg_env
       args = {
-        'username'  => 'PGUSER',
-        'host'      => 'PGHOST',
-        'port'      => 'PGPORT',
-        'password'  => 'PGPASSWORD',
+        username: 'PGUSER',
+        host: 'PGHOST',
+        port: 'PGPORT',
+        password: 'PGPASSWORD',
         # SSL
-        'sslmode'         => 'PGSSLMODE',
-        'sslkey'          => 'PGSSLKEY',
-        'sslcert'         => 'PGSSLCERT',
-        'sslrootcert'     => 'PGSSLROOTCERT',
-        'sslcrl'          => 'PGSSLCRL',
-        'sslcompression'  => 'PGSSLCOMPRESSION'
+        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
@@ -161,7 +165,7 @@ module Backup
     private
 
     def pg_restore_cmd
-      ['psql', config['database']]
+      ['psql', database]
     end
   end
 end

lib/gitlab/database.rb

@@ -72,8 +72,15 @@ module Gitlab
       Gitlab::Application.config.database_configuration[Rails.env].include?(database_name.to_s)
     end
 
+    def self.main_database?(name)
+      # The database is `main` if it is a first entry in `database.yml`
+      # Rails internally names them `primary` to avoid confusion
+      # with broad `primary` usage we use `main` instead
+      ActiveRecord::Base.configurations.primary?(name.to_s)
+    end
+
     def self.ci_database?(name)
-      name == CI_DATABASE_NAME
+      name.to_s == CI_DATABASE_NAME
     end
 
     def self.username

spec/lib/gitlab/database/postgresql_adapter/force_disconnectable_mixin_spec.rb

@@ -14,7 +14,7 @@ RSpec.describe Gitlab::Database::PostgresqlAdapter::ForceDisconnectableMixin do
       end
     end
 
-    let(:config) { Rails.application.config_for(:database).merge(pool: 1) }
+    let(:config) { ActiveRecord::Base.configurations.find_db_config(Rails.env).configuration_hash.merge(pool: 1) }
     let(:pool) { model.establish_connection(config) }
 
     it 'calls the force disconnect callback on checkin' do

spec/lib/gitlab/database/postgresql_adapter/type_map_cache_spec.rb

@@ -3,7 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Database::PostgresqlAdapter::TypeMapCache do
-  let(:db_config) { ActiveRecord::Base.configurations.configs_for(env_name: 'test', name: 'primary').configuration_hash }
+  let(:db_config) { ActiveRecord::Base.configurations.configs_for(env_name: 'test', name: 'main').configuration_hash }
   let(:adapter_class) { ActiveRecord::ConnectionAdapters::PostgreSQLAdapter }
 
   before do

spec/lib/gitlab/database/with_lock_retries_outside_transaction_spec.rb

@@ -37,8 +37,10 @@ RSpec.describe Gitlab::Database::WithLockRetriesOutsideTransaction do
     context 'when lock retry is enabled' do
       let(:lock_fiber) do
         Fiber.new do
+          configuration = ActiveRecordSecond.configurations.find_db_config(Rails.env).configuration_hash
+
           # Initiating a second DB connection for the lock
-          conn = ActiveRecordSecond.establish_connection(Rails.configuration.database_configuration[Rails.env]).connection
+          conn = ActiveRecordSecond.establish_connection(configuration).connection
           conn.transaction do
             conn.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
 

spec/lib/gitlab/database/with_lock_retries_spec.rb

@@ -37,8 +37,10 @@ RSpec.describe Gitlab::Database::WithLockRetries do
     context 'when lock retry is enabled' do
       let(:lock_fiber) do
         Fiber.new do
+          configuration = ActiveRecordSecond.configurations.find_db_config(Rails.env).configuration_hash
+
           # Initiating a second DB connection for the lock
-          conn = ActiveRecordSecond.establish_connection(Rails.configuration.database_configuration[Rails.env]).connection
+          conn = ActiveRecordSecond.establish_connection(configuration).connection
           conn.transaction do
             conn.execute("LOCK TABLE #{Project.table_name} in exclusive mode")
 

spec/lib/gitlab/database_spec.rb

@@ -80,6 +80,40 @@ RSpec.describe Gitlab::Database do
     end
   end
 
+  describe '.main_database?' do
+    using RSpec::Parameterized::TableSyntax
+
+    where(:database_name, :result) do
+      :main     | true
+      'main'    | true
+      :ci       | false
+      'ci'      | false
+      :archive  | false
+      'archive' | false
+    end
+
+    with_them do
+      it { expect(described_class.main_database?(database_name)).to eq(result) }
+    end
+  end
+
+  describe '.ci_database?' do
+    using RSpec::Parameterized::TableSyntax
+
+    where(:database_name, :result) do
+      :main     | false
+      'main'    | false
+      :ci       | true
+      'ci'      | true
+      :archive  | false
+      'archive' | false
+    end
+
+    with_them do
+      it { expect(described_class.ci_database?(database_name)).to eq(result) }
+    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)