1 files changed, 64 insertions, 13 deletions

diff --git a/lib/tasks/gitlab/db/validate_config.rake b/lib/tasks/gitlab/db/validate_config.rake
index 66aa949cc94..2a3a54b5351 100644
--- a/lib/tasks/gitlab/db/validate_config.rake
+++ b/lib/tasks/gitlab/db/validate_config.rake
@@ -4,6 +4,23 @@ databases = ActiveRecord::Tasks::DatabaseTasks.setup_initial_database_yaml
 
 namespace :gitlab do
   namespace :db do
+    DB_CONFIG_NAME_KEY = 'gitlab_db_config_name'
+
+    DB_IDENTIFIER_SQL = <<-SQL
+      SELECT system_identifier, current_database()
+      FROM pg_control_system()
+    SQL
+
+    # We fetch timestamp as a way to properly handle race conditions
+    # fail in such cases, which should not really happen in production environment
+    DB_IDENTIFIER_WITH_DB_CONFIG_NAME_SQL = <<-SQL
+      SELECT
+        system_identifier, current_database(),
+        value as db_config_name, created_at as timestamp
+      FROM pg_control_system()
+      LEFT JOIN ar_internal_metadata ON ar_internal_metadata.key=$1
+    SQL
+
     desc 'Validates `config/database.yml` to ensure a correct behavior is configured'
     task validate_config: :environment do
       original_db_config = ActiveRecord::Base.connection_db_config # rubocop:disable Database/MultipleDatabases
@@ -14,26 +31,22 @@ namespace :gitlab do
       db_configs = ActiveRecord::Base.configurations.configs_for(env_name: Rails.env, include_replicas: true)
       db_configs = db_configs.reject(&:replica?)
 
+      # The `pg_control_system()` is not enough to properly discover matching database systems
+      # since in case of cluster promotion it will return the same identifier as main cluster
+      # We instead set an `ar_internal_metadata` information with configured database name
+      db_configs.reverse_each do |db_config|
+        insert_db_identifier(db_config)
+      end
+
       # Map each database connection into unique identifier of system+database
-      # rubocop:disable Database/MultipleDatabases
       all_connections = db_configs.map do |db_config|
-        identifier =
-          begin
-            ActiveRecord::Base.establish_connection(db_config) # rubocop: disable Database/EstablishConnection
-            ActiveRecord::Base.connection.select_one("SELECT system_identifier, current_database() FROM pg_control_system()")
-          rescue ActiveRecord::ConnectionNotEstablished, PG::ConnectionBad => err
-            warn "WARNING: Could not establish database connection for #{db_config.name}: #{err.message}"
-          rescue ActiveRecord::NoDatabaseError
-          end
-
         {
           name: db_config.name,
           config: db_config,
           database_tasks?: db_config.database_tasks?,
-          identifier: identifier
+          identifier: get_db_identifier(db_config)
         }
-      end.compact
-      # rubocop:enable Database/MultipleDatabases
+      end
 
       unique_connections = all_connections.group_by { |connection| connection[:identifier] }
       primary_connection = all_connections.find { |connection| ActiveRecord::Base.configurations.primary?(connection[:name]) }
@@ -111,5 +124,43 @@ namespace :gitlab do
       Rake::Task["db:schema:load:#{name}"].enhance(['gitlab:db:validate_config'])
       Rake::Task["db:schema:dump:#{name}"].enhance(['gitlab:db:validate_config'])
     end
+
+    def insert_db_identifier(db_config)
+      ActiveRecord::Base.establish_connection(db_config) # rubocop: disable Database/EstablishConnection
+
+      if ActiveRecord::InternalMetadata.table_exists?
+        ts = Time.zone.now
+
+        ActiveRecord::InternalMetadata.upsert(
+          { key: DB_CONFIG_NAME_KEY,
+            value: db_config.name,
+            created_at: ts,
+            updated_at: ts }
+        )
+      end
+    rescue ActiveRecord::ConnectionNotEstablished, PG::ConnectionBad => err
+      warn "WARNING: Could not establish database connection for #{db_config.name}: #{err.message}"
+    rescue ActiveRecord::NoDatabaseError
+    rescue ActiveRecord::StatementInvalid => err
+      raise unless err.cause.is_a?(PG::ReadOnlySqlTransaction)
+
+      warn "WARNING: Could not write to the database #{db_config.name}: #{err.message}"
+    end
+
+    def get_db_identifier(db_config)
+      ActiveRecord::Base.establish_connection(db_config) # rubocop: disable Database/EstablishConnection
+
+      # rubocop:disable Database/MultipleDatabases
+      if ActiveRecord::InternalMetadata.table_exists?
+        ActiveRecord::Base.connection.select_one(
+          DB_IDENTIFIER_WITH_DB_CONFIG_NAME_SQL, nil, [DB_CONFIG_NAME_KEY])
+      else
+        ActiveRecord::Base.connection.select_one(DB_IDENTIFIER_SQL)
+      end
+      # rubocop:enable Database/MultipleDatabases
+    rescue ActiveRecord::ConnectionNotEstablished, PG::ConnectionBad => err
+      warn "WARNING: Could not establish database connection for #{db_config.name}: #{err.message}"
+    rescue ActiveRecord::NoDatabaseError
+    end
   end
 end