switchman 3.0.1 → 4.2.5

data/lib/switchman/active_record/attribute_methods.rb

@@ -25,6 +25,25 @@ module Switchman
           @sharded_column_values[column_name]
         end
 
+        def define_attribute_methods
+          result = super
+          # ensure that we're using the sharded attribute method
+          # and not the silly one in AR::AttributeMethods::PrimaryKey
+          return result 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__)
+          result
+        end
+
         protected
 
         def reflection_for_integer_attribute(attr_name)
@@ -36,110 +55,209 @@ module Switchman
           raise if connection.open_transactions.positive?
         end
 
-        def define_method_global_attribute(attr_name, owner:)
+        def define_cached_method(owner, name, namespace:, as:, &)
+          if ::Rails.version < "7.1.4"
+            # https://github.com/rails/rails/commit/a2a12fc2e3f4e6d06f81d4c74c88f8e6b3369ee6#diff-5b59ece6d9396b596f06271cec0ea726e3360911383511c49b1a66f454bfc2b6L30
+            # These arguments were effectively swapped in Rails 7.1.4, so previous versions need them reversed
+            owner.define_cached_method(as, namespace:, as: name, &)
+          else
+            owner.define_cached_method(name, namespace:, as:, &)
+          end
+        end
+
+        def define_method_global_attribute(attr_name, owner:, as: attr_name)
           if sharded_column?(attr_name)
-            owner << <<-RUBY
-              def global_#{attr_name}
-                ::Switchman::Shard.global_id_for(original_#{attr_name}, shard)
-              end
-            RUBY
+            define_cached_method(owner,
+                                 "sharded_global_#{attr_name}",
+                                 as: "global_#{as}",
+                                 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)
+            define_method_unsharded_column(attr_name, "global", owner)
           end
         end
 
-        def define_method_local_attribute(attr_name, owner:)
+        def define_method_local_attribute(attr_name, owner:, as: attr_name)
           if sharded_column?(attr_name)
-            owner << <<-RUBY
-              def local_#{attr_name}
-                ::Switchman::Shard.local_id_for(original_#{attr_name}).first
-              end
-            RUBY
+            define_cached_method(owner,
+                                 "sharded_local_#{attr_name}",
+                                 as: "local_#{as}",
+                                 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)
+            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
+              <<~RUBY
+                begin
+                  read_attribute(:#{reflection.foreign_type})&.constantize&.connection_class_for_self
+                rescue NameError
+                  ::ActiveRecord::Base
+                end
+              RUBY
             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:, as: attr_name)
+          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,
+                                 "sharded_#{safe_class_name}_#{attr_name}",
+                                 as:,
+                                 namespace: :switchman) do |batch|
+              batch << build_sharded_getter("sharded_#{safe_class_name}_#{attr_name}",
+                                            "original_#{as}",
+                                            class_name)
+            end
+          else
+            define_cached_method(owner, "plain_#{attr_name}", as:, 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 = loaded_from_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 loaded_from_shard.global_id_for(raw_value) if !same_shard && abs_raw_value < ::Switchman::Shard::IDS_PER_SHARD
+
+              ::Switchman::Shard.relative_id_for(raw_value, loaded_from_shard, current_shard)
+            end
+          RUBY
         end
 
-        def define_method_original_attribute(attr_name, owner:)
+        def define_method_attribute=(attr_name, owner:, as: attr_name)
           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,
+                                 "sharded_#{safe_class_name}_#{attr_name}=",
+                                 as: "#{as}=",
+                                 namespace: :switchman) do |batch|
+              batch << build_sharded_setter("sharded_#{safe_class_name}_#{attr_name}", attr_name, class_name)
+            end
+          else
+            define_cached_method(owner, "plain_#{attr_name}=", as: "#{as}=", 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}), loaded_from_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}
-                ::Switchman::Shard.relative_id_for(original_#{attr_name}, shard, ::Switchman::Shard.current(#{connection_classes_code_for_reflection(reflection)}))
-              end
+        def define_method_original_attribute(attr_name, owner:, as: attr_name)
+          if sharded_column?(attr_name)
+            define_cached_method(owner,
+                                 "sharded_original_#{attr_name}",
+                                 as: "original_#{as}",
+                                 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
 
-              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)
+        def define_method_original_attribute=(attr_name, owner:, as: attr_name)
+          return unless sharded_column?(attr_name)
+
+          define_cached_method(owner,
+                               "sharded_original_#{attr_name}=",
+                               as: "original_#{as}=",
+                               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,
+                               "unsharded_#{prefix}_#{attr_name}",
+                               as: "#{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)
-        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
+      def self.prepended(klass)
+        klass.singleton_class.prepend(ClassMethods)
+        klass.attribute_method_prefix "global_", "local_", "original_"
+        klass.attribute_method_affix prefix: "original_", suffix: "="
       end
 
       # these are called if the specific methods haven't been defined yet
@@ -147,7 +265,11 @@ 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)
@@ -157,7 +279,11 @@ 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(attr_name, ::Switchman::Shard.relative_id_for(
+          new_value,
+          ::Switchman::Shard.current(connection_class_for_self_for_reflection(reflection)),
+          shard
+        ))
       end
 
       def global_attribute(attr_name)
@@ -170,25 +296,11 @@ module Switchman
 
       def local_attribute(attr_name)
         if self.class.sharded_column?(attr_name)
-          ::Switchman::Shard.local_id_for(attribute(attr_name), shard).first
+          ::Switchman::Shard.local_id_for(attribute(attr_name)).first
         else
           attribute(attr_name)
         end
       end
-
-      private
-
-      def connection_classes_for_reflection(reflection)
-        if reflection
-          if reflection.options[:polymorphic]
-            read_attribute(reflection.foreign_type)&.constantize&.connection_classes
-          else
-            reflection.klass.connection_classes
-          end
-        else
-          self.class.connection_classes
-        end
-      end
     end
   end
 end