switchman 3.0.6 → 3.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 81aeec21e056026642c6aa15bda349e97c8194212f8b79222b4581d5cb2a9909
-  data.tar.gz: 77bfd14eaf245a399ec6a5b48a9f3106daac6df393d8c55d7dad5eae56966f15
+  metadata.gz: 48b27124d5233adf585aed105d0b2522c76249bdd1e13fbcfd7e904241d3a624
+  data.tar.gz: fea1ccc0ad63f00633fb90bee1d7bd4606f6de0e18564250865c1b1dde61ef8d
 SHA512:
-  metadata.gz: bfc71a265b6aea2a48350bcce77dcdcb24cea206c2dbcd4b4e3e113f0615917b314233da2038197d394d2ae961dbc8be18f12d48ccc240d57ca96eb39d4db9e6
-  data.tar.gz: d392e78986ef050800ad61e064c702d3cb3dc98b97e48cbf42c9a973ae52226672a33ade5944119be7365f9d7f2f17e05097e0e1f2a46b1adab27dfe5169bf87
+  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

@@ -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
@@ -35,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

@@ -0,0 +1,223 @@
+# 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_class_for_self, owner.class.connection_class_for_self].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
+          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_class_for_self == 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
+          # 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 ::Rails.version < '7.0' && 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_class_for_self))
+                  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_class_for_self_for_reflection(reflection)) { super }
+        end
+      end
+    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