switchman 3.0.22 → 3.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: f6ed57a128cb7da74d0df280b924f4369dcd224039343f4476796ae09d8573f5
-  data.tar.gz: 732fbee6c445f210f9b97b3921f30356a2eb7f9df1c925631dccce5e91612904
+  metadata.gz: 48b27124d5233adf585aed105d0b2522c76249bdd1e13fbcfd7e904241d3a624
+  data.tar.gz: fea1ccc0ad63f00633fb90bee1d7bd4606f6de0e18564250865c1b1dde61ef8d
 SHA512:
-  metadata.gz: 8fdbbdc4309ebac36a0c779aefe89163f16419581752307223a367c3e47197d96a64daa8f8c4182b1b94512cf871891e531ea5c99616b0fc2e808c1b698d4f24
-  data.tar.gz: 93bd4a98e64cba28e28a80f4b33e8c93ef00ef3ba1c467b81c2ac93665ae0678acd2715d3781566c4e60df29749a353f08546d7aff9809fef2e53b5eb7f81d43
+  metadata.gz: c05b445111dd0e4b16ecbfaea6ab84557fa2cbadc6102fa0af6fb43cff2f6bbb6606e825f76ed20f6b8a4a2a2966e31babadf66c5eb065a6c4d8a7eea6a8b92a
+  data.tar.gz: 0af9e15706142f90772d24e2187db93ec83331997bccbb25c1b89cec2dde1bd2414d2774fec3785c7bea1a1b06d3fa5ad49e15b716f0aa35794e2e5734c848a8

data/Rakefile

@@ -35,4 +35,4 @@ RuboCop::RakeTask.new do |task|
   task.options = ['-S']
 end
 
-task default: %i[spec rubocop]
+task default: %i[spec]

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

@@ -33,7 +33,7 @@ module Switchman
 
       protected
 
-      def log(*args, &block)
+      def log(...)
         super
       ensure
         @last_query_at = Time.now

data/lib/switchman/active_record/associations.rb

@@ -27,7 +27,7 @@ module Switchman
           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
+          Shard.with_each_shard(shards, [klass.connection_class_for_self, owner.class.connection_class_for_self].uniq) do
             super
           end
         end
@@ -83,10 +83,23 @@ module Switchman
 
       module Preloader
         module Association
+          module LoaderQuery
+            def load_records_in_batch(loaders)
+              # While in theory loading multiple associations that end up being effectively the same would be nice
+              # it's not very switchman compatible, so just don't bother trying to use that logic
+              # raw_records = records_for(loaders)
+
+              loaders.each do |loader|
+                loader.load_records(nil)
+                loader.run
+              end
+            end
+          end
+
           # 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
+              global_key = if model.connection_class_for_self == UnshardedRecord
                              convert_key(record[association_key_name])
                            else
                              Shard.global_id_for(record[association_key_name], record.shard)
@@ -100,13 +113,15 @@ module Switchman
           # significant changes:
           #  * partition_by_shard the records_for call
           #  * re-globalize the fetched owner id before looking up in the map
-          def load_records
+          # TODO: the ignored param currently loads records; we should probably not waste effort double-loading them
+          # Change introduced here: https://github.com/rails/rails/commit/c6c0b2e8af64509b699b782aadfecaa430700ece
+          def load_records(raw_records = nil)
             # 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 = []
+            if ::Rails.version < '7.0' && owner_keys.empty?
+              raw_records ||= []
             else
               # determine the shard to search for each owner
               if reflection.macro == :belongs_to
@@ -123,12 +138,12 @@ module Switchman
                 partition_proc = ->(owner) { owner.shard }
               end
 
-              raw_records = Shard.partition_by_shard(owners, partition_proc) do |partitioned_owners|
+              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))
+                                                Shard.current(klass.connection_class_for_self))
                   end
                   convert_key(key)
                 end
@@ -200,7 +215,7 @@ module Switchman
           # 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 }
+          shard.activate(connection_class_for_self_for_reflection(reflection)) { super }
         end
       end
     end

data/lib/switchman/active_record/attribute_methods.rb

@@ -25,6 +25,16 @@ module Switchman
           @sharded_column_values[column_name]
         end
 
+        def define_attribute_methods
+          super
+          # ensure that we're using the sharded attribute method
+          # and not the silly one in AR::AttributeMethods::PrimaryKey
+          return unless sharded_column?(@primary_key)
+
+          class_eval(build_sharded_getter('id', '_read_attribute(@primary_key)', "::#{connection_class_for_self.name}"), __FILE__, __LINE__)
+          class_eval(build_sharded_setter('id', @primary_key, "::#{connection_class_for_self.name}"), __FILE__, __LINE__)
+        end
+
         protected
 
         def reflection_for_integer_attribute(attr_name)
@@ -36,17 +46,30 @@ module Switchman
           raise if connection.open_transactions.positive?
         end
 
+        # rubocop:disable Naming/MethodParameterName
+        def define_cached_method(owner, name, namespace:, as:, &block)
+          if ::Rails.version < '7.0'
+            yield owner
+            owner.rename_method(as, name)
+          else
+            owner.define_cached_method(name, namespace: namespace, as: as, &block)
+          end
+        end
+        # rubocop:enable Naming/MethodParameterName
+
         def define_method_global_attribute(attr_name, owner:)
           if sharded_column?(attr_name)
-            owner << <<-RUBY
-              def global_#{attr_name}
-                raw_value = original_#{attr_name}
-                return nil if raw_value.nil?
-                return raw_value if raw_value > ::Switchman::Shard::IDS_PER_SHARD
-
-                ::Switchman::Shard.global_id_for(raw_value, shard)
-              end
-            RUBY
+            define_cached_method(owner, "global_#{attr_name}", as: "sharded_global_#{attr_name}", namespace: :switchman) do |batch|
+              batch << <<-RUBY
+                def sharded_global_#{attr_name}
+                  raw_value = original_#{attr_name}
+                  return nil if raw_value.nil?
+                  return raw_value if raw_value > ::Switchman::Shard::IDS_PER_SHARD
+
+                  ::Switchman::Shard.global_id_for(raw_value, shard)
+                end
+              RUBY
+            end
           else
             define_method_unsharded_column(attr_name, 'global', owner)
           end
@@ -54,113 +77,150 @@ module Switchman
 
         def define_method_local_attribute(attr_name, owner:)
           if sharded_column?(attr_name)
-            owner << <<-RUBY
-              def local_#{attr_name}
-                raw_value = original_#{attr_name}
-                return nil if raw_value.nil?
-                return raw_value % ::Switchman::Shard::IDS_PER_SHARD
-              end
-            RUBY
+            define_cached_method(owner, "local_#{attr_name}", as: "sharded_local_#{attr_name}", namespace: :switchman) do |batch|
+              batch << <<-RUBY
+                def sharded_local_#{attr_name}
+                  raw_value = original_#{attr_name}
+                  return nil if raw_value.nil?
+                  return raw_value % ::Switchman::Shard::IDS_PER_SHARD
+                end
+              RUBY
+            end
           else
             define_method_unsharded_column(attr_name, 'local', owner)
           end
         end
 
-        # see also Base#connection_classes_for_reflection
+        # see also Base#connection_class_for_self_for_reflection
         # the difference being this will output static strings for the common cases, making them
         # more performant
-        def connection_classes_code_for_reflection(reflection)
+        def connection_class_for_self_code_for_reflection(reflection)
           if reflection
             if reflection.options[:polymorphic]
               # a polymorphic association has to be discovered at runtime. This code ends up being something like
-              # context_type.&.constantize&.connection_classes
-              "read_attribute(:#{reflection.foreign_type})&.constantize&.connection_classes"
+              # context_type.&.constantize&.connection_class_for_self
+              "read_attribute(:#{reflection.foreign_type})&.constantize&.connection_class_for_self"
             else
               # otherwise we can just return a symbol for the statically known type of the association
-              "::#{reflection.klass.connection_classes.name}"
+              "::#{reflection.klass.connection_class_for_self.name}"
             end
           else
-            "::#{connection_classes.name}"
+            "::#{connection_class_for_self.name}"
           end
         end
 
-        # just a dummy class with the proper interface that calls module_eval immediately
-        class CodeGenerator
-          def initialize(mod, line)
-            @module = mod
-            @line = line
+        def define_method_attribute(attr_name, owner:)
+          if sharded_column?(attr_name)
+            reflection = reflection_for_integer_attribute(attr_name)
+            class_name = connection_class_for_self_code_for_reflection(reflection)
+            safe_class_name = class_name.unpack1('h*')
+            define_cached_method(owner, attr_name, as: "sharded_#{safe_class_name}_#{attr_name}", namespace: :switchman) do |batch|
+              batch << build_sharded_getter("sharded_#{safe_class_name}_#{attr_name}", "original_#{attr_name}", class_name)
+            end
+          else
+            define_cached_method(owner, attr_name, as: "plain_#{attr_name}", namespace: :switchman) do |batch|
+              batch << <<-RUBY
+                def plain_#{attr_name}
+                  _read_attribute("#{attr_name}") { |n| missing_attribute(n, caller) }
+                end
+              RUBY
+            end
           end
+        end
 
-          def <<(string)
-            @module.module_eval(string, __FILE__, @line)
-          end
+        def build_sharded_getter(attr_name, raw_expr, attr_connection_class)
+          <<-RUBY
+            def #{attr_name}
+              raw_value = #{raw_expr}
+              return nil if raw_value.nil?
+
+              abs_raw_value = raw_value.abs
+              current_shard = ::Switchman::Shard.current(#{attr_connection_class})
+              same_shard = shard == current_shard
+              return raw_value if same_shard && abs_raw_value < ::Switchman::Shard::IDS_PER_SHARD
+
+              value_shard_id = abs_raw_value / ::Switchman::Shard::IDS_PER_SHARD
+              # this is a stupid case when someone stuffed a global id for the current shard in instead
+              # of a local id
+              return raw_value % ::Switchman::Shard::IDS_PER_SHARD if value_shard_id == current_shard.id
+              return raw_value if !same_shard && abs_raw_value > ::Switchman::Shard::IDS_PER_SHARD
+              return shard.global_id_for(raw_value) if !same_shard && abs_raw_value < ::Switchman::Shard::IDS_PER_SHARD
+
+              ::Switchman::Shard.relative_id_for(raw_value, shard, current_shard)
+            end
+          RUBY
         end
 
-        def define_method_original_attribute(attr_name, owner:)
+        def define_method_attribute=(attr_name, owner:)
           if sharded_column?(attr_name)
             reflection = reflection_for_integer_attribute(attr_name)
-            if attr_name == 'id'
-              return if method_defined?(:original_id)
+            class_name = connection_class_for_self_code_for_reflection(reflection)
+            safe_class_name = class_name.unpack1('h*')
+            define_cached_method(owner, "#{attr_name}=", as: "sharded_#{safe_class_name}_#{attr_name}=", namespace: :switchman) do |batch|
+              batch << build_sharded_setter("sharded_#{safe_class_name}_#{attr_name}", attr_name, class_name)
+            end
+          else
+            define_cached_method(owner, "#{attr_name}=", as: "plain_#{attr_name}=", namespace: :switchman) do |batch|
+              batch << <<-RUBY
+                def plain_#{attr_name}=(new_value)
+                  _write_attribute('#{attr_name}', new_value)
+                end
+              RUBY
+            end
+          end
+        end
 
-              owner = CodeGenerator.new(self, __LINE__ + 4)
+        def build_sharded_setter(attr_name, attr_field, attr_connection_class)
+          <<-RUBY
+            def #{attr_name}=(new_value)
+              self.original_#{attr_field} = ::Switchman::Shard.relative_id_for(new_value, ::Switchman::Shard.current(#{attr_connection_class}), shard)
             end
+          RUBY
+        end
 
-            owner << <<-RUBY
-              # rename the original method to original_*
-              alias_method 'original_#{attr_name}', '#{attr_name}'
-              # and replace with one that transposes the id
-              def #{attr_name}
-                raw_value = original_#{attr_name}
-                return nil if raw_value.nil?
-
-                abs_raw_value = raw_value.abs
-                current_shard = ::Switchman::Shard.current(#{connection_classes_code_for_reflection(reflection)})
-                same_shard = shard == current_shard
-                return raw_value if same_shard && abs_raw_value < ::Switchman::Shard::IDS_PER_SHARD
-
-                value_shard_id = abs_raw_value / ::Switchman::Shard::IDS_PER_SHARD
-                # this is a stupid case when someone stuffed a global id for the current shard in instead
-                # of a local id
-                return raw_value % ::Switchman::Shard::IDS_PER_SHARD if value_shard_id == current_shard.id
-                return raw_value if !same_shard && abs_raw_value > ::Switchman::Shard::IDS_PER_SHARD
-                return shard.global_id_for(raw_value) if !same_shard && abs_raw_value < ::Switchman::Shard::IDS_PER_SHARD
-
-                ::Switchman::Shard.relative_id_for(raw_value, shard, current_shard)
-              end
+        def define_method_original_attribute(attr_name, owner:)
+          if sharded_column?(attr_name)
+            define_cached_method(owner, "original_#{attr_name}", as: "sharded_original_#{attr_name}", namespace: :switchman) do |batch|
+              batch << <<-RUBY
+                def sharded_original_#{attr_name}
+                  _read_attribute("#{attr_name}") { |n| missing_attribute(n, caller) }
+                end
+              RUBY
+            end
+          else
+            define_method_unsharded_column(attr_name, 'global', owner)
+          end
+        end
+
+        def define_method_original_attribute=(attr_name, owner:)
+          return unless sharded_column?(attr_name)
 
-              alias_method 'original_#{attr_name}=', '#{attr_name}='
-              def #{attr_name}=(new_value)
-                self.original_#{attr_name} = ::Switchman::Shard.relative_id_for(new_value, ::Switchman::Shard.current(#{connection_classes_code_for_reflection(reflection)}), shard)
+          define_cached_method(owner, "original_#{attr_name}=", as: "sharded_original_#{attr_name}=", namespace: :switchman) do |batch|
+            batch << <<-RUBY
+              def sharded_original_#{attr_name}=(new_value)
+                _write_attribute('#{attr_name}', new_value)
               end
             RUBY
-          else
-            define_method_unsharded_column(attr_name, 'global', owner)
           end
         end
 
         def define_method_unsharded_column(attr_name, prefix, owner)
-          return if columns_hash["#{prefix}_#{attr_name}"]
+          return if columns_hash["#{prefix}_#{attr_name}"] || attr_name == 'id'
 
-          owner << <<-RUBY
-            def #{prefix}_#{attr_name}
-              raise NoMethodError, "undefined method `#{prefix}_#{attr_name}'; are you missing an association?"
-            end
-          RUBY
+          define_cached_method(owner, "#{prefix}_#{attr_name}", as: "unsharded_#{prefix}_#{attr_name}", namespace: :switchman) do |batch|
+            batch << <<-RUBY
+              def unsharded_#{prefix}_#{attr_name}
+                raise NoMethodError, "undefined method `#{prefix}_#{attr_name}'; are you missing an association?"
+              end
+            RUBY
+          end
         end
       end
 
-      def self.included(klass)
-        klass.singleton_class.include(ClassMethods)
+      def self.prepended(klass)
+        klass.singleton_class.prepend(ClassMethods)
         klass.attribute_method_prefix 'global_', 'local_', 'original_'
-      end
-
-      # ensure that we're using the sharded attribute method
-      # and not the silly one in AR::AttributeMethods::PrimaryKey
-      def id
-        return super if is_a?(Shard)
-
-        self.class.define_attribute_methods
-        super
+        klass.attribute_method_affix prefix: 'original_', suffix: '='
       end
 
       # these are called if the specific methods haven't been defined yet
@@ -168,7 +228,7 @@ module Switchman
         return super unless self.class.sharded_column?(attr_name)
 
         reflection = self.class.send(:reflection_for_integer_attribute, attr_name)
-        ::Switchman::Shard.relative_id_for(super, shard, ::Switchman::Shard.current(connection_classes_for_reflection(reflection)))
+        ::Switchman::Shard.relative_id_for(super, shard, ::Switchman::Shard.current(connection_class_for_self_for_reflection(reflection)))
       end
 
       def attribute=(attr_name, new_value)
@@ -178,7 +238,7 @@ module Switchman
         end
 
         reflection = self.class.send(:reflection_for_integer_attribute, attr_name)
-        super(::Switchman::Shard.relative_id_for(new_value, ::Switchman::Shard.current(connection_classes_for_reflection(reflection)), shard))
+        super(::Switchman::Shard.relative_id_for(new_value, ::Switchman::Shard.current(connection_class_for_self_for_reflection(reflection)), shard))
       end
 
       def global_attribute(attr_name)
@@ -199,15 +259,15 @@ module Switchman
 
       private
 
-      def connection_classes_for_reflection(reflection)
+      def connection_class_for_self_for_reflection(reflection)
         if reflection
           if reflection.options[:polymorphic]
-            read_attribute(reflection.foreign_type)&.constantize&.connection_classes
+            read_attribute(reflection.foreign_type)&.constantize&.connection_class_for_self
           else
-            reflection.klass.connection_classes
+            reflection.klass.connection_class_for_self
           end
         else
-          self.class.connection_classes
+          self.class.connection_class_for_self
         end
       end
     end

data/lib/switchman/active_record/base.rb

@@ -25,11 +25,11 @@ module Switchman
         def transaction(**)
           if self != ::ActiveRecord::Base && current_scope
             current_scope.activate do
-              db = Shard.current(connection_classes).database_server
+              db = Shard.current(connection_class_for_self).database_server
               db.unguard { super }
             end
           else
-            db = Shard.current(connection_classes).database_server
+            db = Shard.current(connection_class_for_self).database_server
             db.unguard { super }
           end
         end
@@ -58,8 +58,8 @@ 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)
@@ -78,7 +78,7 @@ module Switchman
         end
 
         def connected_to_stack
-          return super if Thread.current.thread_variable?(:ar_connected_to_stack)
+          return super if ::Rails.version < '7.0' ? Thread.current.thread_variable?(:ar_connected_to_stack) : ::ActiveSupport::IsolatedExecutionState.key?(:active_record_connected_to_stack)
 
           ret = super
           DatabaseServer.guard_servers
@@ -86,13 +86,13 @@ 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)
-            if shard_role && (hash[:klasses].include?(::ActiveRecord::Base) || hash[:klasses].include?(connection_classes))
+            shard_role = hash.dig(:shard_roles, target_shard)
+            if shard_role && (hash[:klasses].include?(::ActiveRecord::Base) || hash[:klasses].include?(connection_class_for_self))
               sharded_role = shard_role
               break
             end
@@ -107,7 +107,7 @@ module Switchman
         # i.e. other sharded models don't inherit the current shard of Base
         def current_shard
           connected_to_stack.reverse_each do |hash|
-            return hash[:shard] if hash[:shard] && hash[:klasses].include?(connection_classes)
+            return hash[:shard] if hash[:shard] && hash[:klasses].include?(connection_class_for_self)
           end
 
           default_shard
@@ -115,11 +115,17 @@ module Switchman
 
         def current_switchman_shard
           connected_to_stack.reverse_each do |hash|
-            return hash[:switchman_shard] if hash[:switchman_shard] && hash[:klasses].include?(connection_classes)
+            return hash[:switchman_shard] if hash[:switchman_shard] && hash[:klasses].include?(connection_class_for_self)
           end
 
           Shard.default
         end
+
+        if ::Rails.version < '7.0'
+          def connection_class_for_self
+            connection_classes
+          end
+        end
       end
 
       def self.prepended(klass)
@@ -128,15 +134,15 @@ module Switchman
 
       def _run_initialize_callbacks
         @shard ||= if self.class.sharded_primary_key?
-                     Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_classes))
+                     Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_class_for_self))
                    else
-                     Shard.current(self.class.connection_classes)
+                     Shard.current(self.class.connection_class_for_self)
                    end
         super
       end
 
       def shard
-        @shard || Shard.current(self.class.connection_classes) || Shard.default
+        @shard || Shard.current(self.class.connection_class_for_self) || Shard.default
       end
 
       def shard=(new_shard)
@@ -161,7 +167,7 @@ module Switchman
       end
 
       def destroy
-        shard.activate(self.class.connection_classes) { super }
+        shard.activate(self.class.connection_class_for_self) { super }
       end
 
       def clone
@@ -174,14 +180,14 @@ module Switchman
       end
 
       def transaction(**kwargs, &block)
-        shard.activate(self.class.connection_classes) do
+        shard.activate(self.class.connection_class_for_self) do
           self.class.transaction(**kwargs, &block)
         end
       end
 
       def with_transaction_returning_status
-        shard.activate(self.class.connection_classes) do
-          db = Shard.current(self.class.connection_classes).database_server
+        shard.activate(self.class.connection_class_for_self) do
+          db = Shard.current(self.class.connection_class_for_self).database_server
           db.unguard { super }
         end
       end
@@ -218,22 +224,22 @@ module Switchman
 
       protected
 
-      # see also AttributeMethods#connection_classes_code_for_reflection
-      def connection_classes_for_reflection(reflection)
+      # see also AttributeMethods#connection_class_for_self_code_for_reflection
+      def connection_class_for_self_for_reflection(reflection)
         if reflection
           if reflection.options[:polymorphic]
             begin
-              read_attribute(reflection.foreign_type)&.constantize&.connection_classes || ::ActiveRecord::Base
+              read_attribute(reflection.foreign_type)&.constantize&.connection_class_for_self || ::ActiveRecord::Base
             rescue NameError
               # in case someone is abusing foreign_type to not point to an actual class
               ::ActiveRecord::Base
             end
           else
             # otherwise we can just return a symbol for the statically known type of the association
-            reflection.klass.connection_classes
+            reflection.klass.connection_class_for_self
           end
         else
-          self.class.connection_classes
+          self.class.connection_class_for_self
         end
       end
     end

data/lib/switchman/active_record/calculations.rb

@@ -4,7 +4,7 @@ module Switchman
   module ActiveRecord
     module Calculations
       def pluck(*column_names)
-        target_shard = Shard.current(klass.connection_classes)
+        target_shard = Shard.current(klass.connection_class_for_self)
         shard_count = 0
         result = activate do |relation, shard|
           shard_count += 1
@@ -109,11 +109,18 @@ module Switchman
       private
 
       def type_cast_calculated_value_switchman(value, column_name, operation)
-        type_cast_calculated_value(value, operation) do |val|
+        if ::Rails.version < '7.0'
+          type_cast_calculated_value(value, operation) do |val|
+            column = aggregate_column(column_name)
+            type ||= column.try(:type_caster) ||
+                     lookup_cast_type_from_join_dependencies(column_name.to_s) || ::ActiveRecord::Type.default_value
+            type.deserialize(val)
+          end
+        else
           column = aggregate_column(column_name)
           type ||= column.try(:type_caster) ||
-                   lookup_cast_type_from_join_dependencies(column_name.to_s) || Type.default_value
-          type.deserialize(val)
+                   lookup_cast_type_from_join_dependencies(column_name.to_s) || ::ActiveRecord::Type.default_value
+          type_cast_calculated_value(value, operation, type)
         end
       end
 

data/lib/switchman/active_record/connection_pool.rb

@@ -39,7 +39,7 @@ module Switchman
       private
 
       def current_shard
-        connection_klass.current_switchman_shard
+        ::Rails.version < '7.0' ? connection_klass.current_switchman_shard : connection_class.current_switchman_shard
       end
 
       def tls_key

data/lib/switchman/active_record/finder_methods.rb

@@ -7,7 +7,7 @@ module Switchman
         return super(id) unless klass.integral_id?
 
         if shard_source_value != :implicit
-          current_shard = Shard.current(klass.connection_classes)
+          current_shard = Shard.current(klass.connection_class_for_self)
           result = activate do |relation, shard|
             current_id = Shard.relative_id_for(id, current_shard, shard)
             # current_id will be nil for non-integral id
@@ -33,7 +33,7 @@ module Switchman
       end
 
       def find_some_ordered(ids)
-        current_shard = Shard.current(klass.connection_classes)
+        current_shard = Shard.current(klass.connection_class_for_self)
         ids = ids.map { |id| Shard.relative_id_for(id, current_shard, current_shard) }
         super(ids)
       end

data/lib/switchman/active_record/model_schema.rb

@@ -6,7 +6,7 @@ module Switchman
       module ClassMethods
         def quoted_table_name
           @quoted_table_name ||= {}
-          @quoted_table_name[Shard.current(connection_classes).id] ||= connection.quote_table_name(table_name)
+          @quoted_table_name[Shard.current(connection_class_for_self).id] ||= connection.quote_table_name(table_name)
         end
       end
     end

data/lib/switchman/active_record/persistence.rb

@@ -5,11 +5,11 @@ module Switchman
     module Persistence
       # touch reads the id attribute directly, so it's not relative to the current shard
       def touch(*, **)
-        shard.activate(self.class.connection_classes) { super }
+        shard.activate(self.class.connection_class_for_self) { super }
       end
 
       def update_columns(*)
-        shard.activate(self.class.connection_classes) { super }
+        shard.activate(self.class.connection_class_for_self) { super }
       end
 
       def delete

data/lib/switchman/active_record/postgresql_adapter.rb

@@ -7,7 +7,7 @@ module Switchman
       def create_database(name, options = {})
         options = { encoding: 'utf8' }.merge!(options.symbolize_keys)
 
-        option_string = options.sum do |key, value|
+        option_string = options.sum('') do |key, value|
           case key
           when :owner
             " OWNER = \"#{value}\""

data/lib/switchman/active_record/query_methods.rb

@@ -65,7 +65,7 @@ module Switchman
         when ::ActiveRecord::Relation
           Shard.default
         when nil
-          Shard.current(klass.connection_classes)
+          Shard.current(klass.connection_class_for_self)
         else
           raise ArgumentError, "invalid shard value #{shard_value}"
         end
@@ -79,7 +79,7 @@ module Switchman
         when ::ActiveRecord::Base
           shard_value.respond_to?(:associated_shards) ? shard_value.associated_shards : [shard_value.shard]
         when nil
-          [Shard.current(klass.connection_classes)]
+          [Shard.current(klass.connection_class_for_self)]
         else
           shard_value
         end
@@ -131,7 +131,7 @@ module Switchman
           id_shards = Set.new
           right.each do |value|
             local_id, id_shard = Shard.local_id_for(value)
-            id_shard ||= Shard.current(klass.connection_classes) if local_id
+            id_shard ||= Shard.current(klass.connection_class_for_self) if local_id
             id_shards << id_shard if id_shard
           end
           return if id_shards.empty?
@@ -151,10 +151,13 @@ module Switchman
           end
         when ::Arel::Nodes::BindParam
           local_id, id_shard = Shard.local_id_for(right.value.value_before_type_cast)
-          id_shard ||= Shard.current(klass.connection_classes) if local_id
+          id_shard ||= Shard.current(klass.connection_class_for_self) if local_id
+        when ::ActiveModel::Attribute
+          local_id, id_shard = Shard.local_id_for(right.value_before_type_cast)
+          id_shard ||= Shard.current(klass.connection_class_for_self) if local_id
         else
           local_id, id_shard = Shard.local_id_for(right)
-          id_shard ||= Shard.current(klass.connection_classes) if local_id
+          id_shard ||= Shard.current(klass.connection_class_for_self) if local_id
         end
 
         return if !id_shard || id_shard == primary_shard
@@ -193,9 +196,9 @@ module Switchman
           reflection = model.send(:reflection_for_integer_attribute, column)
           break if reflection
         end
-        return Shard.current(klass.connection_classes) if reflection.options[:polymorphic]
+        return Shard.current(klass.connection_class_for_self) if reflection.options[:polymorphic]
 
-        Shard.current(reflection.klass.connection_classes)
+        Shard.current(reflection.klass.connection_class_for_self)
       end
 
       def relation_and_column(attribute)
@@ -291,7 +294,7 @@ module Switchman
           if source_shard
             source_shard
           elsif type == :primary
-            Shard.current(klass.connection_classes)
+            Shard.current(klass.connection_class_for_self)
           elsif type == :foreign
             source_shard_for_foreign_key(relation, column)
           end
@@ -332,8 +335,9 @@ module Switchman
       end
 
       def transpose_predicate_value(value, current_shard, target_shard, attribute_type, remove_non_local_ids)
-        if value.is_a?(::Arel::Nodes::BindParam)
-          query_att = value.value
+        case value
+        when ::Arel::Nodes::BindParam, ::ActiveModel::Attribute
+          query_att = value.is_a?(::ActiveModel::Attribute) ? value : value.value
           current_id = query_att.value_before_type_cast
           if current_id.is_a?(::ActiveRecord::StatementCache::Substitute)
             current_id.sharded = true # mark for transposition later
@@ -346,7 +350,12 @@ module Switchman
               # make a new bind param
               value
             else
-              ::Arel::Nodes::BindParam.new(query_att.class.new(query_att.name, local_id, query_att.type))
+              new_att = query_att.class.new(query_att.name, local_id, query_att.type)
+              if value.is_a?(::ActiveModel::Attribute)
+                new_att
+              else
+                ::Arel::Nodes::BindParam.new(new_att)
+              end
             end
           end
         else

data/lib/switchman/active_record/reflection.rb

@@ -5,7 +5,7 @@ module Switchman
     module Reflection
       module AbstractReflection
         def shard(owner)
-          if polymorphic? || klass.connection_classes == owner.class.connection_classes
+          if polymorphic? || klass.connection_class_for_self == owner.class.connection_class_for_self
             # polymorphic associations assume the same shard as the owning item
             owner.shard
           else

data/lib/switchman/active_record/relation.rb

@@ -9,13 +9,13 @@ module Switchman
 
       def initialize(*, **)
         super
-        self.shard_value = Shard.current(klass ? klass.connection_classes : :primary) unless shard_value
+        self.shard_value = Shard.current(klass ? klass.connection_class_for_self : :primary) unless shard_value
         self.shard_source_value = :implicit unless shard_source_value
       end
 
       def clone
         result = super
-        result.shard_value = Shard.current(klass ? klass.connection_classes : :primary) unless shard_value
+        result.shard_value = Shard.current(klass ? klass.connection_class_for_self : :primary) unless shard_value
         result
       end
 
@@ -29,35 +29,32 @@ module Switchman
       end
 
       def new(*, &block)
-        primary_shard.activate(klass.connection_classes) { super }
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
       def create(*, &block)
-        primary_shard.activate(klass.connection_classes) { super }
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
       def create!(*, &block)
-        primary_shard.activate(klass.connection_classes) { super }
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
       def to_sql
-        primary_shard.activate(klass.connection_classes) { super }
+        primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
       def explain
         activate { |relation| relation.call_super(:explain, Relation) }
       end
 
-      def records
-        return @records if loaded?
-
-        results = activate { |relation| relation.call_super(:records, Relation) }
-        case shard_value
-        when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
-          @records = results
+      def load(&block)
+        if !loaded? || (::Rails.version >= '7.0' && scheduled?)
+          @records = activate { |relation| relation.send(:exec_queries, &block) }
           @loaded = true
         end
-        results
+
+        self
       end
 
       %I[update_all delete_all].each do |method|
@@ -106,19 +103,19 @@ module Switchman
       def activate(unordered: false, &block)
         shards = all_shards
         if Array === shards && shards.length == 1
-          if shards.first == DefaultShard || shards.first == Shard.current(klass.connection_classes)
+          if shards.first == DefaultShard || shards.first == Shard.current(klass.connection_class_for_self)
             yield(self, shards.first)
           else
-            shards.first.activate(klass.connection_classes) { yield(self, shards.first) }
+            shards.first.activate(klass.connection_class_for_self) { yield(self, shards.first) }
           end
         else
           result_count = 0
           can_order = false
-          result = Shard.with_each_shard(shards, [klass.connection_classes]) do
+          result = Shard.with_each_shard(shards, [klass.connection_class_for_self]) do
             # don't even query other shards if we're already past the limit
             next if limit_value && result_count >= limit_value && order_values.empty?
 
-            relation = shard(Shard.current(klass.connection_classes), :to_a)
+            relation = shard(Shard.current(klass.connection_class_for_self), :to_a)
             # do a minimal query if possible
             relation = relation.limit(limit_value - result_count) if limit_value && !result_count.zero? && order_values.empty?
 

data/lib/switchman/active_record/statement_cache.rb

@@ -33,12 +33,12 @@ module Switchman
           primary_value = params[primary_index]
           target_shard = Shard.local_id_for(primary_value)[1]
         end
-        current_shard = Shard.current(klass.connection_classes)
+        current_shard = Shard.current(klass.connection_class_for_self)
         target_shard ||= current_shard
 
         bind_values = bind_map.bind(params, current_shard, target_shard)
 
-        target_shard.activate(klass.connection_classes) do
+        target_shard.activate(klass.connection_class_for_self) do
           sql = qualified_query_builder(target_shard, klass).sql_for(bind_values, connection)
           klass.find_by_sql(sql, bind_values)
         end

data/lib/switchman/database_server.rb

@@ -150,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/engine.rb

@@ -10,7 +10,16 @@ module Switchman
 
     ::GuardRail.singleton_class.prepend(GuardRail::ClassMethods)
 
-    initializer 'switchman.active_record_patch', before: 'active_record.initialize_database' do
+    # 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
+    if ::Rails.version < '7.0'
+      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
+    end
+
+    initializer 'switchman.active_record_patch',
+                before: 'active_record.initialize_database',
+                after: (::Rails.version < '7.0' ? :initialize_dependency_mechanism : :setup_once_autoloader) do
       ::ActiveSupport.on_load(:active_record) do
         # Switchman requires postgres, so just always load the pg adapter
         require 'active_record/connection_adapters/postgresql_adapter'
@@ -19,7 +28,7 @@ module Switchman
         self.default_role = :primary
 
         prepend ActiveRecord::Base
-        include ActiveRecord::AttributeMethods
+        prepend ActiveRecord::AttributeMethods
         include ActiveRecord::Persistence
         singleton_class.prepend ActiveRecord::ModelSchema::ClassMethods
 
@@ -41,6 +50,7 @@ module Switchman
         ::ActiveRecord::Associations::CollectionProxy.include(ActiveRecord::Associations::CollectionProxy)
 
         ::ActiveRecord::Associations::Preloader::Association.prepend(ActiveRecord::Associations::Preloader::Association)
+        ::ActiveRecord::Associations::Preloader::Association::LoaderQuery.prepend(ActiveRecord::Associations::Preloader::Association::LoaderQuery) unless ::Rails.version < '7.0'
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::AbstractAdapter)
         ::ActiveRecord::ConnectionAdapters::ConnectionPool.prepend(ActiveRecord::ConnectionPool)
         ::ActiveRecord::ConnectionAdapters::AbstractAdapter.prepend(ActiveRecord::QueryCache)
@@ -96,7 +106,7 @@ module Switchman
       end
     end
 
-    initializer 'switchman.initialize_cache', before: 'initialize_cache', after: 'switchman.active_record_patch' do
+    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

data/lib/switchman/guard_rail/relation.rb

@@ -5,7 +5,7 @@ module Switchman
     module Relation
       def exec_queries(*args)
         if lock_value
-          db = Shard.current(connection_classes).database_server
+          db = Shard.current(connection_class_for_self).database_server
           db.unguard { super }
         else
           super
@@ -15,7 +15,7 @@ module Switchman
       %w[update_all delete_all].each do |method|
         class_eval <<-RUBY, __FILE__, __LINE__ + 1
           def #{method}(*args)
-            db = Shard.current(connection_classes).database_server
+            db = Shard.current(connection_class_for_self).database_server
             db.unguard { super }
           end
         RUBY

data/lib/switchman/sharded_instrumenter.rb

@@ -16,7 +16,7 @@ module Switchman
         payload[:shard] = {
           database_server_id: shard.database_server.id,
           id: shard.id,
-          env: @shard_host.pool.connection_klass&.current_role
+          env: ::Rails.version < '7.0' ? @shard_host.pool.connection_klass&.current_role : @shard_host.pool.connection_class&.current_role
         }
       end
       super name, payload

data/lib/switchman/version.rb

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

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: switchman
 version: !ruby/object:Gem::Version
-  version: 3.0.22
+  version: 3.1.0
 platform: ruby
 authors:
 - Cody Cutrer
@@ -10,7 +10,7 @@ authors:
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2022-04-07 00:00:00.000000000 Z
+date: 2022-06-02 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activerecord
@@ -21,7 +21,7 @@ dependencies:
         version: 6.1.4
     - - "<"
       - !ruby/object:Gem::Version
-        version: '6.2'
+        version: '7.1'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
@@ -31,21 +31,21 @@ dependencies:
         version: 6.1.4
     - - "<"
       - !ruby/object:Gem::Version
-        version: '6.2'
+        version: '7.1'
 - !ruby/object:Gem::Dependency
   name: guardrail
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 3.0.0
+        version: 3.0.1
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: 3.0.0
+        version: 3.0.1
 - !ruby/object:Gem::Dependency
   name: parallel
   requirement: !ruby/object:Gem::Requirement
@@ -69,7 +69,7 @@ dependencies:
         version: '6.1'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '6.2'
+        version: '7.1'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
@@ -79,7 +79,7 @@ dependencies:
         version: '6.1'
     - - "<"
       - !ruby/object:Gem::Version
-        version: '6.2'
+        version: '7.1'
 - !ruby/object:Gem::Dependency
   name: appraisal
   requirement: !ruby/object:Gem::Requirement
@@ -308,14 +308,14 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '2.6'
+      version: '2.7'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.1.4
+rubygems_version: 3.1.6
 signing_key:
 specification_version: 4
 summary: Rails sharding magic