BP-2147 | Hotfix: Added support for Rails 6 multi-database switching

[chrisyuska]

In testing the Rails 6 final release, seeding the test database was failing due to the the new way connections are handled. It appears with Rails 6 that establishing the DB connection within the connection pool no longer automatically connects at the ActiveRecord::Base level. This adds a manual call to establish the connection when it's just added to the pool.

In addition to fixing the hard fails in Rails 6, this appears to fix migrations in Rails 6. As of 6.0.0.rc1, migrations in tenant databases has been broken. We'll likely need to look at patching any out-of-sync database schemas due to this issue, but that will be a separate ticket if it's needed.

There was an additional issue discovered where long-running transactions with different databases in concurrent threads will almost always result in a deadlock. I believe I've narrowed the Rails change that causes this to the following commit in v6.0.0.rc2: 853f5680e8c79a51f5c0bf1624b1a5d3d3df2802. This merges this PR: rails/rails#36618 which is meant to "Fix query cache when using shared connections."

Here's the diff for that PR:

diff --git a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
index 2c6dabbb0e..018f2647c7 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb
@@ -428,7 +428,7 @@ def lock_thread=(lock_thread)
       # #connection can be called any number of times; the connection is
       # held in a cache keyed by a thread.
       def connection
-        @thread_cached_conns[connection_cache_key(@lock_thread || Thread.current)] ||= checkout
+        @thread_cached_conns[connection_cache_key(current_thread)] ||= checkout
       end
 
       # Returns true if there is an open connection being used for the current thread.
@@ -437,7 +437,7 @@ def connection
       # #connection or #with_connection methods. Connections obtained through
       # #checkout will not be detected by #active_connection?
       def active_connection?
-        @thread_cached_conns[connection_cache_key(Thread.current)]
+        @thread_cached_conns[connection_cache_key(current_thread)]
       end
 
       # Signal that the thread is finished with the current connection.
@@ -732,6 +732,10 @@ def connection_cache_key(thread)
           thread
         end
 
+        def current_thread
+          @lock_thread || Thread.current
+        end
+
         # Take control of all existing connections so a "group" action such as
         # reload/disconnect can be performed safely. It is no longer enough to
         # wrap it in +synchronize+ because some pool's actions are allowed
diff --git a/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb b/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb
index a7753e3e9c..e6fcd2870f 100644
--- a/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb
+++ b/activerecord/lib/active_record/connection_adapters/abstract/query_cache.rb
@@ -33,17 +33,17 @@ def initialize(*)
         end
 
         def enable_query_cache!
-          @query_cache_enabled[connection_cache_key(Thread.current)] = true
+          @query_cache_enabled[connection_cache_key(current_thread)] = true
           connection.enable_query_cache! if active_connection?
         end
 
         def disable_query_cache!
-          @query_cache_enabled.delete connection_cache_key(Thread.current)
+          @query_cache_enabled.delete connection_cache_key(current_thread)
           connection.disable_query_cache! if active_connection?
         end
 
         def query_cache_enabled
-          @query_cache_enabled[connection_cache_key(Thread.current)]
+          @query_cache_enabled[connection_cache_key(current_thread)]
         end
       end
 
diff --git a/activerecord/test/cases/query_cache_test.rb b/activerecord/test/cases/query_cache_test.rb
index eb32b690aa..4024b26a31 100644
--- a/activerecord/test/cases/query_cache_test.rb
+++ b/activerecord/test/cases/query_cache_test.rb
@@ -540,6 +540,23 @@ def test_clear_query_cache_is_called_on_all_connections
     ActiveRecord::Base.connection_handlers = { writing: ActiveRecord::Base.default_connection_handler }
   end
 
+  test "query cache is enabled in threads with shared connection" do
+    ActiveRecord::Base.connection_pool.lock_thread = true
+
+    assert_cache :off
+
+    thread_a = Thread.new do
+      middleware { |env|
+        assert_cache :clean
+        [200, {}, nil]
+      }.call({})
+    end
+
+    thread_a.join
+
+    ActiveRecord::Base.connection_pool.lock_thread = false
+  end
+
   private
 
     def with_temporary_connection_pool

Specifically, the change in the active_connection? method to also use the @lock_threads from the current_thread method when present broke it. By changing Apartment.connection.clear_query_cache to Apartment.connection_class.clear_query_caches_for_current_thread we should now be clearing all query caches, including shared ones between threads in the test environment.

Additionally, two other PRs were merged into Rails (rails/rails#36870 and rails/rails#36883) that broke the ability to seed the multiple databases with fixture data. It appears the new connection is not automatically established unless the Apartment.connection_class's @connection_specification_name == nil. I believe that rails/rails#36870 established a connection now, which sets the instance variable during the schema check. This is now pointing at the wrong database, so we can fix it by setting it on the connection switch.

[chrisyuska]

I'm giving up on getting this to pass in Travis for now... Would still appreciate a review anyway.

[nickjer]

Googling around, the only thing I worry about when calling Apartment.establish_connection(config) directly is that you may not be able to maintain multiple transactions on different databases.

Do you think it would be possible to test this by opening up two long running transactions on the same model across two different databases and confirming they are both successful?

[nickjer]

Going to trust you on this as you satisfied my only fear on this.