switchman 3.0.21 → 3.0.24

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: df4352d56024adc191b904c1786c43dff9a685be0c7e82e5813339a107ac6df3
-  data.tar.gz: 7f6ce97fdadcf3dff63b64065e94b2c4de5a481b234ec6f36429ff7eb796eccc
+  metadata.gz: 469ef6bca62776e9ba974a5c6d2a6220fc6ddcf1649956f43247465cbfb24f37
+  data.tar.gz: 9fefef5095835cf28a6c72197ed74f6de9974d7c9c9019ad4ea87209d98c37e9
 SHA512:
-  metadata.gz: 26b862b5b2d44c3aebf24d2ce469e332ff3af117b75e05559d1343b399a01d36db1a321f6fb2f48e535d3e0f99517c56c0dd549911d47b9ef024e8652b5f54b3
-  data.tar.gz: 229c8277f8df802f2f5f29221d66afbaba873ff70b0f97ca0ff9356a99b58d8b1dc08bd6fb07c75800275dfc4ef63ec8296451bdfed1204860fbfc1cd764d291
+  metadata.gz: 96d01726ff4e0e560e63964c5607417b7272202109b6e7e718769e9ba2cc04f82000f442c70aef8d6140e6dac00925197f0a95c70b790d47fbc2355d9021ae39
+  data.tar.gz: 4b2edb3c0400f3280ba579bfaac61fce263ea5dd0881e5eeb77e675e0d2c1dfe04168e059f7bb91ac601ce2ef210c50c252945258fe2166ee6cf21c458defac0

data/lib/switchman/action_controller/caching.rb

@@ -13,8 +13,8 @@ module Switchman
         # disallow assigning to ActionController::Base.cache_store or
         # ActionController::Base#cache_store for the same reasons we disallow
         # assigning to Rails.cache
-        def cache_store=(_cache)
-          raise NoMethodError
+        def cache_store=(cache)
+          raise NoMethodError unless cache == ::Rails.cache
         end
       end
 

data/lib/switchman/active_record/abstract_adapter.rb

@@ -1,13 +1,11 @@
 # frozen_string_literal: true
 
-require 'switchman/sharded_instrumenter'
-
 module Switchman
   module ActiveRecord
     module AbstractAdapter
       module ForeignKeyCheck
         def add_column(table, name, type, limit: nil, **)
-          Engine.foreign_key_check(name, type, limit: limit)
+          Switchman.foreign_key_check(name, type, limit: limit)
           super
         end
       end

data/lib/switchman/active_record/associations.rb

@@ -0,0 +1,208 @@
+# frozen_string_literal: true
+
+module Switchman
+  module ActiveRecord
+    module Associations
+      module Association
+        def shard
+          reflection.shard(owner)
+        end
+
+        def build_record(*args)
+          shard.activate { super }
+        end
+
+        def load_target
+          shard.activate { super }
+        end
+
+        def scope
+          shard_value = @reflection.options[:multishard] ? @owner : shard
+          @owner.shard.activate { super.shard(shard_value, :association) }
+        end
+      end
+
+      module CollectionAssociation
+        def find_target
+          shards = reflection.options[:multishard] && owner.respond_to?(:associated_shards) ? owner.associated_shards : [shard]
+          # activate both the owner and the target's shard category, so that Reflection#join_id_for,
+          # when called for the owner, will be returned relative to shard the query will execute on
+          Shard.with_each_shard(shards, [klass.connection_classes, owner.class.connection_classes].uniq) do
+            super
+          end
+        end
+
+        def _create_record(*)
+          shard.activate { super }
+        end
+      end
+
+      module BelongsToAssociation
+        def replace_keys(record, force: false)
+          if record&.class&.sharded_column?(reflection.association_primary_key(record.class))
+            foreign_id = record[reflection.association_primary_key(record.class)]
+            owner[reflection.foreign_key] = Shard.relative_id_for(foreign_id, record.shard, owner.shard)
+          else
+            super
+          end
+        end
+
+        def shard
+          if @owner.class.sharded_column?(@reflection.foreign_key) &&
+             (foreign_id = @owner[@reflection.foreign_key])
+            Shard.shard_for(foreign_id, @owner.shard)
+          else
+            super
+          end
+        end
+      end
+
+      module ForeignAssociation
+        # significant change:
+        #   * transpose the key to the correct shard
+        def set_owner_attributes(record) # rubocop:disable Naming/AccessorMethodName
+          return if options[:through]
+
+          key = owner._read_attribute(reflection.join_foreign_key)
+          key = Shard.relative_id_for(key, owner.shard, shard)
+          record._write_attribute(reflection.join_primary_key, key)
+
+          record._write_attribute(reflection.type, owner.class.polymorphic_name) if reflection.type
+        end
+      end
+
+      module Extension
+        def self.build(_model, _reflection); end
+
+        def self.valid_options
+          [:multishard]
+        end
+      end
+
+      ::ActiveRecord::Associations::Builder::Association.extensions << Extension
+
+      module Preloader
+        module Association
+          # Copypasta from Activerecord but with added global_id_for goodness.
+          def records_for(ids)
+            scope.where(association_key_name => ids).load do |record|
+              global_key = if model.connection_classes == UnshardedRecord
+                             convert_key(record[association_key_name])
+                           else
+                             Shard.global_id_for(record[association_key_name], record.shard)
+                           end
+              owner = owners_by_key[convert_key(global_key)].first
+              association = owner.association(reflection.name)
+              association.set_inverse_instance(record)
+            end
+          end
+
+          # significant changes:
+          #  * partition_by_shard the records_for call
+          #  * re-globalize the fetched owner id before looking up in the map
+          def load_records
+            # owners can be duplicated when a relation has a collection association join
+            # #compare_by_identity makes such owners different hash keys
+            @records_by_owner = {}.compare_by_identity
+
+            if owner_keys.empty?
+              raw_records = []
+            else
+              # determine the shard to search for each owner
+              if reflection.macro == :belongs_to
+                # for belongs_to, it's the shard of the foreign_key
+                partition_proc = lambda do |owner|
+                  if owner.class.sharded_column?(owner_key_name)
+                    Shard.shard_for(owner[owner_key_name], owner.shard)
+                  else
+                    Shard.current
+                  end
+                end
+              elsif !reflection.options[:multishard]
+                # for non-multishard associations, it's *just* the owner's shard
+                partition_proc = ->(owner) { owner.shard }
+              end
+
+              raw_records = Shard.partition_by_shard(owners, partition_proc) do |partitioned_owners|
+                relative_owner_keys = partitioned_owners.map do |owner|
+                  key = owner[owner_key_name]
+                  if key && owner.class.sharded_column?(owner_key_name)
+                    key = Shard.relative_id_for(key, owner.shard,
+                                                Shard.current(klass.connection_classes))
+                  end
+                  convert_key(key)
+                end
+                relative_owner_keys.compact!
+                relative_owner_keys.uniq!
+                records_for(relative_owner_keys)
+              end
+            end
+
+            @preloaded_records = raw_records.select do |record|
+              assignments = false
+
+              owner_key = record[association_key_name]
+              if owner_key && record.class.sharded_column?(association_key_name)
+                owner_key = Shard.global_id_for(owner_key,
+                                                record.shard)
+              end
+
+              owners_by_key[convert_key(owner_key)].each do |owner|
+                entries = (@records_by_owner[owner] ||= [])
+
+                if reflection.collection? || entries.empty?
+                  entries << record
+                  assignments = true
+                end
+              end
+
+              assignments
+            end
+          end
+
+          # significant change: globalize keys on sharded columns
+          def owners_by_key
+            @owners_by_key ||= owners.each_with_object({}) do |owner, result|
+              key = owner[owner_key_name]
+              key = Shard.global_id_for(key, owner.shard) if key && owner.class.sharded_column?(owner_key_name)
+              key = convert_key(key)
+              (result[key] ||= []) << owner if key
+            end
+          end
+
+          # significant change: don't cache scope (since it could be for different shards)
+          def scope
+            build_scope
+          end
+        end
+      end
+
+      module CollectionProxy
+        def initialize(*args)
+          super
+          self.shard_value = scope.shard_value
+          self.shard_source_value = :association
+        end
+
+        def shard(*args)
+          scope.shard(*args)
+        end
+      end
+
+      module AutosaveAssociation
+        def record_changed?(reflection, record, key)
+          record.new_record? ||
+            (record.has_attribute?(reflection.foreign_key) && record.send(reflection.foreign_key) != key) || # have to use send instead of [] because sharding
+            record.attribute_changed?(reflection.foreign_key)
+        end
+
+        def save_belongs_to_association(reflection)
+          # this seems counter-intuitive, but the autosave code will assign to attribute bypassing switchman,
+          # after reading the id attribute _without_ bypassing switchman. So we need Shard.current for the
+          # category of the associated record to match Shard.current for the category of self
+          shard.activate(connection_classes_for_reflection(reflection)) { super }
+        end
+      end
+    end
+  end
+end

data/lib/switchman/active_record/base.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require 'switchman/database_server'
-
 module Switchman
   module ActiveRecord
     module Base
@@ -16,8 +14,6 @@ module Switchman
         def sharded_model
           self.abstract_class = true
 
-          return if self == UnshardedRecord
-
           Shard.send(:add_sharded_model, self)
         end
 
@@ -62,8 +58,23 @@ module Switchman
           end
         end
 
-        def current_role_overriden?
-          current_role != current_role(without_overrides: true)
+        def role_overriden?(shard_id)
+          current_role(target_shard: shard_id) != current_role(without_overrides: true)
+        end
+
+        def establish_connection(config_or_env = nil)
+          raise ArgumentError, 'establish connection cannot be used on the non-current shard/role' if config_or_env.is_a?(Symbol) && config_or_env != ::Rails.env.to_sym
+
+          # Ensure we don't randomly surprise change the connection parms associated with a shard/role
+          config_or_env = nil if config_or_env == ::Rails.env.to_sym
+
+          config_or_env ||= if current_shard == ::Rails.env.to_sym && current_role == :primary
+                              :primary
+                            else
+                              "#{current_shard}/#{current_role}".to_sym
+                            end
+
+          super(config_or_env)
         end
 
         def connected_to_stack
@@ -75,12 +86,12 @@ module Switchman
         end
 
         # significant change: Allow per-shard roles
-        def current_role(without_overrides: false)
+        def current_role(without_overrides: false, target_shard: current_shard)
           return super() if without_overrides
 
           sharded_role = nil
           connected_to_stack.reverse_each do |hash|
-            shard_role = hash.dig(:shard_roles, current_shard)
+            shard_role = hash.dig(:shard_roles, target_shard)
             if shard_role && (hash[:klasses].include?(::ActiveRecord::Base) || hash[:klasses].include?(connection_classes))
               sharded_role = shard_role
               break

data/lib/switchman/active_record/connection_pool.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require 'switchman/errors'
-
 module Switchman
   module ActiveRecord
     module ConnectionPool
@@ -20,7 +18,7 @@ module Switchman
 
       def connection(switch_shard: true)
         conn = super()
-        raise NonExistentShardError if current_shard.new_record?
+        raise Errors::NonExistentShardError if current_shard.new_record?
 
         switch_database(conn) if conn.shard != current_shard && switch_shard
         conn

data/lib/switchman/active_record/table_definition.rb

@@ -4,7 +4,7 @@ module Switchman
   module ActiveRecord
     module TableDefinition
       def column(name, type, limit: nil, **)
-        Engine.foreign_key_check(name, type, limit: limit)
+        Switchman.foreign_key_check(name, type, limit: limit)
         super
       end
     end

data/lib/switchman/active_support/cache.rb

@@ -4,6 +4,22 @@ module Switchman
   module ActiveSupport
     module Cache
       module ClassMethods
+        def 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
+
         def lookup_store(*store_options)
           store = super
           # can't use defined?, because it's a _ruby_ autoloaded constant,

data/lib/switchman/database_server.rb

@@ -8,15 +8,12 @@ module Switchman
 
     class << self
       attr_accessor :creating_new_shard
+      attr_reader :all_roles
 
       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)
@@ -38,7 +35,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
 
@@ -59,25 +56,32 @@ module Switchman
         return if all_roles.include?(role)
 
         @all_roles << role
-        Shard.send(:initialize_sharding)
+        Shard.send(:configure_connects_to)
       end
 
       def 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('/')
             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.to_sym
+              @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
+          Shard.send(:configure_connects_to)
         end
         @database_servers
       end
@@ -146,7 +150,7 @@ module Switchman
     end
 
     def unguard
-      return yield unless ::ActiveRecord::Base.current_role_overriden?
+      return yield unless ::ActiveRecord::Base.role_overriden?(id.to_sym)
 
       begin
         unguard!

data/lib/switchman/default_shard.rb

@@ -1,7 +1,5 @@
 # frozen_string_literal: true
 
-require 'switchman/database_server'
-
 module Switchman
   class DefaultShard
     def id

data/lib/switchman/engine.rb

@@ -8,101 +8,18 @@ module Switchman
     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)
+    # after :initialize_dependency_mechanism to ensure autoloading is configured for any downstream initializers that care
+    # In rails 7.0 we should be able to just use an explicit after on configuring the once autoloaders and not need to go monkey around with initializer order
+    initialize_dependency_mechanism = ::Rails::Application::Bootstrap.initializers.find { |i| i.name == :initialize_dependency_mechanism }
+    initialize_dependency_mechanism.instance_variable_get(:@options)[:after] = :set_autoload_paths
 
-        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: :initialize_dependency_mechanism do
       ::ActiveSupport.on_load(:active_record) do
         # Switchman requires postgres, so just always load the pg adapter
         require 'active_record/connection_adapters/postgresql_adapter'
 
-        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/postgresql_adapter'
-        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/active_record/test_fixtures'
-        require 'switchman/arel'
-        require 'switchman/call_super'
-        require 'switchman/rails'
-        require 'switchman/guard_rail/relation'
-        require 'switchman/standard_error'
-
-        ::StandardError.include(StandardError)
-
         self.default_shard = ::Rails.env.to_sym
         self.default_role = :primary
 
@@ -116,19 +33,19 @@ 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::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::AbstractAdapter)
         ::ActiveRecord::ConnectionAdapters::ConnectionPool.prepend(ActiveRecord::ConnectionPool)
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::QueryCache)
@@ -165,49 +82,65 @@ module Switchman
         ::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)
-
-        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)
 
-        ::GuardRail.singleton_class.prepend(GuardRail::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] = ::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.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]
+      ::Rails.singleton_class.prepend(Rails::ClassMethods)
+
+      ::ActiveSupport.on_load(:action_controller) do
         ::ActionController::Base.include(ActionController::Caching)
       end
     end

data/lib/switchman/errors.rb

@@ -1,7 +1,9 @@
 # frozen_string_literal: true
 
 module Switchman
-  class NonExistentShardError < RuntimeError; end
+  module Errors
+    class NonExistentShardError < RuntimeError; end
 
-  class ParallelShardExecError < RuntimeError; end
+    class ParallelShardExecError < RuntimeError; end
+  end
 end

data/lib/switchman/rails.rb

@@ -4,16 +4,13 @@ module Switchman
   module Rails
     module ClassMethods
       def self.prepended(klass)
-        # in Rails 4+, the Rails.cache= method was used during bootstrap to set
-        # Rails.cache(_without_sharding) to the value from the config file. but now
-        # that that's done (the bootstrap happened before this module is included
-        # into Rails), we want to make sure no one tries to assign to Rails.cache,
+        # we want to make sure no one tries to assign to Rails.cache,
         # because it would be wrong w.r.t. sharding.
         klass.send(:remove_method, :cache=)
       end
 
       def cache
-        Switchman::Shard.current ? Switchman::Shard.current.database_server.cache_store : super
+        Switchman::Shard.current.database_server.cache_store
       end
     end
   end

data/{app/models → lib}/switchman/shard.rb

@@ -1,10 +1,5 @@
 # frozen_string_literal: true
 
-require 'switchman/database_server'
-require 'switchman/default_shard'
-require 'switchman/environment'
-require 'switchman/errors'
-
 module Switchman
   class Shard < UnshardedRecord
     # ten trillion possible ids per shard. yup.
@@ -42,12 +37,9 @@ module Switchman
           # Now find the actual record, if it exists
           @default = begin
             find_cached('default_shard') { Shard.where(default: true).take } || default
-          # If we are *super* early in boot, the connection pool won't exist; we don't want to fill in the default shard yet
-          # Otherwise, rescue the fake default if the table doesn't exist
           rescue
-            sharding_initialized ? default : nil
+            default
           end
-          return default unless @default
 
           # make sure this is not erroneously cached
           @default.database_server.remove_instance_variable(:@primary_shard) if @default.database_server.instance_variable_defined?(:@primary_shard)
@@ -193,7 +185,7 @@ module Switchman
           unless errors.empty?
             raise errors.first.exception if errors.length == 1
 
-            raise ParallelShardExecError,
+            raise Errors::ParallelShardExecError,
                   "The following database server(s) did not finish processing cleanly: #{errors.map(&:name).sort.join(', ')}",
                   cause: errors.first.exception
           end
@@ -374,18 +366,14 @@ module Switchman
         shard || source_shard || Shard.current
       end
 
-      def sharding_initialized
-        @sharding_initialized ||= false
-      end
-
       private
 
       def add_sharded_model(klass)
         @sharded_models = (sharded_models + [klass]).freeze
-        initialize_sharding
+        configure_connects_to
       end
 
-      def initialize_sharding
+      def configure_connects_to
         full_connects_to_hash = DatabaseServer.all.to_h { |db| [db.id.to_sym, db.connects_to_hash] }
         sharded_models.each do |klass|
           connects_to_hash = full_connects_to_hash.deep_dup
@@ -407,13 +395,6 @@ module Switchman
 
           klass.connects_to shards: connects_to_hash
         end
-
-        return if @sharding_initialized
-
-        # If we hadn't initialized sharding yet, the servers won't be guarded
-        # The order matters here or guard_servers will be a noop
-        @sharding_initialized = true
-        DatabaseServer.guard_servers
       end
 
       # in-process caching

data/lib/switchman/standard_error.rb

@@ -12,11 +12,7 @@ module Switchman
       begin
         Thread.current[:switchman_error_handler] = true
 
-        begin
-          @active_shards ||= Shard.active_shards if defined?(Shard)
-        rescue
-          # If we hit an error really early in boot, activerecord may not be initialized yet
-        end
+        @active_shards ||= Shard.active_shards
       ensure
         Thread.current[:switchman_error_handler] = nil
       end

data/{app/models → lib}/switchman/unsharded_record.rb

@@ -2,6 +2,6 @@
 
 module Switchman
   class UnshardedRecord < ::ActiveRecord::Base
-    sharded_model
+    self.abstract_class = true
   end
 end

data/lib/switchman/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Switchman
-  VERSION = '3.0.21'
+  VERSION = '3.0.24'
 end

data/lib/switchman.rb

@@ -1,8 +1,21 @@
 # frozen_string_literal: true
 
 require 'guard_rail'
-require 'switchman/parallel'
-require 'switchman/engine'
+require 'zeitwerk'
+
+class SwitchmanInflector < Zeitwerk::GemInflector
+  def camelize(basename, abspath)
+    if basename =~ /\Apostgresql_(.*)/
+      'PostgreSQL' + super($1, abspath)
+    else
+      super
+    end
+  end
+end
+
+loader = Zeitwerk::Loader.for_gem
+loader.inflector = SwitchmanInflector.new(__FILE__)
+loader.setup
 
 module Switchman
   def self.config
@@ -18,5 +31,12 @@ module Switchman
     @cache = cache
   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
+
   class OrderOnMultiShardQuery < RuntimeError; end
 end
+
+# Load the engine and everything associated at gem load time
+Switchman::Engine

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: switchman
 version: !ruby/object:Gem::Version
-  version: 3.0.21
+  version: 3.0.24
 platform: ruby
 authors:
 - Cody Cutrer
@@ -10,7 +10,7 @@ authors:
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2022-04-06 00:00:00.000000000 Z
+date: 2022-04-28 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -242,8 +242,6 @@ extensions: []
 extra_rdoc_files: []
 files:
 - Rakefile
-- app/models/switchman/shard.rb
-- app/models/switchman/unsharded_record.rb
 - db/migrate/20130328212039_create_switchman_shards.rb
 - db/migrate/20130328224244_create_default_shard.rb
 - db/migrate/20161206323434_add_back_default_string_limits_switchman.rb
@@ -253,7 +251,7 @@ files:
 - lib/switchman.rb
 - lib/switchman/action_controller/caching.rb
 - lib/switchman/active_record/abstract_adapter.rb
-- lib/switchman/active_record/association.rb
+- lib/switchman/active_record/associations.rb
 - lib/switchman/active_record/attribute_methods.rb
 - lib/switchman/active_record/base.rb
 - lib/switchman/active_record/calculations.rb
@@ -290,9 +288,11 @@ files:
 - lib/switchman/parallel.rb
 - lib/switchman/r_spec_helper.rb
 - lib/switchman/rails.rb
+- lib/switchman/shard.rb
 - lib/switchman/sharded_instrumenter.rb
 - lib/switchman/standard_error.rb
 - lib/switchman/test_helper.rb
+- lib/switchman/unsharded_record.rb
 - lib/switchman/version.rb
 - lib/tasks/switchman.rake
 homepage: http://www.instructure.com/

data/lib/switchman/active_record/association.rb

@@ -1,206 +0,0 @@
-# frozen_string_literal: true
-
-module Switchman
-  module ActiveRecord
-    module Association
-      def shard
-        reflection.shard(owner)
-      end
-
-      def build_record(*args)
-        shard.activate { super }
-      end
-
-      def load_target
-        shard.activate { super }
-      end
-
-      def scope
-        shard_value = @reflection.options[:multishard] ? @owner : shard
-        @owner.shard.activate { super.shard(shard_value, :association) }
-      end
-    end
-
-    module CollectionAssociation
-      def find_target
-        shards = reflection.options[:multishard] && owner.respond_to?(:associated_shards) ? owner.associated_shards : [shard]
-        # activate both the owner and the target's shard category, so that Reflection#join_id_for,
-        # when called for the owner, will be returned relative to shard the query will execute on
-        Shard.with_each_shard(shards, [klass.connection_classes, owner.class.connection_classes].uniq) do
-          super
-        end
-      end
-
-      def _create_record(*)
-        shard.activate { super }
-      end
-    end
-
-    module BelongsToAssociation
-      def replace_keys(record, force: false)
-        if record&.class&.sharded_column?(reflection.association_primary_key(record.class))
-          foreign_id = record[reflection.association_primary_key(record.class)]
-          owner[reflection.foreign_key] = Shard.relative_id_for(foreign_id, record.shard, owner.shard)
-        else
-          super
-        end
-      end
-
-      def shard
-        if @owner.class.sharded_column?(@reflection.foreign_key) &&
-           (foreign_id = @owner[@reflection.foreign_key])
-          Shard.shard_for(foreign_id, @owner.shard)
-        else
-          super
-        end
-      end
-    end
-
-    module ForeignAssociation
-      # significant change:
-      #   * transpose the key to the correct shard
-      def set_owner_attributes(record) # rubocop:disable Naming/AccessorMethodName
-        return if options[:through]
-
-        key = owner._read_attribute(reflection.join_foreign_key)
-        key = Shard.relative_id_for(key, owner.shard, shard)
-        record._write_attribute(reflection.join_primary_key, key)
-
-        record._write_attribute(reflection.type, owner.class.polymorphic_name) if reflection.type
-      end
-    end
-
-    module Extension
-      def self.build(_model, _reflection); end
-
-      def self.valid_options
-        [:multishard]
-      end
-    end
-
-    ::ActiveRecord::Associations::Builder::Association.extensions << Extension
-
-    module Preloader
-      module Association
-        # Copypasta from Activerecord but with added global_id_for goodness.
-        def records_for(ids)
-          scope.where(association_key_name => ids).load do |record|
-            global_key = if model.connection_classes == UnshardedRecord
-                           convert_key(record[association_key_name])
-                         else
-                           Shard.global_id_for(record[association_key_name], record.shard)
-                         end
-            owner = owners_by_key[convert_key(global_key)].first
-            association = owner.association(reflection.name)
-            association.set_inverse_instance(record)
-          end
-        end
-
-        # significant changes:
-        #  * partition_by_shard the records_for call
-        #  * re-globalize the fetched owner id before looking up in the map
-        def load_records
-          # owners can be duplicated when a relation has a collection association join
-          # #compare_by_identity makes such owners different hash keys
-          @records_by_owner = {}.compare_by_identity
-
-          if owner_keys.empty?
-            raw_records = []
-          else
-            # determine the shard to search for each owner
-            if reflection.macro == :belongs_to
-              # for belongs_to, it's the shard of the foreign_key
-              partition_proc = lambda do |owner|
-                if owner.class.sharded_column?(owner_key_name)
-                  Shard.shard_for(owner[owner_key_name], owner.shard)
-                else
-                  Shard.current
-                end
-              end
-            elsif !reflection.options[:multishard]
-              # for non-multishard associations, it's *just* the owner's shard
-              partition_proc = ->(owner) { owner.shard }
-            end
-
-            raw_records = Shard.partition_by_shard(owners, partition_proc) do |partitioned_owners|
-              relative_owner_keys = partitioned_owners.map do |owner|
-                key = owner[owner_key_name]
-                if key && owner.class.sharded_column?(owner_key_name)
-                  key = Shard.relative_id_for(key, owner.shard,
-                                              Shard.current(klass.connection_classes))
-                end
-                convert_key(key)
-              end
-              relative_owner_keys.compact!
-              relative_owner_keys.uniq!
-              records_for(relative_owner_keys)
-            end
-          end
-
-          @preloaded_records = raw_records.select do |record|
-            assignments = false
-
-            owner_key = record[association_key_name]
-            if owner_key && record.class.sharded_column?(association_key_name)
-              owner_key = Shard.global_id_for(owner_key,
-                                              record.shard)
-            end
-
-            owners_by_key[convert_key(owner_key)].each do |owner|
-              entries = (@records_by_owner[owner] ||= [])
-
-              if reflection.collection? || entries.empty?
-                entries << record
-                assignments = true
-              end
-            end
-
-            assignments
-          end
-        end
-
-        # significant change: globalize keys on sharded columns
-        def owners_by_key
-          @owners_by_key ||= owners.each_with_object({}) do |owner, result|
-            key = owner[owner_key_name]
-            key = Shard.global_id_for(key, owner.shard) if key && owner.class.sharded_column?(owner_key_name)
-            key = convert_key(key)
-            (result[key] ||= []) << owner if key
-          end
-        end
-
-        # significant change: don't cache scope (since it could be for different shards)
-        def scope
-          build_scope
-        end
-      end
-    end
-
-    module CollectionProxy
-      def initialize(*args)
-        super
-        self.shard_value = scope.shard_value
-        self.shard_source_value = :association
-      end
-
-      def shard(*args)
-        scope.shard(*args)
-      end
-    end
-
-    module AutosaveAssociation
-      def record_changed?(reflection, record, key)
-        record.new_record? ||
-          (record.has_attribute?(reflection.foreign_key) && record.send(reflection.foreign_key) != key) || # have to use send instead of [] because sharding
-          record.attribute_changed?(reflection.foreign_key)
-      end
-
-      def save_belongs_to_association(reflection)
-        # this seems counter-intuitive, but the autosave code will assign to attribute bypassing switchman,
-        # after reading the id attribute _without_ bypassing switchman. So we need Shard.current for the
-        # category of the associated record to match Shard.current for the category of self
-        shard.activate(connection_classes_for_reflection(reflection)) { super }
-      end
-    end
-  end
-end