Eagerly establish test transactions for shards

activerecord/lib/active_record/connection_adapters/abstract/connection_pool.rb

@@ -1040,6 +1040,7 @@ def establish_connection(config, owner_name: Base.name, shard: Base.default_shar
         payload = {}
         if pool_config
           payload[:spec_name] = pool_config.connection_specification_name
+          payload[:shard] = shard
           payload[:config] = db_config.configuration_hash
         end
 

activerecord/lib/active_record/test_fixtures.rb

@@ -131,11 +131,12 @@ def setup_fixtures(config = ActiveRecord::Base)
         # When connections are established in the future, begin a transaction too
         @connection_subscriber = ActiveSupport::Notifications.subscribe("!connection.active_record") do |_, _, _, _, payload|
           spec_name = payload[:spec_name] if payload.key?(:spec_name)
+          shard = payload[:shard] if payload.key?(:shard)
           setup_shared_connection_pool
 
           if spec_name
             begin
-              connection = ActiveRecord::Base.connection_handler.retrieve_connection(spec_name)
+              connection = ActiveRecord::Base.connection_handler.retrieve_connection(spec_name, shard: shard)
             rescue ConnectionNotEstablished
               connection = nil
             end

activerecord/test/cases/connection_pool_test.rb

@@ -492,6 +492,7 @@ def test_anonymous_class_exception
       end
 
       class ConnectionTestModel < ActiveRecord::Base
+        self.abstract_class = true
       end
 
       def test_connection_notification_is_called
@@ -501,8 +502,23 @@ def test_connection_notification_is_called
         end
         ConnectionTestModel.establish_connection :arunit
 
-        assert_equal [:config, :spec_name], payloads[0].keys.sort
+        assert_equal [:config, :shard, :spec_name], payloads[0].keys.sort
         assert_equal "ActiveRecord::ConnectionAdapters::ConnectionPoolTest::ConnectionTestModel", payloads[0][:spec_name]
+        assert_equal ActiveRecord::Base.default_shard, payloads[0][:shard]
+      ensure
+        ActiveSupport::Notifications.unsubscribe(subscription) if subscription
+      end
+
+      def test_connection_notification_is_called_for_shard
+        payloads = []
+        subscription = ActiveSupport::Notifications.subscribe("!connection.active_record") do |name, started, finished, unique_id, payload|
+          payloads << payload
+        end
+        ConnectionTestModel.connects_to shards: { shard_two: { writing: :arunit } }
+
+        assert_equal [:config, :shard, :spec_name], payloads[0].keys.sort
+        assert_equal "ActiveRecord::ConnectionAdapters::ConnectionPoolTest::ConnectionTestModel", payloads[0][:spec_name]
+        assert_equal :shard_two, payloads[0][:shard]
       ensure
         ActiveSupport::Notifications.unsubscribe(subscription) if subscription
       end

activerecord/test/cases/fixtures_test.rb

@@ -951,12 +951,31 @@ def lock_thread=(lock_thread); end
     assert(connection.rollback_transaction_called, "Expected <mock connection>#rollback_transaction to be called but was not")
   end
 
+  def test_transaction_created_on_connection_notification_for_shard
+    connection = Class.new do
+      attr_accessor :pool
+
+      def transaction_open?; end
+      def begin_transaction(*args); end
+      def rollback_transaction(*args); end
+    end.new
+
+    connection.pool = Class.new do
+      def lock_thread=(lock_thread); end
+    end.new
+
+    assert_called_with(connection, :begin_transaction, [joinable: false, _lazy: false]) do
+      fire_connection_notification(connection, shard: :shard_two)
+    end
+  end
+
   private
-    def fire_connection_notification(connection)
-      assert_called_with(ActiveRecord::Base.connection_handler, :retrieve_connection, ["book"], returns: connection) do
+    def fire_connection_notification(connection, shard: ActiveRecord::Base.default_shard)
+      assert_called_with(ActiveRecord::Base.connection_handler, :retrieve_connection, ["book", { shard: shard }], returns: connection) do
         message_bus = ActiveSupport::Notifications.instrumenter
         payload = {
           spec_name: "book",
+          shard: shard,
           config: nil,
         }