switchman 3.0.2 → 4.2.5

data/lib/switchman/active_record/query_methods.rb

@@ -3,6 +3,19 @@
 module Switchman
   module ActiveRecord
     module QueryMethods
+      # Use this class to prevent a value from getting transposed across shards
+      class NonTransposingValue < SimpleDelegator
+        def class
+          __getobj__.class
+        end
+
+        def is_a?(other)
+          return true if other == NonTransposingValue
+
+          __getobj__.is_a?(other)
+        end
+      end
+
       # shard_value is one of:
       #   A shard
       #   An array or relation of shards
@@ -12,8 +25,6 @@ module Switchman
       #   :explicit    - explicit set on the relation
       #   :association - a special value that scopes from associations use to use slightly different logic
       #                  for foreign key transposition
-      #   :to_a        - a special value that Relation#to_a uses when querying multiple shards to
-      #                  remove primary keys from conditions that aren't applicable to the current shard
       def shard_value
         @values[:shard]
       end
@@ -23,13 +34,29 @@ module Switchman
       end
 
       def shard_value=(value)
-        raise ::ActiveRecord::ImmutableRelation if @loaded
+        if @loaded
+          error_class = if ::Rails.version < "7.2"
+                          ::ActiveRecord::ImmutableRelation
+                        else
+                          ::ActiveRecord::UnmodifiableRelation
+                        end
+
+          raise error_class
+        end
 
         @values[:shard] = value
       end
 
       def shard_source_value=(value)
-        raise ::ActiveRecord::ImmutableRelation if @loaded
+        if @loaded
+          error_class = if ::Rails.version < "7.2"
+                          ::ActiveRecord::ImmutableRelation
+                        else
+                          ::ActiveRecord::UnmodifiableRelation
+                        end
+
+          raise error_class
+        end
 
         @values[:shard_source] = value
       end
@@ -44,10 +71,7 @@ module Switchman
         old_primary_shard = primary_shard
         self.shard_value = value
         self.shard_source_value = source
-        if old_primary_shard != primary_shard || source == :to_a
-          transpose_clauses(old_primary_shard, primary_shard,
-                            remove_nonlocal_primary_keys: source == :to_a)
-        end
+        transpose_predicates(nil, old_primary_shard, primary_shard) if old_primary_shard != primary_shard
         self
       end
 
@@ -65,7 +89,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 +103,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
@@ -89,35 +113,41 @@ module Switchman
         super(other.shard(primary_shard))
       end
 
-      private
+      # use a temp variable so that the new where clause is built before self.where_clause is read,
+      # since build_where_clause might mutate self.where_clause
+      def where!(opts, *rest)
+        new_clause = build_where_clause(opts, rest)
+        self.where_clause += new_clause
+        self
+      end
 
-      %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
-                #{type}_clause.instance_variable_set(:@predicates, new_predicates)
-              end
-            end
-          end
-        end
-        RUBY
+      protected
+
+      def arel_columns(columns)
+        connection.with_local_table_name { super }
       end
 
-      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)
+      def remove_nonlocal_primary_keys!
+        each_transposable_predicate_value do |value, predicate, _relation, _column, type|
+          next value unless
+            type == :primary &&
+            predicate.left.relation.klass == klass &&
+            (predicate.is_a?(::Arel::Nodes::Equality) || predicate.is_a?(::Arel::Nodes::HomogeneousIn))
+
+          (value.is_a?(Integer) && value > Shard::IDS_PER_SHARD) ? [] : value
+        end
+        self
       end
 
+      private
+
       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.is_a?(::Arel::Nodes::HomogeneousIn)) &&
+          (predicate.is_a?(::Arel::Nodes::Equality) ||
+           predicate.is_a?(::Arel::Nodes::In) ||
+           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
@@ -131,7 +161,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?
@@ -145,21 +175,24 @@ module Switchman
             return
           else
             id_shards = id_shards.to_a
-            transpose_clauses(primary_shard, id_shards.first)
+            transpose_predicates(nil, 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
+          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
 
-        transpose_clauses(primary_shard, id_shard)
+        transpose_predicates(nil, primary_shard, id_shard)
         self.shard_value = id_shard
       end
 
@@ -177,25 +210,25 @@ module Switchman
       end
 
       def sharded_foreign_key?(relation, column)
-        models_for_table(relation.table_name).any? { |m| m.sharded_column?(column) }
+        models_for_table(relation.name).any? { |m| m.sharded_column?(column) }
       end
 
       def sharded_primary_key?(relation, column)
         column = column.to_s
-        return column == 'id' if relation.klass == ::ActiveRecord::Base
+        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)
         reflection = nil
-        models_for_table(relation.table_name).each do |model|
+        models_for_table(relation.name).each do |model|
           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)
@@ -209,12 +242,11 @@ module Switchman
 
         case opts
         when String, Array
-          values = Hash === rest.first ? rest.first.values : rest
+          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
+          if shard_source_value != :explicit && values.grep(ActiveRecord::Relation).first
+            raise "Sub-queries are not allowed as simple substitutions; " \
+                  "please build your relation with more structured methods so that Switchman is able to introspect it."
           end
 
           super
@@ -231,124 +263,203 @@ module Switchman
         end
       end
 
-      def arel_columns(columns)
-        connection.with_local_table_name { super }
-      end
-
       def arel_column(columns)
         connection.with_local_table_name { super }
       end
 
-      def transpose_predicates(predicates,
-                               source_shard,
-                               target_shard,
-                               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)
+      def table_name_matches?(from)
+        if ::Rails.version < "7.2"
+          connection.with_global_table_name { super }
+        else
+          connection.with_global_table_name do
+            table_name = Regexp.escape(table.name)
+            # INST: adapter_class -> connection
+            quoted_table_name = Regexp.escape(connection.quote_table_name(table.name))
+            /(?:\A|(?<!FROM)\s)(?:\b#{table_name}\b|#{quoted_table_name})(?!\.)/i.match?(from.to_s)
+          end
         end
       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)
-
-          # 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) || predicate.is_a?(::Arel::Nodes::HomogeneousIn))
-
-        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)
+      unless ::Rails.version < "7.2"
+        def order_column(field)
+          arel_column(field) do |attr_name|
+            if attr_name == "count" && !group_values.empty?
+              table[attr_name]
+            else
+              # INST: adapter_class -> connection
+              ::Arel.sql(connection.quote_table_name(attr_name), retryable: true)
+            end
           end
+        end
+      end
 
-        right = if predicate.is_a?(::Arel::Nodes::HomogeneousIn)
-                  predicate.values
-                else
-                  predicate.right
-                end
+      def each_predicate(predicates = nil, &)
+        return predicates.map(&) if predicates
 
-        new_right_value =
-          case right
-          when Array
-            right.map { |val| transpose_predicate_value(val, current_source_shard, target_shard, type, remove).presence }.compact
-          else
-            transpose_predicate_value(right, current_source_shard, target_shard, type, remove)
+        each_predicate_cb(:having_clause, :having_clause=, &)
+        each_predicate_cb(:where_clause, :where_clause=, &)
+      end
+
+      def each_predicate_cb(clause_getter, clause_setter, &)
+        old_clause = send(clause_getter)
+        old_predicates = old_clause.send(:predicates)
+        return if old_predicates.empty?
+
+        new_predicates = old_predicates.map(&)
+        return if new_predicates == old_predicates
+
+        new_clause = old_clause.dup
+        new_clause.instance_variable_set(:@predicates, new_predicates)
+
+        send(clause_setter, new_clause)
+      end
+
+      def each_transposable_predicate(predicates, &block)
+        each_predicate(predicates) do |predicate|
+          case predicate
+          when ::Arel::Nodes::Grouping
+            next predicate unless predicate.expr.is_a?(::Arel::Nodes::Or)
+
+            or_expr = predicate.expr
+            old_left = or_expr.left
+            old_right = or_expr.right
+            new_left, new_right = each_transposable_predicate([old_left, old_right], &block)
+
+            next predicate if new_left == old_left && new_right == old_right
+
+            next predicate.class.new predicate.expr.class.new(new_left, new_right) if ::Rails.version < "7.2"
+
+            next predicate.class.new predicate.expr.class.new([new_left, new_right])
+
+          when ::Arel::Nodes::SelectStatement
+            new_cores = predicate.cores.map do |core|
+              next core unless core.is_a?(::Arel::Nodes::SelectCore) # just in case something weird is going on
+
+              new_wheres = each_transposable_predicate(core.wheres, &block)
+              new_havings = each_transposable_predicate(core.havings, &block)
+
+              next core if core.wheres == new_wheres && core.havings == new_havings
+
+              new_core = core.clone
+              new_core.wheres = new_wheres
+              new_core.havings = new_havings
+              new_core
+            end
+
+            next predicate if predicate.cores == new_cores
+
+            new_node = predicate.clone
+            new_node.instance_variable_set(:@cores, new_cores)
+            next new_node
+          when ::Arel::Nodes::Not
+            old_value = predicate.expr
+            new_value = each_transposable_predicate([old_value], &block).first
+
+            next predicate if old_value == new_value
+
+            next predicate.class.new(new_value)
+          when ::Arel::Nodes::Exists
+            old_value = predicate.expressions
+            new_value = each_transposable_predicate([old_value], &block).first
+
+            next predicate if old_value == new_value
+
+            next predicate.class.new(new_value)
           end
 
-        if new_right_value == right
-          predicate
-        elsif predicate.right.is_a?(::Arel::Nodes::Casted)
-          if new_right_value == right.value
-            predicate
-          else
-            predicate.class.new(predicate.left, right.class.new(new_right_value, right.attribute))
+          next predicate unless predicate.is_a?(::Arel::Nodes::Binary) || predicate.is_a?(::Arel::Nodes::HomogeneousIn)
+          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))
+
+          yield(predicate, relation, column, type)
+        end
+      end
+
+      def each_transposable_predicate_value(predicates = nil, &block)
+        each_transposable_predicate(predicates) do |predicate, relation, column, type|
+          each_transposable_predicate_value_cb(predicate, block) do |value|
+            yield(value, predicate, relation, column, type)
           end
-        elsif predicate.is_a?(::Arel::Nodes::HomogeneousIn)
+        end
+      end
+
+      def each_transposable_predicate_value_cb(node, original_block, &)
+        case node
+        when Array
+          node.filter_map { |val| each_transposable_predicate_value_cb(val, original_block, &).presence }
+        when ::ActiveModel::Attribute
+          old_value = node.value_before_type_cast
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
+
+          (old_value == new_value) ? node : node.class.new(node.name, new_value, node.type)
+        when ::Arel::Nodes::And
+          old_value = node.children
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
+
+          (old_value == new_value) ? node : node.class.new(new_value)
+        when ::Arel::Nodes::BindParam
+          old_value = node.value
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
+
+          (old_value == new_value) ? node : node.class.new(new_value)
+        when ::Arel::Nodes::Casted
+          old_value = node.value
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
+
+          (old_value == new_value) ? node : node.class.new(new_value, node.attribute)
+        when ::Arel::Nodes::HomogeneousIn
+          old_value = node.values
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
+
           # switch to a regular In, so that Relation::WhereClause#contradiction? knows about it
-          if new_right_value.empty?
-            klass = predicate.type == :in ? ::Arel::Nodes::In : ::Arel::Nodes::NotIn
-            klass.new(predicate.attribute, new_right_value)
+          if new_value.empty?
+            klass = (node.type == :in) ? ::Arel::Nodes::In : ::Arel::Nodes::NotIn
+            klass.new(node.attribute, new_value)
           else
-            predicate.class.new(new_right_value, predicate.attribute, predicate.type)
+            (old_value == new_value) ? node : node.class.new(new_value, node.attribute, node.type)
           end
+        when ::Arel::Nodes::Binary
+          old_value = node.right
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
+
+          (old_value == new_value) ? node : node.class.new(node.left, new_value)
+        when ::Arel::Nodes::SelectStatement
+          each_transposable_predicate_value([node], &original_block).first
         else
-          predicate.class.new(predicate.left, new_right_value)
+          yield(node)
         end
       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
-          current_id = query_att.value_before_type_cast
-          if current_id.is_a?(::ActiveRecord::StatementCache::Substitute)
-            current_id.sharded = true # mark for transposition later
-            current_id.primary = true if attribute_type == :primary
-            value
-          else
-            local_id = Shard.relative_id_for(current_id, current_shard, target_shard) || current_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
-              value
+      def transpose_predicates(predicates,
+                               source_shard,
+                               target_shard)
+        each_transposable_predicate_value(predicates) do |value, _predicate, relation, column, type|
+          current_source_shard =
+            if source_shard
+              source_shard
+            elsif type == :primary
+              Shard.current(klass.connection_class_for_self)
+            elsif type == :foreign
+              source_shard_for_foreign_key(relation, column)
             else
-              ::Arel::Nodes::BindParam.new(query_att.class.new(query_att.name, local_id, query_att.type))
+              primary_shard
             end
-          end
+
+          transpose_predicate_value(value, current_source_shard, target_shard, type)
+        end
+      end
+
+      def transpose_predicate_value(value, current_shard, target_shard, attribute_type)
+        if value.is_a?(NonTransposingValue)
+          value
+        elsif value.is_a?(::ActiveRecord::StatementCache::Substitute)
+          value.sharded = true # mark for transposition later
+          value.primary = true if attribute_type == :primary
+          value
         else
-          local_id = Shard.relative_id_for(value, current_shard, target_shard) || value
-          local_id = [] if remove_non_local_ids && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
-          local_id
+          Shard.relative_id_for(value, current_shard, target_shard) || value
         end
       end
     end

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
@@ -28,11 +28,18 @@ 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`
-        def association_scope_cache(klass, owner, &block)
+        def association_scope_cache(klass, owner, &)
           key = self
           key = [key, owner._read_attribute(@foreign_type)] if polymorphic?
           key = [key, shard(owner).id].flatten
-          klass.cached_find_by_statement(key, &block)
+
+          if ::Rails.version < "7.2"
+            klass.cached_find_by_statement(key, &)
+          else
+            klass.with_connection do |connection|
+              klass.cached_find_by_statement(connection, key, &)
+            end
+          end
         end
       end