From 31021a8c85bb80300deb01db96876858ff417f4a Mon Sep 17 00:00:00 2001 From: Eileen Uchitelle Date: Fri, 28 Sep 2018 16:52:54 -0400 Subject: Basic API for connection switching This PR adds the ability to 1) connect to multiple databases in a model, and 2) switch between those connections using a block. To connect a model to a set of databases for writing and reading use the following API. This API supercedes `establish_connection`. The `writing` and `reading` keys represent handler / role names and `animals` and `animals_replica` represents the database key to look up the configuration hash from. ``` class AnimalsBase < ApplicationRecord connects_to database: { writing: :animals, reading: :animals_replica } end ``` Inside the application - outside the model declaration - we can switch connections with a block call to `connected_to`. If we want to connect to a db that isn't default (ie readonly_slow) we can connect like this: Outside the model we may want to connect to a new database (one that is not in the default writing/reading set) - for example a slow replica for making slow queries. To do this we have the `connected_to` method that takes a `database` hash that matches the signature of `connects_to`. The `connected_to` method also takes a block. ``` AcitveRecord::Base.connected_to(database: { slow_readonly: :primary_replica_slow }) do ModelInPrimary.do_something_thats_slow end ``` For models that are already loaded and connections that are already connected, `connected_to` doesn't need to pass in a `database` because you may want to run queries against multiple databases using a specific role/handler. In this case `connected_to` can take a `role` and use that to swap on the connection passed. This simplies queries - and matches how we do it in GitHub. Once you're connected to the database you don't need to re-connect, we assume the connection is in the pool and simply pass the handler we'd like to swap on. ``` ActiveRecord::Base.connected_to(role: :reading) do Dog.read_something_from_dog ModelInPrimary.do_something_from_model_in_primary end ``` --- .../lib/active_record/connection_handling.rb | 97 +++++++++++++++++++++- activerecord/lib/active_record/core.rb | 3 + 2 files changed, 99 insertions(+), 1 deletion(-) (limited to 'activerecord/lib') 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 -- cgit v1.2.3