switchman 2.2.2 → 3.0.0

data/lib/switchman/active_record/query_methods.rb

@@ -17,15 +17,20 @@ module Switchman
       def shard_value
         @values[:shard]
       end
+
       def shard_source_value
         @values[:shard_source]
       end
+
       def shard_value=(value)
         raise ::ActiveRecord::ImmutableRelation if @loaded
+
         @values[:shard] = value
       end
+
       def shard_source_value=(value)
         raise ::ActiveRecord::ImmutableRelation if @loaded
+
         @values[:shard_source] = value
       end
 
@@ -35,11 +40,13 @@ module Switchman
 
       def shard!(value, source = :explicit)
         raise ArgumentError, "shard can't be nil" unless value
-        old_primary_shard = self.primary_shard
+
+        old_primary_shard = primary_shard
         self.shard_value = value
         self.shard_source_value = source
-        if (old_primary_shard != self.primary_shard || source == :to_a)
-          transpose_clauses(old_primary_shard, self.primary_shard, source == :to_a)
+        if old_primary_shard != primary_shard || source == :to_a
+          transpose_clauses(old_primary_shard, primary_shard,
+                            remove_nonlocal_primary_keys: source == :to_a)
         end
         self
       end
@@ -58,7 +65,7 @@ module Switchman
         when ::ActiveRecord::Relation
           Shard.default
         when nil
-          Shard.current(klass.shard_category)
+          Shard.current(klass.connection_classes)
         else
           raise ArgumentError, "invalid shard value #{shard_value}"
         end
@@ -72,114 +79,88 @@ module Switchman
         when ::ActiveRecord::Base
           shard_value.respond_to?(:associated_shards) ? shard_value.associated_shards : [shard_value.shard]
         when nil
-          [Shard.current(klass.shard_category)]
+          [Shard.current(klass.connection_classes)]
         else
           shard_value
         end
       end
 
+      def or(other)
+        super(other.shard(primary_shard))
+      end
+
       private
 
-      if ::Rails.version >= '5.2'
-        [:where, :having].each do |type|
-          class_eval <<-RUBY, __FILE__, __LINE__ + 1
-          def transpose_#{type}_clauses(source_shard, target_shard, remove_nonlocal_primary_keys)
-            unless (predicates = #{type}_clause.send(:predicates)).empty?
-              new_predicates, _binds = transpose_predicates(predicates, source_shard,
-                                                               target_shard, remove_nonlocal_primary_keys)
+      %i[where having].each do |type|
+        class_eval <<-RUBY, __FILE__, __LINE__ + 1
+        def transpose_#{type}_clauses(source_shard, target_shard, remove_nonlocal_primary_keys:)
+          unless (predicates = #{type}_clause.send(:predicates)).empty?
+            new_predicates = transpose_predicates(predicates, source_shard,
+                                                              target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
+            if new_predicates != predicates
+              self.#{type}_clause = #{type}_clause.dup
               if new_predicates != predicates
-                self.#{type}_clause = #{type}_clause.dup
-                if new_predicates != predicates
-                  #{type}_clause.instance_variable_set(:@predicates, new_predicates)
-                end
+                #{type}_clause.instance_variable_set(:@predicates, new_predicates)
               end
             end
           end
-          RUBY
-        end
-      else
-        [:where, :having].each do |type|
-          class_eval <<-RUBY, __FILE__, __LINE__ + 1
-            def transpose_#{type}_clauses(source_shard, target_shard, remove_nonlocal_primary_keys)
-              unless (predicates = #{type}_clause.send(:predicates)).empty?
-                new_predicates, new_binds = transpose_predicates(predicates, source_shard,
-                                                                 target_shard, remove_nonlocal_primary_keys,
-                                                                 binds: #{type}_clause.binds,
-                                                                 dup_binds_on_mutation: true)
-                if new_predicates != predicates || !new_binds.equal?(#{type}_clause.binds)
-                  self.#{type}_clause = #{type}_clause.dup
-                  if new_predicates != predicates
-                    #{type}_clause.instance_variable_set(:@predicates, new_predicates)
-                  end
-                  if !new_binds.equal?(#{type}_clause.binds)
-                    #{type}_clause.instance_variable_set(:@binds, new_binds)
-                  end
-                end
-              end
-            end
-          RUBY
         end
+        RUBY
       end
 
-      def transpose_clauses(source_shard, target_shard, remove_nonlocal_primary_keys = false)
-        transpose_where_clauses(source_shard, target_shard, remove_nonlocal_primary_keys)
-        transpose_having_clauses(source_shard, target_shard, remove_nonlocal_primary_keys)
+      def transpose_clauses(source_shard, target_shard, remove_nonlocal_primary_keys: false)
+        transpose_where_clauses(source_shard, target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
+        transpose_having_clauses(source_shard, target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
       end
 
-      def infer_shards_from_primary_key(predicates, binds = nil)
+      def infer_shards_from_primary_key(predicates)
         return unless klass.integral_id?
 
         primary_key = predicates.detect do |predicate|
-          predicate.is_a?(::Arel::Nodes::Binary) && predicate.left.is_a?(::Arel::Attributes::Attribute) &&
-            predicate.left.relation.is_a?(::Arel::Table) && predicate.left.relation.model == klass &&
+          (predicate.is_a?(::Arel::Nodes::Binary) || predicate.is_a?(::Arel::Nodes::HomogeneousIn)) &&
+            predicate.left.is_a?(::Arel::Attributes::Attribute) &&
+            predicate.left.relation.is_a?(::Arel::Table) && predicate.left.relation.klass == klass &&
             klass.primary_key == predicate.left.name
         end
-        if primary_key
-          case primary_key.right
-          when Array
-            id_shards = Set.new
-            primary_key.right.each do |value|
-              local_id, id_shard = Shard.local_id_for(value)
-              id_shard ||= Shard.current(klass.shard_category) if local_id
-              id_shards << id_shard if id_shard
-            end
-            if id_shards.empty?
-              return
-            elsif id_shards.length == 1
-              id_shard = id_shards.first
-            # prefer to not change the shard
-            elsif id_shards.include?(primary_shard)
-              id_shards.delete(primary_shard)
-              self.shard_value = [primary_shard] + id_shards.to_a
-              return
-            else
-              id_shards = id_shards.to_a
-              transpose_clauses(primary_shard, id_shards.first)
-              self.shard_value = id_shards
-              return
-            end
-          when ::Arel::Nodes::BindParam
-            if ::Rails.version >= "5.2"
-              local_id, id_shard = Shard.local_id_for(primary_key.right.value.value_before_type_cast)
-              id_shard ||= Shard.current(klass.shard_category) if local_id
-            else
-              # look for a bind param with a matching column name
-              if binds && bind = binds.detect{|b| b&.name.to_s == klass.primary_key.to_s}
-                unless bind.value.is_a?(::ActiveRecord::StatementCache::Substitute)
-                  local_id, id_shard = Shard.local_id_for(bind.value)
-                  id_shard ||= Shard.current(klass.shard_category) if local_id
-                end
-              end
-            end
-          else
-            local_id, id_shard = Shard.local_id_for(primary_key.right)
-            id_shard ||= Shard.current(klass.shard_category) if local_id
+        return unless primary_key
+
+        right = primary_key.is_a?(::Arel::Nodes::HomogeneousIn) ? primary_key.values : primary_key.right
+
+        case right
+        when Array
+          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_shards << id_shard if id_shard
           end
+          return if id_shards.empty?
 
-          return if !id_shard || id_shard == primary_shard
-          transpose_clauses(primary_shard, id_shard)
-          self.shard_value = id_shard
+          if id_shards.length == 1
+            id_shard = id_shards.first
+          # prefer to not change the shard
+          elsif id_shards.include?(primary_shard)
+            id_shards.delete(primary_shard)
+            self.shard_value = [primary_shard] + id_shards.to_a
+            return
+          else
+            id_shards = id_shards.to_a
+            transpose_clauses(primary_shard, id_shards.first)
+            self.shard_value = id_shards
+            return
+          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
+        else
+          local_id, id_shard = Shard.local_id_for(right)
+          id_shard ||= Shard.current(klass.connection_classes) if local_id
         end
+
+        return if !id_shard || id_shard == primary_shard
+
+        transpose_clauses(primary_shard, id_shard)
+        self.shard_value = id_shard
       end
 
       def transposable_attribute_type(relation, column)
@@ -201,8 +182,9 @@ module Switchman
 
       def sharded_primary_key?(relation, column)
         column = column.to_s
-        return column == 'id' if relation.model == ::ActiveRecord::Base
-        relation.model.primary_key == column && relation.model.integral_id?
+        return column == 'id' if relation.klass == ::ActiveRecord::Base
+
+        relation.klass.primary_key == column && relation.klass.integral_id?
       end
 
       def source_shard_for_foreign_key(relation, column)
@@ -211,8 +193,9 @@ module Switchman
           reflection = model.send(:reflection_for_integer_attribute, column)
           break if reflection
         end
-        return Shard.current(klass.shard_category) if reflection.options[:polymorphic]
-        Shard.current(reflection.klass.shard_category)
+        return Shard.current(klass.connection_classes) if reflection.options[:polymorphic]
+
+        Shard.current(reflection.klass.connection_classes)
       end
 
       def relation_and_column(attribute)
@@ -221,8 +204,31 @@ module Switchman
         [attribute.relation, column]
       end
 
-      def where_clause_factory
-        super.tap { |factory| factory.scope = self }
+      def build_where_clause(opts, rest = [])
+        opts = sanitize_forbidden_attributes(opts)
+
+        case opts
+        when String, Array
+          values = Hash === rest.first ? rest.first.values : rest
+
+          values.grep(ActiveRecord::Relation) do |rel|
+            # serialize subqueries against the same shard as the outer query is currently
+            # targeted to run against
+            rel.shard!(primary_shard) if rel.shard_source_value == :implicit && rel.primary_shard != primary_shard
+          end
+
+          super
+        when Hash, ::Arel::Nodes::Node
+          where_clause = super
+
+          predicates = where_clause.send(:predicates)
+          infer_shards_from_primary_key(predicates) if shard_source_value == :implicit && shard_value.is_a?(Shard)
+          predicates = transpose_predicates(predicates, nil, primary_shard)
+          where_clause.instance_variable_set(:@predicates, predicates)
+          where_clause
+        else
+          super
+        end
       end
 
       def arel_columns(columns)
@@ -233,117 +239,86 @@ module Switchman
         connection.with_local_table_name { super }
       end
 
-      def table_name_matches?(from)
-        connection.with_global_table_name { super }
-      end
-
-      # semi-private
-      public
       def transpose_predicates(predicates,
                                source_shard,
                                target_shard,
-                               remove_nonlocal_primary_keys = false,
-                               binds: nil,
-                               dup_binds_on_mutation: false)
-        result = predicates.map do |predicate|
-          if ::Rails.version >= '5.2' && predicate.is_a?(::Arel::Nodes::And)
-            new_predicates, _binds = transpose_predicates(predicate.children, source_shard, target_shard,
-              remove_nonlocal_primary_keys,
-              binds: binds,
-              dup_binds_on_mutation: dup_binds_on_mutation)
-            next (if new_predicates == predicate.children
-              predicate
-            else
-              ::Arel::Nodes::And.new(new_predicates)
-            end)
-          end
+                               remove_nonlocal_primary_keys: false)
+        predicates.map do |predicate|
+          transpose_single_predicate(predicate, source_shard, target_shard,
+                                     remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
+        end
+      end
 
-          next predicate unless predicate.is_a?(::Arel::Nodes::Binary)
-          next predicate unless predicate.left.is_a?(::Arel::Attributes::Attribute)
-          relation, column = relation_and_column(predicate.left)
-          next predicate unless (type = transposable_attribute_type(relation, column))
-
-          remove = true if type == :primary &&
-              remove_nonlocal_primary_keys &&
-              predicate.left.relation.model == klass &&
-              (predicate.is_a?(::Arel::Nodes::Equality) || predicate.is_a?(::Arel::Nodes::In))
-
-          current_source_shard =
-              if source_shard
-                source_shard
-              elsif type == :primary
-                Shard.current(klass.shard_category)
-              elsif type == :foreign
-                source_shard_for_foreign_key(relation, column)
-              end
+      def transpose_single_predicate(predicate,
+                                     source_shard,
+                                     target_shard,
+                                     remove_nonlocal_primary_keys: false)
+        if predicate.is_a?(::Arel::Nodes::Grouping)
+          return predicate unless predicate.expr.is_a?(::Arel::Nodes::Or)
 
-          if ::Rails.version >= "5.2"
-            new_right_value =
-              case predicate.right
-              when Array
-                predicate.right.map {|val| transpose_predicate_value(val, current_source_shard, target_shard, type, remove).presence }.compact
-              else
-                transpose_predicate_value(predicate.right, current_source_shard, target_shard, type, remove)
-              end
-          else
-            new_right_value = case predicate.right
-            when Array
-              local_ids = []
-              predicate.right.each do |value|
-                local_id = Shard.relative_id_for(value, current_source_shard, target_shard)
-                next unless local_id
-                unless remove && local_id > Shard::IDS_PER_SHARD
-                  if value.is_a?(::Arel::Nodes::Casted)
-                    if local_id == value.val
-                      local_id = value
-                    elsif local_id != value
-                      local_id = value.class.new(local_id, value.attribute)
-                    end
-                  end
-                  local_ids << local_id
-                end
-              end
-              local_ids
-            when ::Arel::Nodes::BindParam
-              # look for a bind param with a matching column name
-              if binds && bind = binds.detect{|b| b&.name.to_s == predicate.left.name.to_s}
-                # before we mutate, dup
-                if dup_binds_on_mutation
-                  binds = binds.map(&:dup)
-                  dup_binds_on_mutation = false
-                  bind = binds.find { |b| b&.name.to_s == predicate.left.name.to_s }
-                end
-                if bind.value.is_a?(::ActiveRecord::StatementCache::Substitute)
-                  bind.value.sharded = true # mark for transposition later
-                  bind.value.primary = true if type == :primary
+          # Dang, we have an OR.  OK, that means we have other epxressions below this
+          # level, perhaps many, that may need transposition.
+          # the left side and right side must each be treated as predicate lists and
+          # transformed in kind, if neither of them changes we can just return the grouping as is.
+          # hold on, it's about to get recursive...
+          or_expr = predicate.expr
+          left_node = or_expr.left
+          right_node = or_expr.right
+          new_left_predicates = transpose_single_predicate(left_node, source_shard,
+                                                           target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
+          or_expr.instance_variable_set(:@left, new_left_predicates) if new_left_predicates != left_node
+          new_right_predicates = transpose_single_predicate(right_node, source_shard,
+                                                            target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
+          or_expr.instance_variable_set(:@right, new_right_predicates) if new_right_predicates != right_node
+          return predicate
+        end
+        return predicate unless predicate.is_a?(::Arel::Nodes::Binary) || predicate.is_a?(::Arel::Nodes::HomogeneousIn)
+        return predicate unless predicate.left.is_a?(::Arel::Attributes::Attribute)
+
+        relation, column = relation_and_column(predicate.left)
+        return predicate unless (type = transposable_attribute_type(relation, column))
+
+        remove = true if type == :primary &&
+                         remove_nonlocal_primary_keys &&
+                         predicate.left.relation.klass == klass &&
+                         predicate.is_a?(::Arel::Nodes::Equality)
+
+        current_source_shard =
+          if source_shard
+            source_shard
+          elsif type == :primary
+            Shard.current(klass.connection_classes)
+          elsif type == :foreign
+            source_shard_for_foreign_key(relation, column)
+          end
+
+        right = if predicate.is_a?(::Arel::Nodes::HomogeneousIn)
+                  predicate.values
                 else
-                  local_id = Shard.relative_id_for(bind.value, current_source_shard, target_shard)
-                  local_id = [] if remove && local_id > Shard::IDS_PER_SHARD
-                  bind.instance_variable_set(:@value, local_id)
-                  bind.instance_variable_set(:@value_for_database, nil)
+                  predicate.right
                 end
-              end
-              predicate.right
-            else
-              local_id = Shard.relative_id_for(predicate.right, current_source_shard, target_shard) || predicate.right
-              local_id = [] if remove && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
-              local_id
-            end
+
+        new_right_value =
+          case right
+          when Array
+            right.map { |val| transpose_predicate_value(val, current_source_shard, target_shard, type, remove) }
+          else
+            transpose_predicate_value(right, current_source_shard, target_shard, type, remove)
           end
-          if new_right_value == predicate.right
+
+        if new_right_value == right
+          predicate
+        elsif predicate.right.is_a?(::Arel::Nodes::Casted)
+          if new_right_value == right.value
             predicate
-          elsif predicate.right.is_a?(::Arel::Nodes::Casted)
-            if new_right_value == predicate.right.val
-              predicate
-            else
-              predicate.class.new(predicate.left, predicate.right.class.new(new_right_value, predicate.right.attribute))
-            end
           else
-            predicate.class.new(predicate.left, new_right_value)
+            predicate.class.new(predicate.left, right.class.new(new_right_value, right.attribute))
           end
+        elsif predicate.is_a?(::Arel::Nodes::HomogeneousIn)
+          predicate.class.new(new_right_value, predicate.attribute, predicate.type)
+        else
+          predicate.class.new(predicate.left, new_right_value)
         end
-        result = [result, binds]
-        result
       end
 
       def transpose_predicate_value(value, current_shard, target_shard, attribute_type, remove_non_local_ids)
@@ -356,12 +331,12 @@ module Switchman
             value
           else
             local_id = Shard.relative_id_for(current_id, current_shard, target_shard) || current_id
-            return nil if remove_non_local_ids && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
-            if current_id != local_id
+            local_id = [] if remove_non_local_ids && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
+            if current_id == local_id
               # make a new bind param
-              ::Arel::Nodes::BindParam.new(query_att.class.new(query_att.name, local_id, query_att.type))
-            else
               value
+            else
+              ::Arel::Nodes::BindParam.new(query_att.class.new(query_att.name, local_id, query_att.type))
             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.shard_category == owner.class.shard_category
+          if polymorphic? || klass.connection_classes == owner.class.connection_classes
             # polymorphic associations assume the same shard as the owning item
             owner.shard
           else
@@ -28,32 +28,17 @@ module Switchman
         # this technically belongs on AssociationReflection, but we put it on
         # ThroughReflection as well, instead of delegating to its internal
         # HasManyAssociation, losing its proper `klass`
-        if ::Rails.version < '6.0.4'
-          def association_scope_cache(conn, owner, &block)
-            key = conn.prepared_statements
-            if polymorphic?
-              key = [key, owner._read_attribute(@foreign_type)]
-            end
-            key = [key, shard(owner).id].flatten
-            @association_scope_cache[key] ||= @scope_lock.synchronize {
-              @association_scope_cache[key] ||= (::Rails.version >= "5.2" ? ::ActiveRecord::StatementCache.create(conn, &block) : block.call)
-            }
-          end
-        else
-          def association_scope_cache(klass, owner, &block)
-            key = self
-            if polymorphic?
-              key = [key, owner._read_attribute(@foreign_type)]
-            end
-            key = [key, shard(owner).id].flatten
-            klass.cached_find_by_statement(key, &block)
-          end
+        def association_scope_cache(klass, owner, &block)
+          key = self
+          key = [key, owner._read_attribute(@foreign_type)] if polymorphic?
+          key = [key, shard(owner).id].flatten
+          klass.cached_find_by_statement(key, &block)
         end
       end
 
       module AssociationReflection
         def join_id_for(owner)
-          owner.send(::Rails.version >= "5.2" ? join_foreign_key : active_record_primary_key) # use sharded id values in association binds
+          owner.send(join_foreign_key) # use sharded id values in association binds
         end
       end
     end