From a2a5af3a4b673da1cee3ad01c9943d8bc8365de9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kamil=20Trzci=C5=84ski?= <ayufan@ayufan.eu>
Date: Fri, 20 May 2022 17:25:07 +0200
Subject: [PATCH] Fix handling for `GitlabSchemasValidateConnection` for early
 start

It appears that each connection needs to be adopted. However,
in it's initial phase (before being adopted) the Rails
might execute SQL queries to configure the connection settings
(collation, timeouts, etc.). This will happen before
the connection is actually assigned to the connection pool,
as if such connection will be failed to be initialized
it will be orphaned.
---
 lib/gitlab/database.rb           |  9 +++++++--
 spec/lib/gitlab/database_spec.rb | 13 +++++++++----
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index 677b4485288fe..909dfc4f92d9a 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -203,8 +203,13 @@ def self.db_config_names
     # This does not look at literal connection names, but rather compares
     # models that are holders for a given db_config_name
     def self.gitlab_schemas_for_connection(connection)
-      db_name = self.db_config_name(connection)
-      primary_model = self.database_base_models.fetch(db_name.to_sym)
+      db_config = self.db_config_for_connection(connection)
+
+      # connection might not be yet adopted (returning NullPool, and no connection_klass)
+      # in such cases it is fine to ignore such connections
+      return unless db_config
+
+      primary_model = self.database_base_models.fetch(db_config.name.to_sym)
 
       self.schemas_to_base_models.select do |_, child_models|
         child_models.any? do |child_model|
diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb
index 23f4f0e708933..025cf05424e02 100644
--- a/spec/lib/gitlab/database_spec.rb
+++ b/spec/lib/gitlab/database_spec.rb
@@ -222,10 +222,6 @@
   end
 
   describe '.gitlab_schemas_for_connection' do
-    it 'does raise exception for invalid connection' do
-      expect { described_class.gitlab_schemas_for_connection(:invalid) }.to raise_error /key not found: "unknown"/
-    end
-
     it 'does return a valid schema depending on a base model used', :request_store do
       # FF due to lib/gitlab/database/load_balancing/configuration.rb:92
       stub_feature_flags(force_no_sharing_primary_model: true)
@@ -282,6 +278,15 @@
         end
       end
     end
+
+    it 'does return empty for non-adopted connections' do
+      new_connection = ActiveRecord::Base.postgresql_connection(
+        ActiveRecord::Base.connection_db_config.configuration_hash)
+
+      expect(described_class.gitlab_schemas_for_connection(new_connection)).to be_nil
+    ensure
+      new_connection&.disconnect!
+    end
   end
 
   describe '#true_value' do
-- 
GitLab