active_record_host_pool 1.0.1 → 1.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: a62e23b1cbccc5365b119d31329f4c227a6efebf729ce49edb5f2dd63d704e39
-  data.tar.gz: 815dcd119bebc0ddb5ad7308ca41b4906e15cc72ede2048d164623cdd539df7f
+  metadata.gz: 7baaf5e8c49a8308e3be6a9f52b41af880f0d8182a98a12be2ce8fa6d7ccbe60
+  data.tar.gz: aa68d84f75c00f608d13d475761f936722171f3d3d59a4ede471513279e8cda2
 SHA512:
-  metadata.gz: 8c4022f240849bdc6be2140d5b6abda363264e1f1011e6781436537516c356c2ac4a05b82d30891143274113cf18575b0817b76ccc5f22a6fcb0ebd190148db1
-  data.tar.gz: dd7d2ffed45316302a921ab8f70c7ae067cef09e5ae01a0961b1f899cc34808e5b6077d488182e590bc88b767dd670d29ebb6b83aaf18222c370a5fdd3c06f12
+  metadata.gz: 072b56581266326d5decfd1221ccb9bf2ffde945cd04e1ce52278dd89d3b9e84b4d59d0a8d8e21808352d25a0f360ed67a006f71dc112751b7937a42f21aa455
+  data.tar.gz: c6075399e78ec309aeca872079044b042a5208dc1071c068eac01928a83b588d23cfc0d268907c0a9aa5c13a7d9b99005443e644b8142cb7ad9cb8a0e7cf013a

data/Readme.md

@@ -1,10 +1,46 @@
-[![Build status](https://circleci.com/gh/zendesk/active_record_host_pool.svg?style=svg)](https://circleci.com/gh/zendesk/active_record_host_pool)
+[![Build Status](https://github.com/zendesk/active_record_host_pool/workflows/CI/badge.svg)](https://github.com/zendesk/active_record_host_pool/actions?query=workflow%3ACI)
 
 # ActiveRecord host pooling
 
 This gem allows for one ActiveRecord connection to be used to connect to multiple databases on a server.
 It accomplishes this by calling select_db() as necessary to switch databases between database calls.
 
+## How Connections Are Pooled
+
+ARHP creates separate connection pools based on the pool key.
+
+The pool key is defined as:
+
+`host / port / socket / username / replica`
+
+Therefore two databases with identical host, port, socket, username, and replica status will share a connection pool.
+If any part (host, port, etc.) of the pool key differ, two databases will _not_ share a connection pool.
+
+`replica` in the pool key is a boolean indicating if the database is a replica/reader (true) or writer database (false).
+
+Below, `test_pool_1` and `test_pool_2` have identical host, username, socket, and replica status but the port information differs.
+Here the database configurations are formatted as a table to give a visual example:
+
+|          |  test_pool_1   |  test_pool_2   |
+|----------|----------------|----------------|
+| host     | 127.0.0.1      | 127.0.0.1      |
+| port     |                | 3306           |
+| socket   |                |                |
+| username | root           | root           |
+| replica  | false          | false          |
+
+The configuration items must be explicitly defined or they will be blank in the pool key.
+Configurations with matching _implicit_ items but differing _explicit_ items will create separate pools.
+e.g. `test_pool_1` will default to port 3306 but because it is not explicitly defined it will not share a pool with `test_pool_2`
+
+ARHP will therefore create the following pool keys:
+
+```
+test_pool_1 => 127.0.0.1///root/false
+test_pool_2 => 127.0.0.1/3306//root/false
+```
+
+
 ## Support
 
 For now, the only backend known to work is MySQL, with the mysql2 gem.

data/lib/active_record_host_pool/clear_query_cache_patch.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+if ActiveRecord.version >= Gem::Version.new('6.0')
+  module ActiveRecordHostPool
+    # ActiveRecord 6.0 introduced multiple database support. With that, an update
+    # has been made in https://github.com/rails/rails/pull/35089 to ensure that
+    # all query caches are cleared across connection handlers and pools. If you
+    # write on one connection, the other connection will have the update that
+    # occurred.
+    #
+    # This broke ARHP which implements its own pool, allowing you to access
+    # multiple databases with the same connection (e.g. 1 connection for 100
+    # shards on the same server).
+    #
+    # This patch maintains the reference to the database during the cache clearing
+    # to ensure that the database doesn't get swapped out mid-way into an
+    # operation.
+    #
+    # This is a private Rails API and may change in future releases as they're
+    # actively working on sharding in Rails 6 and above.
+    module ClearQueryCachePatch
+      def clear_query_caches_for_current_thread
+        host_pool_current_database_was = connection.unproxied._host_pool_current_database
+        super
+      ensure
+        # restore in case clearing the cache changed the database
+        connection.unproxied._host_pool_current_database = host_pool_current_database_was
+      end
+
+      def clear_on_handler(handler)
+        handler.all_connection_pools.each do |pool|
+          db_was = pool.connection.unproxied._host_pool_current_database
+          pool.connection.clear_query_cache if pool.active_connection?
+        ensure
+          pool.connection.unproxied._host_pool_current_database = db_was
+        end
+      end
+    end
+  end
+
+  ActiveRecord::Base.singleton_class.prepend ActiveRecordHostPool::ClearQueryCachePatch
+end

data/lib/active_record_host_pool/connection_adapter_mixin.rb

@@ -10,7 +10,7 @@ module ActiveRecordHostPool
 
         def _host_pool_current_database=(database)
           @_host_pool_current_database = database
-          @config[:database] = _host_pool_current_database if ActiveRecord::VERSION::MAJOR >= 5
+          @config[:database] = _host_pool_current_database
         end
 
         alias_method :execute_without_switching, :execute
@@ -83,13 +83,24 @@ module ActiveRecordHostPool
       super(_host_pool_current_database.to_s + "/" + sql, *args)
     end
   end
+
+  module PoolConfigPatch
+    def pool
+      ActiveSupport::ForkTracker.check!
+
+      @pool || synchronize { @pool ||= ActiveRecordHostPool::PoolProxy.new(self) }
+    end
+  end
 end
 
-# rubocop:disable Lint/DuplicateMethods
 module ActiveRecord
   module ConnectionAdapters
     class ConnectionHandler
       case "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}"
+      when '6.1'
+
+        :noop
+
       when '5.1', '5.2', '6.0'
 
         def establish_connection(spec)
@@ -99,15 +110,6 @@ module ActiveRecord
           owner_to_pool[spec.name] = ActiveRecordHostPool::PoolProxy.new(spec)
         end
 
-      when '4.2'
-
-        def establish_connection(owner, spec)
-          @class_to_pool.clear
-          raise "Anonymous class is not allowed." unless owner.name
-
-          owner_to_pool[owner.name] = ActiveRecordHostPool::PoolProxy.new(spec)
-        end
-
       else
 
         raise "Unsupported version of Rails (v#{ActiveRecord::VERSION::STRING})"
@@ -115,6 +117,11 @@ module ActiveRecord
     end
   end
 end
-# rubocop:enable Lint/DuplicateMethods
 
 ActiveRecord::ConnectionAdapters::Mysql2Adapter.include(ActiveRecordHostPool::DatabaseSwitch)
+
+# In Rails 6.1 Connection Pools are no longer instantiated in #establish_connection but in a
+# new pool method.
+if "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" == '6.1'
+  ActiveRecord::ConnectionAdapters::PoolConfig.prepend(ActiveRecordHostPool::PoolConfigPatch)
+end

data/lib/active_record_host_pool/pool_proxy.rb

@@ -1,158 +1,7 @@
 # frozen_string_literal: true
 
-require 'delegate'
-require 'active_record'
-require 'active_record_host_pool/connection_adapter_mixin'
-
-# this module sits in between ConnectionHandler and a bunch of different ConnectionPools (one per host).
-# when a connection is requested, it goes like:
-# ActiveRecordClass -> ConnectionHandler#connection
-# ConnectionHandler#connection -> (find or create PoolProxy)
-# PoolProxy -> shared list of Pools
-# Pool actually gives back a connection, then PoolProxy turns this
-# into a ConnectionProxy that can inform (on execute) which db we should be on.
-
-module ActiveRecordHostPool
-  class PoolProxy < Delegator
-    def initialize(spec)
-      super(spec)
-      @spec = spec
-      @config = spec.config
-    end
-
-    def __getobj__
-      _connection_pool
-    end
-
-    def __setobj__(spec)
-      @spec = spec
-      @config = spec.config
-      @_pool_key = nil
-    end
-
-    def spec
-      @spec
-    end
-
-    def connection(*args)
-      real_connection = _connection_pool.connection(*args)
-      _connection_proxy_for(real_connection, @config[:database])
-    rescue Exception => e
-      if rescuable_errors.any? { |r| e.is_a?(r) }
-        _connection_pools.delete(_pool_key)
-      end
-      Kernel.raise(e)
-    end
-
-    # by the time we are patched into ActiveRecord, the current thread has already established
-    # a connection.  thus we need to patch both connection and checkout/checkin
-    def checkout(*args, &block)
-      cx = _connection_pool.checkout(*args, &block)
-      _connection_proxy_for(cx, @config[:database])
-    end
-
-    def checkin(cx)
-      cx = cx.unproxied
-      _connection_pool.checkin(cx)
-    end
-
-    def with_connection
-      cx = checkout
-      yield cx
-    ensure
-      checkin cx
-    end
-
-    def disconnect!
-      p = _connection_pool(false)
-      return unless p
-
-      p.disconnect!
-      p.automatic_reconnect = true if p.respond_to?(:automatic_reconnect=)
-      _clear_connection_proxy_cache
-    end
-
-    def automatic_reconnect=(value)
-      p = _connection_pool(false)
-      return unless p
-
-      p.automatic_reconnect = value if p.respond_to?(:automatic_reconnect=)
-    end
-
-    def clear_reloadable_connections!
-      _connection_pool.clear_reloadable_connections!
-      _clear_connection_proxy_cache
-    end
-
-    def release_connection(*args)
-      p = _connection_pool(false)
-      return unless p
-
-      p.release_connection(*args)
-    end
-
-    def flush!
-      p = _connection_pool(false)
-      return unless p
-
-      p.flush!
-    end
-
-    def discard!
-      p = _connection_pool(false)
-      return unless p
-
-      p.discard!
-
-      # All connections in the pool (even if they're currently
-      # leased!) have just been discarded, along with the pool itself.
-      # Any further interaction with the pool (except #spec and #schema_cache)
-      # is undefined.
-      # Remove the connection for the given key so a new one can be created in its place
-      _connection_pools.delete(_pool_key)
-    end
-
-    private
-
-    def rescuable_errors
-      @rescuable_errors ||= begin
-        e = [Mysql2::Error]
-        if ActiveRecord.const_defined?("NoDatabaseError")
-          e << ActiveRecord::NoDatabaseError
-        end
-        e
-      end
-    end
-
-    def _connection_pools
-      @@_connection_pools ||= {}
-    end
-
-    def _pool_key
-      @_pool_key ||= "#{@config[:host]}/#{@config[:port]}/#{@config[:socket]}/#{@config[:username]}/#{@config[:slave] && 'slave'}"
-    end
-
-    def _connection_pool(auto_create = true)
-      pool = _connection_pools[_pool_key]
-      if pool.nil? && auto_create
-        pool = _connection_pools[_pool_key] = ActiveRecord::ConnectionAdapters::ConnectionPool.new(@spec)
-      end
-      pool
-    end
-
-    def _connection_proxy_for(connection, database)
-      @connection_proxy_cache ||= {}
-      key = [connection, database]
-
-      @connection_proxy_cache[key] ||= begin
-        cx = ActiveRecordHostPool::ConnectionProxy.new(connection, database)
-        cx.execute('select 1')
-        cx
-      end
-    end
-
-    def _clear_connection_proxy_cache
-      @connection_proxy_cache = {}
-    end
-  end
+if "#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}" == '6.1'
+  require 'active_record_host_pool/pool_proxy_6_1'
+else
+  require 'active_record_host_pool/pool_proxy_legacy'
 end

data/lib/active_record_host_pool/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module ActiveRecordHostPool
-  VERSION = "1.0.1"
+  VERSION = "1.1.0"
 end

data/lib/active_record_host_pool.rb

@@ -4,6 +4,7 @@ require 'active_record'
 require 'active_record/base'
 require 'active_record/connection_adapters/abstract_adapter'
 
+require 'active_record_host_pool/clear_query_cache_patch'
 require 'active_record_host_pool/connection_proxy'
 require 'active_record_host_pool/pool_proxy'
 require 'active_record_host_pool/connection_adapter_mixin'

data/test/database.yml

@@ -1,67 +1,107 @@
 <% mysql = URI(ENV['MYSQL_URL'] || 'mysql://root@127.0.0.1:3306') %>
-test_host_1_db_1:
+# ARHP creates separate connection pools based on the pool key.
+# The pool key is defined as:
+#   host / port / socket / username / replica
+#
+# Therefore two databases with identical host, port, socket, username, and replica status will share a connection pool.
+# If any part (host, port, etc.) of the pool key differ, two databases will _not_ share a connection pool.
+#
+# Below, "test_pool_1..." and "test_pool_2..." have identical host, username, socket, and replica status but the port information differs.
+# Here the yml configurations are reformatted as a table to give a visual example:
+#
+# |----------+----------------+----------------|
+# |          |  test_pool_1   |  test_pool_2   |
+# |----------+----------------+----------------+
+# | host     | 127.0.0.1      | 127.0.0.1      |
+# | port     |                | 3306           |
+# | socket   |                |                |
+# | username | root           | root           |
+# | replica  | false          | false          |
+# |----------+----------------+----------------|
+#
+# Note: The configuration items must be explicitly defined or will be blank in the pool key.
+#       Configurations with matching _implicit_ items but differing _explicit_ items will create separate pools.
+#       e.g. "test_pool_1" will default to port 3306 but because it is not explicitly defined it will not share a pool with test_pool_2
+#
+#       ARHP will therefore create the following pool keys:
+#       test_pool_1 => 127.0.0.1///root/false
+#       test_pool_2 => 127.0.0.1/3306//root/false
+
+test_pool_1_db_a:
   adapter: mysql2
   encoding: utf8
-  database: arhp_test_1
+  database: arhp_test_db_a
   username: <%= mysql.user %>
   password: <%= mysql.password %>
   host: <%= mysql.host %>
   reconnect: true
 
 # Mimic configurations as read by active_record_shards/ar_flexmaster
-test_host_1_db_1_slave:
+test_pool_1_db_a_replica:
   adapter: mysql2
   encoding: utf8
-  database: arhp_test_1_slave
+  database: arhp_test_db_a_replica
   username: <%= mysql.user %>
   password: <%= mysql.password %>
   host: <%= mysql.host %>
   reconnect: true
   slave: true
 
-test_host_1_db_2:
+test_pool_1_db_b:
+  adapter: mysql2
+  encoding: utf8
+  database: arhp_test_db_b
+  username: <%= mysql.user %>
+  password: <%= mysql.password %>
+  host: <%= mysql.host %>
+  reconnect: true
+
+test_pool_1_db_not_there:
   adapter: mysql2
   encoding: utf8
-  database: arhp_test_2
+  database: arhp_test_db_not_there
   username: <%= mysql.user %>
   password: <%= mysql.password %>
   host: <%= mysql.host %>
   reconnect: true
 
-test_host_2_db_3:
+test_pool_2_db_d:
   adapter: mysql2
   encoding: utf8
-  database: arhp_test_3
+  database: arhp_test_db_d
   username: <%= mysql.user %>
   password: <%= mysql.password %>
   host: <%= mysql.host %>
   port: <%= mysql.port %>
   reconnect: true
 
-test_host_2_db_4:
+test_pool_2_db_e:
   adapter: mysql2
   encoding: utf8
-  database: arhp_test_4
+  database: arhp_test_db_e
   username: <%= mysql.user %>
   password: <%= mysql.password %>
   host: <%= mysql.host %>
   port: <%= mysql.port %>
   reconnect: true
 
-test_host_2_db_5:
+test_pool_3_db_e:
   adapter: mysql2
   encoding: utf8
-  database: arhp_test_4
+  database: arhp_test_db_e
   username: john-doe
   password:
   host: <%= mysql.host %>
   port: <%= mysql.port %>
   reconnect: true
 
-test_host_1_db_not_there:
+# test_pool_1_db_c needs to be the last database defined in the file
+# otherwise the test_models_with_matching_hosts_and_non_matching_databases_issue_exists_without_arhp_patch
+# test fails
+test_pool_1_db_c:
   adapter: mysql2
   encoding: utf8
-  database: arhp_test_no_create
+  database: arhp_test_db_c
   username: <%= mysql.user %>
   password: <%= mysql.password %>
   host: <%= mysql.host %>

data/test/helper.rb

@@ -5,15 +5,41 @@ require 'minitest/autorun'
 
 require 'active_record_host_pool'
 require 'logger'
-require 'mocha/setup'
+require 'mocha/minitest'
 require 'phenix'
 
-RAILS_ENV = 'test'
+ENV['RAILS_ENV'] = 'test'
+ENV['LEGACY_CONNECTION_HANDLING'] = 'true' if ENV['LEGACY_CONNECTION_HANDLING'].nil?
 
-Minitest::Test = MiniTest::Unit::TestCase unless defined?(::Minitest::Test)
+if ActiveRecord.version >= Gem::Version.new('6.1')
+  ActiveRecord::Base.legacy_connection_handling = (ENV['LEGACY_CONNECTION_HANDLING'] == 'true')
+end
+
+RAILS_6_1_WITH_NON_LEGACY_CONNECTION_HANDLING =
+  ActiveRecord.version >= Gem::Version.new('6.1') && !ActiveRecord::Base.legacy_connection_handling
 
 ActiveRecord::Base.logger = Logger.new(File.dirname(__FILE__) + '/test.log')
 
+# BEGIN preventing_writes? patch
+## Rails 6.1 by default does not allow writing to replica databases which prevents
+## us from properly setting up the test databases. This patch is used in test/schema.rb
+## to allow us to write to the replicas but only during migrations
+module ActiveRecordHostPool
+  cattr_accessor :allowing_writes
+  module PreventWritesPatch
+    def preventing_writes?
+      return false if ActiveRecordHostPool.allowing_writes && replica?
+
+      super
+    end
+  end
+end
+
+if ActiveRecord.version >= Gem::Version.new('6.1')
+  ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend ActiveRecordHostPool::PreventWritesPatch
+end
+# END preventing_writes? patch
+
 Phenix.configure do |config|
   config.skip_database = ->(name, conf) { name =~ /not_there/ || conf['username'] == 'john-doe' }
 end
@@ -22,42 +48,148 @@ module ARHPTestSetup
   private
 
   def arhp_create_models
-    return if ARHPTestSetup.const_defined?('Test1')
-
-    eval <<-RUBY
-      class Test1 < ActiveRecord::Base
-        self.table_name = "tests"
-        establish_connection(:test_host_1_db_1)
-      end
-
-      class Test1Slave < ActiveRecord::Base
-        self.table_name = "tests"
-        establish_connection(:test_host_1_db_1_slave)
-      end
-
-      class Test2 < ActiveRecord::Base
-        self.table_name =  "tests"
-        establish_connection(:test_host_1_db_2)
-      end
-
-      class Test3 < ActiveRecord::Base
-        self.table_name = "tests"
-        establish_connection(:test_host_2_db_3)
-      end
-
-      class Test4 < ActiveRecord::Base
-        self.table_name = "tests"
-        establish_connection(:test_host_2_db_4)
-      end
-
-      class Test5 < ActiveRecord::Base
-        self.table_name = "tests"
-        establish_connection(:test_host_2_db_5)
-      end
-    RUBY
+    return if ARHPTestSetup.const_defined?('Pool1DbA')
+
+    if RAILS_6_1_WITH_NON_LEGACY_CONNECTION_HANDLING
+      eval <<-RUBY
+        class AbstractPool1DbC < ActiveRecord::Base
+          self.abstract_class = true
+          connects_to database: { writing: :test_pool_1_db_c }
+        end
+
+        # The placement of the Pool1DbC class is important so that its
+        # connection will not be the most recent connection established
+        # for test_pool_1.
+        class Pool1DbC < AbstractPool1DbC
+        end
+
+        class AbstractPool1DbA < ActiveRecord::Base
+          self.abstract_class = true
+          connects_to database: { writing: :test_pool_1_db_a, reading: :test_pool_1_db_a_replica }
+        end
+
+        class Pool1DbA < AbstractPool1DbA
+          self.table_name = "tests"
+        end
+
+        class AbstractPool1DbB < ActiveRecord::Base
+          self.abstract_class = true
+          connects_to database: { writing: :test_pool_1_db_b }
+        end
+
+        class Pool1DbB < AbstractPool1DbB
+          self.table_name = "tests"
+        end
+
+        class AbstractPool2DbD < ActiveRecord::Base
+          self.abstract_class = true
+          connects_to database: { writing: :test_pool_2_db_d }
+        end
+
+        class Pool2DbD < AbstractPool2DbD
+          self.table_name = "tests"
+        end
+
+        class AbstractPool2DbE < ActiveRecord::Base
+          self.abstract_class = true
+          connects_to database: { writing: :test_pool_2_db_e }
+        end
+
+        class Pool2DbE < AbstractPool2DbE
+          self.table_name = "tests"
+        end
+
+        class AbstractPool3DbE < ActiveRecord::Base
+          self.abstract_class = true
+          connects_to database: { writing: :test_pool_3_db_e }
+        end
+
+        class Pool3DbE < AbstractPool3DbE
+          self.table_name = "tests"
+        end
+
+        # Test ARHP with Rails 6.1+ horizontal sharding functionality
+        class AbstractShardedModel < ActiveRecord::Base
+          self.abstract_class = true
+          connects_to shards: {
+                        default: { writing: :test_pool_1_db_shard_a },
+                        shard_b: { writing: :test_pool_1_db_shard_b, reading: :test_pool_1_db_shard_b_replica },
+                        shard_c: { writing: :test_pool_1_db_shard_c, reading: :test_pool_1_db_shard_c_replica },
+                        shard_d: { writing: :test_pool_2_db_shard_d, reading: :test_pool_2_db_shard_d_replica }
+                      }
+        end
+
+        class ShardedModel < AbstractShardedModel
+          self.table_name = "tests"
+        end
+      RUBY
+    else
+      eval <<-RUBY
+        # The placement of the Pool1DbC class is important so that its
+        # connection will not be the most recent connection established
+        # for test_pool_1.
+        class Pool1DbC < ActiveRecord::Base
+          establish_connection(:test_pool_1_db_c)
+        end
+
+        class Pool1DbA < ActiveRecord::Base
+          self.table_name = "tests"
+          establish_connection(:test_pool_1_db_a)
+        end
+
+        class Pool1DbAReplica < ActiveRecord::Base
+          self.table_name = "tests"
+          establish_connection(:test_pool_1_db_a_replica)
+        end
+
+        class Pool1DbB < ActiveRecord::Base
+          self.table_name =  "tests"
+          establish_connection(:test_pool_1_db_b)
+        end
+
+        class Pool2DbD < ActiveRecord::Base
+          self.table_name = "tests"
+          establish_connection(:test_pool_2_db_d)
+        end
+
+        class Pool2DbE < ActiveRecord::Base
+          self.table_name = "tests"
+          establish_connection(:test_pool_2_db_e)
+        end
+
+        class Pool3DbE < ActiveRecord::Base
+          self.table_name = "tests"
+          establish_connection(:test_pool_3_db_e)
+        end
+      RUBY
+    end
   end
 
   def current_database(klass)
     klass.connection.select_value('select DATABASE()')
   end
+
+  # Remove a method from a given module that fixes something.
+  # Execute the passed in block.
+  # Re-add the method back to the module.
+  def without_module_patch(mod, method_name)
+    method_body = mod.instance_method(method_name)
+    mod.remove_method(method_name)
+    yield if block_given?
+  ensure
+    mod.define_method(method_name, method_body)
+  end
+
+  def simulate_rails_app_active_record_railties
+    if ActiveRecord.version >= Gem::Version.new('6.0') && !RAILS_6_1_WITH_NON_LEGACY_CONNECTION_HANDLING
+      # Necessary for testing ActiveRecord 6.0 which uses the connection
+      # handlers when clearing query caches across all handlers when
+      # an operation that dirties the cache is involved (e.g. create/insert,
+      # update, delete/destroy, truncate, etc.)
+      # In Rails 6.1 this is only present when legacy_connection_handling=true
+      ActiveRecord::Base.connection_handlers = {
+        ActiveRecord::Base.writing_role => ActiveRecord::Base.default_connection_handler
+      }
+    end
+  end
 end

data/test/schema.rb

@@ -1,8 +1,23 @@
 # frozen_string_literal: true
 
 require_relative 'helper'
-ActiveRecord::Schema.define(version: 1) do
-  create_table 'tests' do |t|
-    t.string   'val'
+
+begin
+  ActiveRecordHostPool.allowing_writes = true
+
+  ActiveRecord::Schema.define(version: 1) do
+    create_table 'tests', force: true do |t|
+      t.string   'val'
+    end
+
+    # Add a table only the shard database will have. Conditional
+    # exists since Phenix loads the schema file for every database.
+    if ActiveRecord::Base.connection.current_database == 'arhp_test_db_c'
+      create_table 'pool1_db_cs' do |t|
+        t.string 'name'
+      end
+    end
   end
+ensure
+  ActiveRecordHostPool.allowing_writes = false
 end

data/test/test_arhp.rb

@@ -5,11 +5,17 @@ require_relative 'helper'
 class ActiveRecordHostPoolTest < Minitest::Test
   include ARHPTestSetup
   def setup
-    Phenix.rise!
+    if RAILS_6_1_WITH_NON_LEGACY_CONNECTION_HANDLING
+      Phenix.rise! config_path: 'test/three_tier_database.yml'
+    else
+      Phenix.rise!
+    end
     arhp_create_models
   end
 
   def teardown
+    ActiveRecord::Base.connection.disconnect!
+    ActiveRecordHostPool::PoolProxy.class_variable_set(:@@_connection_pools, {})
     Phenix.burn!
   end
 
@@ -19,7 +25,7 @@ class ActiveRecordHostPoolTest < Minitest::Test
 
     # Verify that when we fork, the process doesn't crash
     pid = Process.fork do
-      if ActiveRecord::VERSION::MAJOR == 5 && ActiveRecord::VERSION::MINOR == 2
+      if ActiveRecord.version >= Gem::Version.new('5.2')
         assert_equal(false, ActiveRecord::Base.connected?) # New to Rails 5.2
       else
         assert_equal(true, ActiveRecord::Base.connected?)
@@ -30,78 +36,80 @@ class ActiveRecordHostPoolTest < Minitest::Test
     ActiveRecord::Base.connection_handler.clear_all_connections!
   end
 
-  def test_models_with_matching_hosts_should_share_a_connection
-    assert_equal(Test1.connection.raw_connection, Test2.connection.raw_connection)
-    assert_equal(Test3.connection.raw_connection, Test4.connection.raw_connection)
-  end
-
-  def test_models_without_matching_hosts_should_not_share_a_connection
-    refute_equal(Test1.connection.raw_connection, Test4.connection.raw_connection)
+  def test_models_with_matching_hosts_ports_sockets_usernames_and_replica_status_should_share_a_connection
+    assert_equal(Pool1DbA.connection.raw_connection, Pool1DbB.connection.raw_connection)
+    assert_equal(Pool2DbD.connection.raw_connection, Pool2DbE.connection.raw_connection)
   end
 
-  def test_models_without_matching_usernames_should_not_share_a_connection
-    refute_equal(Test4.connection.raw_connection, Test5.connection.raw_connection)
+  def test_models_with_different_ports_should_not_share_a_connection
+    refute_equal(Pool1DbA.connection.raw_connection, Pool2DbD.connection.raw_connection)
   end
 
-  def test_models_without_match_slave_status_should_not_share_a_connection
-    refute_equal(Test1.connection.raw_connection, Test1Slave.connection.raw_connection)
+  def test_models_with_different_usernames_should_not_share_a_connection
+    refute_equal(Pool2DbE.connection.raw_connection, Pool3DbE.connection.raw_connection)
   end
 
   def test_should_select_on_correct_database
-    assert_action_uses_correct_database(:select_all, 'select 1')
+    Pool1DbA.connection.send(:select_all, 'select 1')
+    assert_equal 'arhp_test_db_a', current_database(Pool1DbA)
+
+    Pool2DbD.connection.send(:select_all, 'select 1')
+    assert_equal 'arhp_test_db_d', current_database(Pool2DbD)
+
+    Pool3DbE.connection.send(:select_all, 'select 1')
+    assert_equal 'arhp_test_db_e', current_database(Pool3DbE)
   end
 
   def test_should_insert_on_correct_database
-    assert_action_uses_correct_database(:insert, "insert into tests values(NULL, 'foo')")
+    Pool1DbA.connection.send(:insert, "insert into tests values(NULL, 'foo')")
+    assert_equal 'arhp_test_db_a', current_database(Pool1DbA)
+
+    Pool2DbD.connection.send(:insert, "insert into tests values(NULL, 'foo')")
+    assert_equal 'arhp_test_db_d', current_database(Pool2DbD)
+
+    Pool3DbE.connection.send(:insert, "insert into tests values(NULL, 'foo')")
+    assert_equal 'arhp_test_db_e', current_database(Pool3DbE)
   end
 
   def test_connection_returns_a_proxy
-    assert_kind_of ActiveRecordHostPool::ConnectionProxy, Test1.connection
+    assert_kind_of ActiveRecordHostPool::ConnectionProxy, Pool1DbA.connection
   end
 
   def test_connection_proxy_handles_private_methods
     # Relies on connection.class returning the real class
-    Test1.connection.class.class_eval do
+    Pool1DbA.connection.class.class_eval do
       private
 
       def test_private_method
         true
       end
     end
-    assert Test1.connection.respond_to?(:test_private_method, true)
-    refute Test1.connection.respond_to?(:test_private_method)
-    assert_includes(Test1.connection.private_methods, :test_private_method)
-    assert_equal true, Test1.connection.send(:test_private_method)
-  end
-
-  def test_should_not_share_a_query_cache
-    Test1.create(val: 'foo')
-    Test2.create(val: 'foobar')
-    Test1.connection.cache do
-      refute_equal Test1.first.val, Test2.first.val
-    end
+    assert Pool1DbA.connection.respond_to?(:test_private_method, true)
+    refute Pool1DbA.connection.respond_to?(:test_private_method)
+    assert_includes(Pool1DbA.connection.private_methods, :test_private_method)
+    assert_equal true, Pool1DbA.connection.send(:test_private_method)
   end
 
   def test_object_creation
-    Test1.create(val: 'foo')
-    assert_equal('arhp_test_1', current_database(Test1))
+    Pool1DbA.create(val: 'foo')
+    assert_equal('arhp_test_db_a', current_database(Pool1DbA))
 
-    Test3.create(val: 'bar')
-    assert_equal('arhp_test_1', current_database(Test1))
-    assert_equal('arhp_test_3', current_database(Test3))
+    Pool2DbD.create(val: 'bar')
+    assert_equal('arhp_test_db_a', current_database(Pool1DbA))
+    assert_equal('arhp_test_db_d', current_database(Pool2DbD))
 
-    Test2.create!(val: 'bar_distinct')
-    assert_equal('arhp_test_2', current_database(Test2))
-    assert Test2.find_by_val('bar_distinct')
-    refute Test1.find_by_val('bar_distinct')
+    Pool1DbB.create!(val: 'bar_distinct')
+    assert_equal('arhp_test_db_b', current_database(Pool1DbB))
+    assert Pool1DbB.find_by_val('bar_distinct')
+    refute Pool1DbA.find_by_val('bar_distinct')
   end
 
   def test_disconnect
-    Test1.create(val: 'foo')
-    unproxied = Test1.connection.unproxied
-    Test1.connection_handler.clear_all_connections!
-    Test1.create(val: 'foo')
-    assert(unproxied != Test1.connection.unproxied)
+    Pool1DbA.create(val: 'foo')
+    unproxied = Pool1DbA.connection.unproxied
+    Pool1DbA.connection_handler.clear_all_connections!
+    Pool1DbA.create(val: 'foo')
+    assert(unproxied != Pool1DbA.connection.unproxied)
   end
 
   def test_checkout
@@ -113,7 +121,7 @@ class ActiveRecordHostPoolTest < Minitest::Test
   end
 
   def test_no_switch_when_creating_db
-    conn = Test1.connection
+    conn = Pool1DbA.connection
     conn.expects(:execute_without_switching)
     conn.expects(:_switch_connection).never
     assert conn._host_pool_current_database
@@ -121,7 +129,7 @@ class ActiveRecordHostPoolTest < Minitest::Test
   end
 
   def test_no_switch_when_dropping_db
-    conn = Test1.connection
+    conn = Pool1DbA.connection
     conn.expects(:execute_without_switching)
     conn.expects(:_switch_connection).never
     assert conn._host_pool_current_database
@@ -131,17 +139,17 @@ class ActiveRecordHostPoolTest < Minitest::Test
   def test_underlying_assumption_about_test_db
     debug_me = false
     # ensure connection
-    Test1.first
+    Pool1DbA.first
 
     # which is the "default" DB to connect to?
-    first_db = Test1.connection.unproxied.instance_variable_get(:@_cached_current_database)
+    first_db = Pool1DbA.connection.unproxied.instance_variable_get(:@_cached_current_database)
     puts "\nOk, we started on #{first_db}" if debug_me
 
     switch_to_klass = case first_db
-    when 'arhp_test_2'
-      Test1
-    when 'arhp_test_1'
-      Test2
+    when 'arhp_test_db_b'
+      Pool1DbA
+    when 'arhp_test_db_a'
+      Pool1DbB
     else
       raise "Expected a database name, got #{first_db.inspect}"
     end
@@ -156,7 +164,7 @@ class ActiveRecordHostPoolTest < Minitest::Test
 
     # now, disable our auto-switching and trigger a mysql reconnect
     switch_to_klass.connection.unproxied.stubs(:_switch_connection).returns(true)
-    Test3.connection.execute("KILL #{thread_id}")
+    Pool2DbD.connection.execute("KILL #{thread_id}")
 
     # and finally, did mysql reconnect correctly?
     puts "\nAnd now we end up on #{current_database(switch_to_klass)}" if debug_me
@@ -169,15 +177,4 @@ class ActiveRecordHostPoolTest < Minitest::Test
     pool.expects(:checkin).with(conn)
     pool.release_connection
   end
-
-  private
-
-  def assert_action_uses_correct_database(action, sql)
-    (1..4).each do |i|
-      klass = ARHPTestSetup.const_get("Test#{i}")
-      desired_db = "arhp_test_#{i}"
-      klass.connection.send(action, sql)
-      assert_equal desired_db, current_database(klass)
-    end
-  end
 end

metadata

@@ -1,17 +1,17 @@
 --- !ruby/object:Gem::Specification
 name: active_record_host_pool
 version: !ruby/object:Gem::Version
-  version: 1.0.1
+  version: 1.1.0
 platform: ruby
 authors:
 - Benjamin Quorning
 - Gabe Martin-Dempesy
 - Pierre Schambacher
 - Ben Osheroff
-autorequire: 
+autorequire:
 bindir: bin
 cert_chain: []
-date: 2020-03-30 00:00:00.000000000 Z
+date: 2022-08-26 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -19,20 +19,20 @@ dependencies:
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: 4.2.0
+        version: 5.1.0
     - - "<"
       - !ruby/object:Gem::Version
-        version: '6.1'
+        version: '7.0'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: 4.2.0
+        version: 5.1.0
     - - "<"
       - !ruby/object:Gem::Version
-        version: '6.1'
+        version: '7.0'
 - !ruby/object:Gem::Dependency
   name: mysql2
   requirement: !ruby/object:Gem::Requirement
@@ -61,34 +61,48 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: minitest
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 5.10.0
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: 5.10.0
 - !ruby/object:Gem::Dependency
   name: mocha
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0'
+        version: 1.4.0
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0'
+        version: 1.4.0
 - !ruby/object:Gem::Dependency
   name: phenix
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0'
+        version: 1.0.1
   type: :development
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0'
+        version: 1.0.1
 - !ruby/object:Gem::Dependency
   name: rake
   requirement: !ruby/object:Gem::Requirement
@@ -145,6 +159,7 @@ files:
 - MIT-LICENSE
 - Readme.md
 - lib/active_record_host_pool.rb
+- lib/active_record_host_pool/clear_query_cache_patch.rb
 - lib/active_record_host_pool/connection_adapter_mixin.rb
 - lib/active_record_host_pool/connection_proxy.rb
 - lib/active_record_host_pool/pool_proxy.rb
@@ -157,7 +172,7 @@ homepage: https://github.com/zendesk/active_record_host_pool
 licenses:
 - MIT
 metadata: {}
-post_install_message: 
+post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -165,15 +180,15 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '0'
+      version: 2.6.0
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.1.1
-signing_key: 
+rubygems_version: 3.1.6
+signing_key:
 specification_version: 4
 summary: Allow ActiveRecord to share a connection to multiple databases on the same
   host