switchman 3.0.2 → 4.2.5

data/lib/switchman/database_server.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require 'securerandom'
+require "securerandom"
 
 module Switchman
   class DatabaseServer
@@ -8,24 +8,25 @@ module Switchman
 
     class << self
       attr_accessor :creating_new_shard
+      attr_reader :all_roles
+
+      include Enumerable
+
+      delegate :each, to: :all
 
       def all
         database_servers.values
       end
 
-      def all_roles
-        @all_roles ||= all.map(&:roles).flatten.uniq
-      end
-
       def find(id_or_all)
         return all if id_or_all == :all
-        return id_or_all.map { |id| database_servers[id || ::Rails.env] }.compact.uniq if id_or_all.is_a?(Array)
+        return id_or_all.filter_map { |id| database_servers[id || ::Rails.env] }.uniq if id_or_all.is_a?(Array)
 
         database_servers[id_or_all || ::Rails.env]
       end
 
       def create(settings = {})
-        raise 'database servers should be set up in database.yml' unless ::Rails.env.test?
+        raise "database servers should be set up in database.yml" unless ::Rails.env.test?
 
         id = settings[:id]
         unless id
@@ -38,7 +39,7 @@ module Switchman
         database_servers[server.id] = server
         ::ActiveRecord::Base.configurations.configurations <<
           ::ActiveRecord::DatabaseConfigurations::HashConfig.new(::Rails.env, "#{server.id}/primary", settings)
-        Shard.send(:initialize_sharding)
+        Shard.send(:configure_connects_to)
         server
       end
 
@@ -49,31 +50,48 @@ module Switchman
         servers[rand(servers.length)]
       end
 
+      def guard_servers
+        all.each { |db| db.guard! if db.config[:prefer_secondary] }
+      end
+
+      def regions
+        @regions ||= all.filter_map(&:region).uniq.sort
+      end
+
       private
 
       def reference_role(role)
         return if all_roles.include?(role)
 
         @all_roles << role
-        Shard.send(:initialize_sharding)
+        Shard.send(:configure_connects_to)
       end
 
       def database_servers
-        unless @database_servers
+        if !@database_servers || @database_servers.empty?
           @database_servers = {}.with_indifferent_access
+          roles = []
           ::ActiveRecord::Base.configurations.configurations.each do |config|
-            if config.name.include?('/')
-              name, role = config.name.split('/')
+            if config.name.include?("/")
+              name, role = config.name.split("/")
             else
               name, role = config.env_name, config.name
             end
+            role = role.to_sym
 
-            if role == 'primary'
+            roles << role
+            if role == :primary
               @database_servers[name] = DatabaseServer.new(config.env_name, config.configuration_hash)
             else
               @database_servers[name].roles << role
             end
           end
+          # Do this after so that all database servers for all roles are established and we won't prematurely
+          # configure a connection for the wrong role
+          @all_roles = roles.uniq
+          return @database_servers if @database_servers.empty?
+
+          Shard.send(:configure_connects_to)
         end
         @database_servers
       end
@@ -89,21 +107,22 @@ module Switchman
     end
 
     def connects_to_hash
-      self.class.all_roles.map do |role|
+      self.class.all_roles.to_h do |role|
         config_role = role
         config_role = :primary unless roles.include?(role)
         config_name = :"#{id}/#{config_role}"
         config_name = :primary if id == ::Rails.env && config_role == :primary
         [role.to_sym, config_name]
-      end.to_h
+      end
     end
 
     def destroy
       self.class.send(:database_servers).delete(id) if id
       Shard.sharded_models.each do |klass|
         self.class.all_roles.each do |role|
-          klass.connection_handler.remove_connection_pool(klass.connection_specification_name, role: role,
-                                                                                               shard: id.to_sym)
+          klass.connection_handler.remove_connection_pool(klass.connection_specification_name,
+                                                          role:,
+                                                          shard: id.to_sym)
         end
       end
     end
@@ -129,32 +148,57 @@ module Switchman
         end
     end
 
-    def guard_rail_environment
-      @guard_rail_environment || ::GuardRail.environment
+    def region
+      config[:region]
+    end
+
+    # @param region [String, Array<String>] the region(s) to check against
+    # @return true if the database server doesn't have a region, or it
+    #   matches the specified region
+    def in_region?(region)
+      !self.region || (region.is_a?(Array) ? region.include?(self.region) : self.region == region)
+    end
+
+    # @return true if the database server doesn't have a region, Switchman is
+    #   not configured with a region, or the database server's region matches
+    #   Switchman's current region
+    def in_current_region?
+      unless instance_variable_defined?(:@in_current_region)
+        @in_current_region = !region ||
+                             !Switchman.region ||
+                             region == Switchman.region
+      end
+      @in_current_region
     end
 
     # locks this db to a specific environment, except for
     # when doing writes (then it falls back to the current
     # value of GuardRail.environment)
     def guard!(environment = :secondary)
-      @guard_rail_environment = environment
+      DatabaseServer.send(:reference_role, environment)
+      ::ActiveRecord::Base.connected_to_stack << { shard_roles: { id.to_sym => environment },
+                                                   klasses: [::ActiveRecord::Base] }
     end
 
     def unguard!
-      @guard_rail_environment = nil
+      ::ActiveRecord::Base.connected_to_stack << { shard_roles: { id.to_sym => :_switchman_inherit },
+                                                   klasses: [::ActiveRecord::Base] }
     end
 
     def unguard
-      old_env = @guard_rail_environment
-      unguard!
-      yield
-    ensure
-      guard!(old_env)
+      return yield unless ::ActiveRecord::Base.role_overriden?(id.to_sym)
+
+      begin
+        unguard!
+        yield
+      ensure
+        ::ActiveRecord::Base.connected_to_stack.pop
+      end
     end
 
     def shards
       if id == ::Rails.env
-        Shard.where('database_server_id IS NULL OR database_server_id=?', id)
+        Shard.where("database_server_id IS NULL OR database_server_id=?", id)
       else
         Shard.where(database_server_id: id)
       end
@@ -174,70 +218,66 @@ module Switchman
       if config_create_statement
         create_commands = Array(config_create_statement).dup
         create_statement = lambda {
-          create_commands.map { |statement| format(statement, name: name, password: password) }
+          create_commands.map { |statement| format(statement, name:, password:) }
         }
       end
 
       id ||= begin
-        id_seq = Shard.connection.quote(Shard.connection.quote_table_name('switchman_shards_id_seq'))
+        id_seq = Shard.connection.quote(Shard.connection.quote_table_name("switchman_shards_id_seq"))
         next_id = Shard.connection.select_value("SELECT nextval(#{id_seq})")
         next_id.to_i
       end
 
       name ||= "#{config[:database]}_shard_#{id}"
 
+      schema_already_existed = false
+      shard = nil
       Shard.connection.transaction do
-        shard = Shard.create!(id: id,
-                              name: name,
-                              database_server_id: self.id)
-        schema_already_existed = false
-
-        begin
-          self.class.creating_new_shard = true
-          DatabaseServer.send(:reference_role, :deploy)
-          ::ActiveRecord::Base.connected_to(shard: self.id.to_sym, role: :deploy) do
-            if create_statement
-              if ::ActiveRecord::Base.connection.select_value("SELECT 1 FROM pg_namespace WHERE nspname=#{::ActiveRecord::Base.connection.quote(name)}")
-                schema_already_existed = true
-                raise 'This schema already exists; cannot overwrite'
-              end
-              Array(create_statement.call).each do |stmt|
-                ::ActiveRecord::Base.connection.execute(stmt)
-              end
+        self.class.creating_new_shard = true
+        DatabaseServer.send(:reference_role, :deploy)
+        ::ActiveRecord::Base.connected_to(shard: self.id.to_sym, role: :deploy) do
+          shard = Shard.create!(id:,
+                                name:,
+                                database_server_id: self.id)
+          if create_statement
+            if ::ActiveRecord::Base.connection.select_value(
+              "SELECT 1 FROM pg_namespace WHERE nspname=#{::ActiveRecord::Base.connection.quote(name)}"
+            )
+              schema_already_existed = true
+              raise "This schema already exists; cannot overwrite"
             end
-            if config[:adapter] == 'postgresql'
-              old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor do
-              end
+            Array(create_statement.call).each do |stmt|
+              ::ActiveRecord::Base.connection.execute(stmt)
             end
-            old_verbose = ::ActiveRecord::Migration.verbose
-            ::ActiveRecord::Migration.verbose = false
-
-            unless schema == false
-              shard.activate(*Shard.sharded_models) do
-                reset_column_information
-
-                ::ActiveRecord::Base.connection.transaction(requires_new: true) do
-                  ::ActiveRecord::Base.connection.migration_context.migrate
-                end
-                reset_column_information
-                ::ActiveRecord::Base.descendants.reject do |m|
-                  m <= UnshardedRecord || !m.table_exists?
-                end.each(&:define_attribute_methods)
+          end
+          if config[:adapter] == "postgresql"
+            old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor {}
+          end
+          old_verbose = ::ActiveRecord::Migration.verbose
+          ::ActiveRecord::Migration.verbose = false
+
+          unless schema == false
+            shard.activate do
+              ::ActiveRecord::Base.connection.transaction(requires_new: true) do
+                ::ActiveRecord::MigrationContext.new(::ActiveRecord::Migrator.migrations_paths).migrate
               end
+
+              ::ActiveRecord::Base.descendants.reject do |m|
+                m <= UnshardedRecord || !m.table_exists?
+              end.each(&:define_attribute_methods)
             end
-          ensure
-            ::ActiveRecord::Migration.verbose = old_verbose
-            ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
           end
-          shard
-        rescue
-          shard.destroy
-          shard.drop_database rescue nil unless schema_already_existed
-          reset_column_information unless schema == false rescue nil
-          raise
         ensure
-          self.class.creating_new_shard = false
+          ::ActiveRecord::Migration.verbose = old_verbose
+          ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
         end
+        shard
+      rescue
+        shard&.destroy
+        shard&.drop_database rescue nil unless schema_already_existed
+        raise
+      ensure
+        self.class.creating_new_shard = false
       end
     end
 
@@ -259,18 +299,18 @@ module Switchman
     end
 
     def primary_shard
-      unless instance_variable_defined?(:@primary_shard)
-        # if sharding isn't fully set up yet, we may not be able to query the shards table
-        @primary_shard = Shard.default if Shard.default.database_server == self
-        @primary_shard ||= shards.where(name: nil).first
-      end
-      @primary_shard
-    end
+      return nil unless primary_shard_id
 
-    private
+      Shard.lookup(primary_shard_id)
+    end
 
-    def reset_column_information
-      ::ActiveRecord::Base.descendants.reject { |m| m <= UnshardedRecord }.each(&:reset_column_information)
+    def primary_shard_id
+      unless instance_variable_defined?(:@primary_shard_id)
+        # if sharding isn't fully set up yet, we may not be able to query the shards table
+        @primary_shard_id = Shard.default.id if Shard.default.database_server == self
+        @primary_shard_id ||= shards.where(name: nil).first&.id
+      end
+      @primary_shard_id
     end
   end
 end

data/lib/switchman/default_shard.rb

@@ -1,13 +1,12 @@
 # frozen_string_literal: true
 
-require 'switchman/database_server'
-
 module Switchman
   class DefaultShard
     def id
-      'default'
+      "default"
     end
-    alias cache_key id
+    alias_method :cache_key, :id
+
     def activate(*_classes)
       yield
     end
@@ -59,8 +58,18 @@ module Switchman
       self
     end
 
+    def region; end
+
+    def in_region?(_region)
+      true
+    end
+
+    def in_current_region?
+      true
+    end
+
     def _dump(_depth)
-      ''
+      ""
     end
 
     def self._load(_str)

data/lib/switchman/engine.rb

@@ -4,105 +4,22 @@ module Switchman
   class Engine < ::Rails::Engine
     isolate_namespace Switchman
 
-    # enable Rails 6.1 style connection handling
-    config.active_record.legacy_connection_handling = false
     config.active_record.writing_role = :primary
 
-    config.autoload_once_paths << File.expand_path('app/models', config.paths.path)
+    ::GuardRail.singleton_class.prepend(GuardRail::ClassMethods)
 
-    def self.lookup_stores(cache_store_config)
-      result = {}
-      cache_store_config.each do |key, value|
-        next if value.is_a?(String)
-
-        result[key] = ::ActiveSupport::Cache.lookup_store(value)
-      end
-
-      cache_store_config.each do |key, value| # rubocop:disable Style/CombinableLoops
-        next unless value.is_a?(String)
-
-        result[key] = result[value]
-      end
-      result
-    end
-
-    initializer 'switchman.initialize_cache', before: 'initialize_cache' do
-      require 'switchman/active_support/cache'
-      ::ActiveSupport::Cache.singleton_class.prepend(ActiveSupport::Cache::ClassMethods)
-
-      # if we haven't already setup our cache map out-of-band, set it up from
-      # config.cache_store now. behaves similarly to Rails' default
-      # initialize_cache initializer, but for each value in the map, rather
-      # than just Rails.cache. if config.cache_store is a flat value, uses it
-      # to fill just the Rails.env entry in the cache map.
-      unless Switchman.config[:cache_map].present?
-        cache_store_config = ::Rails.configuration.cache_store
-        cache_store_config = { ::Rails.env => cache_store_config } unless cache_store_config.is_a?(Hash)
-
-        Switchman.config[:cache_map] = Engine.lookup_stores(cache_store_config)
-      end
-
-      # if the configured cache map (either from before, or as populated from
-      # config.cache_store) didn't have an entry for Rails.env, add one using
-      # lookup_store(nil); matches the behavior of Rails' default
-      # initialize_cache initializer when config.cache_store is nil.
-      unless Switchman.config[:cache_map].key?(::Rails.env)
-        value = ::ActiveSupport::Cache.lookup_store(nil)
-        Switchman.config[:cache_map][::Rails.env] = value
-      end
-
-      middlewares = Switchman.config[:cache_map].values.map do |store|
-        store.middleware if store.respond_to?(:middleware)
-      end.compact.uniq
-      middlewares.each do |middleware|
-        config.middleware.insert_before('Rack::Runtime', middleware)
-      end
-
-      # prevent :initialize_cache from trying to (or needing to) set
-      # Rails.cache. once our switchman.extend_ar initializer (below) runs
-      # Rails.cache will be overridden to pull appropriate values from the
-      # cache map, but between now and then, Rails.cache should return the
-      # Rails.env entry in the cache map.
-      ::Rails.cache = Switchman.config[:cache_map][::Rails.env]
-    end
-
-    initializer 'switchman.extend_ar', before: 'active_record.initialize_database' do
+    initializer "switchman.active_record_patch",
+                before: "active_record.initialize_database",
+                after: :setup_once_autoloader do
       ::ActiveSupport.on_load(:active_record) do
-        require 'switchman/active_record/abstract_adapter'
-        require 'switchman/active_record/association'
-        require 'switchman/active_record/attribute_methods'
-        require 'switchman/active_record/base'
-        require 'switchman/active_record/calculations'
-        require 'switchman/active_record/connection_pool'
-        require 'switchman/active_record/database_configurations'
-        require 'switchman/active_record/database_configurations/database_config'
-        require 'switchman/active_record/finder_methods'
-        require 'switchman/active_record/log_subscriber'
-        require 'switchman/active_record/migration'
-        require 'switchman/active_record/model_schema'
-        require 'switchman/active_record/persistence'
-        require 'switchman/active_record/predicate_builder'
-        require 'switchman/active_record/query_cache'
-        require 'switchman/active_record/query_methods'
-        require 'switchman/active_record/reflection'
-        require 'switchman/active_record/relation'
-        require 'switchman/active_record/spawn_methods'
-        require 'switchman/active_record/statement_cache'
-        require 'switchman/active_record/tasks/database_tasks'
-        require 'switchman/active_record/type_caster'
-        require 'switchman/arel'
-        require 'switchman/call_super'
-        require 'switchman/rails'
-        require 'switchman/guard_rail/relation'
-        require 'switchman/standard_error'
-
-        ::StandardError.include(StandardError)
+        # Switchman requires postgres, so just always load the pg adapter
+        require "active_record/connection_adapters/postgresql_adapter"
 
         self.default_shard = ::Rails.env.to_sym
         self.default_role = :primary
 
-        include ActiveRecord::Base
-        include ActiveRecord::AttributeMethods
+        prepend ActiveRecord::Base
+        prepend ActiveRecord::AttributeMethods
         include ActiveRecord::Persistence
         singleton_class.prepend ActiveRecord::ModelSchema::ClassMethods
 
@@ -111,25 +28,40 @@ module Switchman
         ::ActiveRecord::StatementCache::BindMap.prepend(ActiveRecord::StatementCache::BindMap)
         ::ActiveRecord::StatementCache::Substitute.send(:attr_accessor, :primary, :sharded)
 
-        ::ActiveRecord::Associations::CollectionAssociation.prepend(ActiveRecord::CollectionAssociation)
-        ::ActiveRecord::Associations::HasOneAssociation.prepend(ActiveRecord::ForeignAssociation)
-        ::ActiveRecord::Associations::HasManyAssociation.prepend(ActiveRecord::ForeignAssociation)
+        ::ActiveRecord::Associations::CollectionAssociation.prepend(ActiveRecord::Associations::CollectionAssociation)
+        ::ActiveRecord::Associations::HasOneAssociation.prepend(ActiveRecord::Associations::ForeignAssociation)
+        ::ActiveRecord::Associations::HasManyAssociation.prepend(ActiveRecord::Associations::ForeignAssociation)
 
         ::ActiveRecord::PredicateBuilder.singleton_class.prepend(ActiveRecord::PredicateBuilder)
 
-        prepend(ActiveRecord::AutosaveAssociation)
+        prepend(ActiveRecord::Associations::AutosaveAssociation)
 
-        ::ActiveRecord::Associations::Association.prepend(ActiveRecord::Association)
-        ::ActiveRecord::Associations::BelongsToAssociation.prepend(ActiveRecord::BelongsToAssociation)
-        ::ActiveRecord::Associations::CollectionProxy.include(ActiveRecord::CollectionProxy)
+        ::ActiveRecord::Associations::Association.prepend(ActiveRecord::Associations::Association)
+        ::ActiveRecord::Associations::BelongsToAssociation.prepend(ActiveRecord::Associations::BelongsToAssociation)
+        ::ActiveRecord::Associations::CollectionProxy.include(ActiveRecord::Associations::CollectionProxy)
 
-        ::ActiveRecord::Associations::Preloader::Association.prepend(ActiveRecord::Preloader::Association)
+        ::ActiveRecord::Associations::Preloader::Association.prepend(ActiveRecord::Associations::Preloader::Association)
+        ::ActiveRecord::Associations::Preloader::Association::LoaderRecords.prepend(
+          ActiveRecord::Associations::Preloader::Association::LoaderRecords
+        )
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::AbstractAdapter)
+        ::ActiveRecord::ConnectionAdapters::ConnectionHandler.prepend(ActiveRecord::ConnectionHandler)
         ::ActiveRecord::ConnectionAdapters::ConnectionPool.prepend(ActiveRecord::ConnectionPool)
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::QueryCache)
+        ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(ActiveRecord::PostgreSQLAdapter)
+        # https://github.com/rails/rails/commit/0016280f4fde55d96738887093dc333aae0d107b
+        if ::Rails.version < "7.2"
+          ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(ActiveRecord::PostgreSQLAdapter::ClassMethods)
+        else
+          ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.singleton_class.prepend(
+            ActiveRecord::PostgreSQLAdapter::ClassMethods
+          )
+        end
 
         ::ActiveRecord::DatabaseConfigurations.prepend(ActiveRecord::DatabaseConfigurations)
-        ::ActiveRecord::DatabaseConfigurations::DatabaseConfig.prepend(ActiveRecord::DatabaseConfigurations::DatabaseConfig)
+        ::ActiveRecord::DatabaseConfigurations::DatabaseConfig.prepend(
+          ActiveRecord::DatabaseConfigurations::DatabaseConfig
+        )
 
         ::ActiveRecord::LogSubscriber.prepend(ActiveRecord::LogSubscriber)
         ::ActiveRecord::Migration.prepend(ActiveRecord::Migration)
@@ -137,6 +69,11 @@ module Switchman
         ::ActiveRecord::MigrationContext.prepend(ActiveRecord::MigrationContext)
         ::ActiveRecord::Migrator.prepend(ActiveRecord::Migrator)
 
+        if ::Rails.version > "7.1.3"
+          ::ActiveRecord::PendingMigrationConnection.singleton_class
+                                                    .include(ActiveRecord::PendingMigrationConnection::ClassMethods)
+        end
+
         ::ActiveRecord::Reflection::AbstractReflection.include(ActiveRecord::Reflection::AbstractReflection)
         ::ActiveRecord::Reflection::AssociationReflection.prepend(ActiveRecord::Reflection::AssociationScopeCache)
         ::ActiveRecord::Reflection::ThroughReflection.prepend(ActiveRecord::Reflection::AssociationScopeCache)
@@ -146,64 +83,81 @@ module Switchman
         ::ActiveRecord::Relation.include(ActiveRecord::QueryMethods)
         ::ActiveRecord::Relation.prepend(GuardRail::Relation)
         ::ActiveRecord::Relation.prepend(ActiveRecord::Relation)
+        ::ActiveRecord::Relation.prepend(ActiveRecord::Relation::InsertUpsertAll) if ::Rails.version >= "7.2"
         ::ActiveRecord::Relation.include(ActiveRecord::SpawnMethods)
         ::ActiveRecord::Relation.include(CallSuper)
 
-        ::ActiveRecord::PredicateBuilder::AssociationQueryValue.prepend(ActiveRecord::PredicateBuilder::AssociationQueryValue)
+        ::ActiveRecord::PredicateBuilder::PolymorphicArrayValue.prepend(
+          ActiveRecord::PredicateBuilder::PolymorphicArrayValue
+        )
 
         ::ActiveRecord::Tasks::DatabaseTasks.singleton_class.prepend(ActiveRecord::Tasks::DatabaseTasks)
 
+        ::ActiveRecord::TestFixtures.prepend(ActiveRecord::TestFixtures)
+
         ::ActiveRecord::TypeCaster::Map.include(ActiveRecord::TypeCaster::Map)
         ::ActiveRecord::TypeCaster::Connection.include(ActiveRecord::TypeCaster::Connection)
 
-        ::Rails.singleton_class.prepend(Rails::ClassMethods)
-
         ::Arel::Table.prepend(Arel::Table)
         ::Arel::Visitors::ToSql.prepend(Arel::Visitors::ToSql)
-      end
-    end
-
-    def self.foreign_key_check(name, type, limit: nil)
-      puts "WARNING: All foreign keys need to be 8-byte integers. #{name} looks like a foreign key. If so, please add the option: `:limit => 8`" if name.to_s =~ /_id\z/ && type.to_s == 'integer' && limit.to_i < 8
-    end
 
-    initializer 'switchman.extend_connection_adapters', after: 'active_record.initialize_database' do
-      ::ActiveSupport.on_load(:active_record) do
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.descendants.each do |klass|
           klass.prepend(ActiveRecord::AbstractAdapter::ForeignKeyCheck)
         end
 
-        require 'switchman/active_record/table_definition'
         ::ActiveRecord::ConnectionAdapters::TableDefinition.prepend(ActiveRecord::TableDefinition)
-
-        if defined?(::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
-          require 'switchman/active_record/postgresql_adapter'
-          ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.prepend(ActiveRecord::PostgreSQLAdapter)
-        end
-
-        Shard.send(:initialize_sharding)
       end
+      # Ensure that ActiveRecord::Base is always loaded before any app-level
+      # initializers can go try to load Switchman::Shard or we get a loop
+      ::ActiveRecord::Base
     end
 
-    initializer 'switchman.eager_load' do
-      ::ActiveSupport.on_load(:before_eager_load) do
-        # This needs to be loaded before Switchman::Shard, otherwise it won't autoload it correctly
-        require 'active_record/base'
+    initializer "switchman.error_patch", after: "active_record.initialize_database" do
+      ::ActiveSupport.on_load(:active_record) do
+        ::StandardError.include(StandardError)
       end
     end
 
-    initializer 'switchman.extend_guard_rail', before: 'switchman.extend_ar' do
-      ::ActiveSupport.on_load(:active_record) do
-        require 'switchman/guard_rail'
+    initializer "switchman.initialize_cache", before: :initialize_cache, after: "active_record.initialize_database" do
+      ::ActiveSupport::Cache.singleton_class.prepend(ActiveSupport::Cache::ClassMethods)
+
+      # if we haven't already setup our cache map out-of-band, set it up from
+      # config.cache_store now. behaves similarly to Rails' default
+      # initialize_cache initializer, but for each value in the map, rather
+      # than just Rails.cache. if config.cache_store is a flat value, uses it
+      # to fill just the Rails.env entry in the cache map.
+      unless Switchman.config[:cache_map].present?
+        cache_store_config = ::Rails.configuration.cache_store
+        cache_store_config = { ::Rails.env => cache_store_config } unless cache_store_config.is_a?(Hash)
 
-        ::GuardRail.singleton_class.prepend(GuardRail::ClassMethods)
+        Switchman.config[:cache_map] = ::ActiveSupport::Cache.lookup_stores(cache_store_config)
       end
-    end
 
-    initializer 'switchman.extend_controller', after: 'guard_rail.extend_ar' do
-      ::ActiveSupport.on_load(:action_controller) do
-        require 'switchman/action_controller/caching'
+      # if the configured cache map (either from before, or as populated from
+      # config.cache_store) didn't have an entry for Rails.env, add one using
+      # lookup_store(nil); matches the behavior of Rails' default
+      # initialize_cache initializer when config.cache_store is nil.
+      unless Switchman.config[:cache_map].key?(::Rails.env)
+        value = ::ActiveSupport::Cache.lookup_store(nil)
+        Switchman.config[:cache_map][::Rails.env] = value
+      end
+
+      middlewares = Switchman.config[:cache_map].values.filter_map do |store|
+        store.middleware if store.respond_to?(:middleware)
+      end.uniq
+      middlewares.each do |middleware|
+        config.middleware.insert_before("Rack::Runtime", middleware)
+      end
 
+      # prevent :initialize_cache from trying to (or needing to) set
+      # Rails.cache. once our switchman.extend_ar initializer (below) runs
+      # Rails.cache will be overridden to pull appropriate values from the
+      # cache map, but between now and then, Rails.cache should return the
+      # Rails.env entry in the cache map.
+      ::Rails.cache = Switchman.config[:cache_map][::Rails.env]
+      ::Rails.singleton_class.prepend(Rails::ClassMethods)
+
+      ::ActiveSupport.on_load(:action_controller) do
         ::ActionController::Base.include(ActionController::Caching)
       end
     end