switchman 3.0.19 → 3.0.22

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: eea092610f18550e6ab0d24baa9a2746e8452341116c0b5e3a94c15b1bcb9be0
-  data.tar.gz: 53e0c2f7735838283311e57017de6bf020aae61f38dd394821a5bab5fdda869e
+  metadata.gz: f6ed57a128cb7da74d0df280b924f4369dcd224039343f4476796ae09d8573f5
+  data.tar.gz: 732fbee6c445f210f9b97b3921f30356a2eb7f9df1c925631dccce5e91612904
 SHA512:
-  metadata.gz: da4096ca84676fecda1bee5e99a9e7e34c1d4cc2233e4a32f391d68ccd9d93a8d641b5d4a5ecd330e5df340daca8611bad0226709ba9f9fc6653282821216e0a
-  data.tar.gz: d8a6f2245ac58e5ab64398245f7b14188cdb4675dea430fc2dd256b5aea73f68a81a30f5bcd3cd626634c492932bf85f3becab5dfa0a6e773abc6639615ab247
+  metadata.gz: 8fdbbdc4309ebac36a0c779aefe89163f16419581752307223a367c3e47197d96a64daa8f8c4182b1b94512cf871891e531ea5c99616b0fc2e808c1b698d4f24
+  data.tar.gz: 93bd4a98e64cba28e28a80f4b33e8c93ef00ef3ba1c467b81c2ac93665ae0678acd2715d3781566c4e60df29749a353f08546d7aff9809fef2e53b5eb7f81d43

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

@@ -14,8 +14,6 @@ module Switchman
         def sharded_model
           self.abstract_class = true
 
-          return if self == UnshardedRecord
-
           Shard.send(:add_sharded_model, self)
         end
 
@@ -64,6 +62,29 @@ module Switchman
           current_role != 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
+          return super if Thread.current.thread_variable?(:ar_connected_to_stack)
+
+          ret = super
+          DatabaseServer.guard_servers
+          ret
+        end
+
         # significant change: Allow per-shard roles
         def current_role(without_overrides: false)
           return super() if without_overrides

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
 
@@ -49,31 +46,42 @@ module Switchman
         servers[rand(servers.length)]
       end
 
+      def guard_servers
+        all.each { |db| db.guard! if db.config[:prefer_secondary] }
+      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
         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
@@ -133,6 +141,7 @@ module Switchman
     # when doing writes (then it falls back to the current
     # value of GuardRail.environment)
     def guard!(environment = :secondary)
+      DatabaseServer.send(:reference_role, environment)
       ::ActiveRecord::Base.connected_to_stack << { shard_roles: { id.to_sym => environment }, klasses: [::ActiveRecord::Base] }
     end
 

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,13 @@ 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)
-
-        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' 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 +28,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 +77,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: 'switchman.active_record_patch' 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
@@ -376,16 +368,12 @@ module Switchman
 
       private
 
-      def sharding_initialized
-        @sharding_initialized ||= false
-      end
-
       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,9 +395,6 @@ module Switchman
 
           klass.connects_to shards: connects_to_hash
         end
-        DatabaseServer.all.each { |db| db.guard! if db.config[:prefer_secondary] } unless @sharding_initialized
-
-        @sharding_initialized = true
       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.19'
+  VERSION = '3.0.22'
 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.19
+  version: 3.0.22
 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-07 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