Merge pull request #34052 from eileencodes/connection-switching

Part 4: Multi db improvements, Basic API for connection switching

2 files changed, 99 insertions, 1 deletions

diff --git a/activerecord/lib/active_record/connection_handling.rb b/activerecord/lib/active_record/connection_handling.rb
index 18114f9e1c..5141271db9 100644
--- a/activerecord/lib/active_record/connection_handling.rb
+++ b/activerecord/lib/active_record/connection_handling.rb
@@ -47,6 +47,92 @@ module ActiveRecord
     # The exceptions AdapterNotSpecified, AdapterNotFound and +ArgumentError+
     # may be returned on an error.
     def establish_connection(config_or_env = nil)
+      config_hash = resolve_config_for_connection(config_or_env)
+      connection_handler.establish_connection(config_hash)
+    end
+
+    # Connects a model to the databases specified. The +database+ keyword
+    # takes a hash consisting of a +role+ and a +database_key+.
+    #
+    # This will create a connection handler for switching between connections,
+    # look up the config hash using the +database_key+ and finally
+    # establishes a connection to that config.
+    #
+    #   class AnimalsModel < ApplicationRecord
+    #     self.abstract_class = true
+    #
+    #     connects_to database: { writing: :primary, reading: :primary_replica }
+    #   end
+    #
+    # Returns an array of established connections.
+    def connects_to(database: {})
+      connections = []
+
+      database.each do |role, database_key|
+        config_hash = resolve_config_for_connection(database_key)
+        handler = lookup_connection_handler(role.to_sym)
+
+        connections << handler.establish_connection(config_hash)
+      end
+
+      connections
+    end
+
+    # Connects to a database or role (ex writing, reading, or another
+    # custom role) for the duration of the block.
+    #
+    # If a role is passed, Active Record will look up the connection
+    # based on the requested role:
+    #
+    #   ActiveRecord::Base.connected_to(role: :writing) do
+    #     Dog.create! # creates dog using dog connection
+    #   end
+    #
+    #   ActiveRecord::Base.connected_to(role: :reading) do
+    #     Dog.create! # throws exception because we're on a replica
+    #   end
+    #
+    #   ActiveRecord::Base.connected_to(role: :unknown_ode) do
+    #     # raises exception due to non-existent role
+    #   end
+    #
+    # For cases where you may want to connect to a database outside of the model,
+    # you can use +connected_to+ with a +database+ argument. The +database+ argument
+    # expects a symbol that corresponds to the database key in your config.
+    #
+    # This will connect to a new database for the queries inside the block.
+    #
+    #   ActiveRecord::Base.connected_to(database: :animals_slow_replica) do
+    #     Dog.run_a_long_query # runs a long query while connected to the +animals_slow_replica+
+    #   end
+    def connected_to(database: nil, role: nil, &blk)
+      if database && role
+        raise ArgumentError, "connected_to can only accept a database or role argument, but not both arguments."
+      elsif database
+        config_hash = resolve_config_for_connection(database)
+        handler = lookup_connection_handler(database.to_sym)
+
+        with_handler(database.to_sym) do
+          handler.establish_connection(config_hash)
+          return yield
+        end
+      elsif role
+        with_handler(role.to_sym, &blk)
+      else
+        raise ArgumentError, "must provide a `database` or a `role`."
+      end
+    end
+
+    def lookup_connection_handler(handler_key) # :nodoc:
+      connection_handlers[handler_key] ||= ActiveRecord::ConnectionAdapters::ConnectionHandler.new
+    end
+
+    def with_handler(handler_key, &blk) # :nodoc:
+      handler = lookup_connection_handler(handler_key)
+      swap_connection_handler(handler, &blk)
+    end
+
+    def resolve_config_for_connection(config_or_env) # :nodoc:
       raise "Anonymous class is not allowed." unless name
 
       config_or_env ||= DEFAULT_ENV.call.to_sym
@@ -57,7 +143,7 @@ module ActiveRecord
       config_hash = resolver.resolve(config_or_env, pool_name).symbolize_keys
       config_hash[:name] = pool_name
 
-      connection_handler.establish_connection(config_hash)
+      config_hash
     end
 
     # Returns the connection currently associated with the class. This can
@@ -118,5 +204,14 @@ module ActiveRecord
 
     delegate :clear_active_connections!, :clear_reloadable_connections!,
       :clear_all_connections!, :flush_idle_connections!, to: :connection_handler
+
+    private
+
+      def swap_connection_handler(handler, &blk) # :nodoc:
+        old_handler, ActiveRecord::Base.connection_handler = ActiveRecord::Base.connection_handler, handler
+        yield
+      ensure
+        ActiveRecord::Base.connection_handler = old_handler
+      end
   end
 end
diff --git a/activerecord/lib/active_record/core.rb b/activerecord/lib/active_record/core.rb
index 392602bc0f..b53681ee20 100644
--- a/activerecord/lib/active_record/core.rb
+++ b/activerecord/lib/active_record/core.rb
@@ -124,6 +124,8 @@ module ActiveRecord
 
       mattr_accessor :belongs_to_required_by_default, instance_accessor: false
 
+      mattr_accessor :connection_handlers, instance_accessor: false, default: {}
+
       class_attribute :default_connection_handler, instance_writer: false
 
       self.filter_attributes = []
@@ -137,6 +139,7 @@ module ActiveRecord
       end
 
       self.default_connection_handler = ConnectionAdapters::ConnectionHandler.new
+      self.connection_handlers = { writing: ActiveRecord::Base.default_connection_handler }
     end
 
     module ClassMethods