switchman 3.4.1 → 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
@@ -21,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
@@ -84,8 +113,20 @@ module Switchman
         super(other.shard(primary_shard))
       end
 
+      # 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
+
       protected
 
+      def arel_columns(columns)
+        connection.with_local_table_name { super }
+      end
+
       def remove_nonlocal_primary_keys!
         each_transposable_predicate_value do |value, predicate, _relation, _column, type|
           next value unless
@@ -93,7 +134,7 @@ module Switchman
             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
+          (value.is_a?(Integer) && value > Shard::IDS_PER_SHARD) ? [] : value
         end
         self
       end
@@ -104,7 +145,9 @@ module Switchman
         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
@@ -167,19 +210,19 @@ 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
@@ -199,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
@@ -221,31 +263,49 @@ 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 table_name_matches?(from)
-        connection.with_global_table_name { super }
+        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
+
+      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
 
-      def each_predicate(predicates = nil, &block)
-        return predicates.map(&block) if predicates
+      def each_predicate(predicates = nil, &)
+        return predicates.map(&) if predicates
 
-        each_predicate_cb(:having_clause, :having_clause=, &block)
-        each_predicate_cb(:where_clause, :where_clause=, &block)
+        each_predicate_cb(:having_clause, :having_clause=, &)
+        each_predicate_cb(:where_clause, :where_clause=, &)
       end
 
-      def each_predicate_cb(clause_getter, clause_setter, &block)
+      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(&block)
+        new_predicates = old_predicates.map(&)
         return if new_predicates == old_predicates
 
         new_clause = old_clause.dup
@@ -254,9 +314,10 @@ module Switchman
         send(clause_setter, new_clause)
       end
 
-      def each_transposable_predicate(predicates = nil, &block)
+      def each_transposable_predicate(predicates, &block)
         each_predicate(predicates) do |predicate|
-          if predicate.is_a?(::Arel::Nodes::Grouping)
+          case predicate
+          when ::Arel::Nodes::Grouping
             next predicate unless predicate.expr.is_a?(::Arel::Nodes::Or)
 
             or_expr = predicate.expr
@@ -266,7 +327,44 @@ module Switchman
 
             next predicate if new_left == old_left && new_right == old_right
 
-            next predicate.class.new predicate.expr.class.new(new_left, new_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
 
           next predicate unless predicate.is_a?(::Arel::Nodes::Binary) || predicate.is_a?(::Arel::Nodes::HomogeneousIn)
@@ -279,54 +377,56 @@ module Switchman
         end
       end
 
-      def each_transposable_predicate_value(predicates = nil)
+      def each_transposable_predicate_value(predicates = nil, &block)
         each_transposable_predicate(predicates) do |predicate, relation, column, type|
-          each_transposable_predicate_value_cb(predicate) do |value|
+          each_transposable_predicate_value_cb(predicate, block) do |value|
             yield(value, predicate, relation, column, type)
           end
         end
       end
 
-      def each_transposable_predicate_value_cb(node, &block)
+      def each_transposable_predicate_value_cb(node, original_block, &)
         case node
         when Array
-          node.map { |val| each_transposable_predicate_value_cb(val, &block).presence }.compact
+          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, &block)
+          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)
+          (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, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
-          old_value == new_value ? node : node.class.new(new_value)
+          (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, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
-          old_value == new_value ? node : node.class.new(new_value)
+          (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, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
-          old_value == new_value ? node : node.class.new(new_value, node.attribute)
+          (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, &block)
+          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_value.empty?
-            klass = node.type == :in ? ::Arel::Nodes::In : ::Arel::Nodes::NotIn
+            klass = (node.type == :in) ? ::Arel::Nodes::In : ::Arel::Nodes::NotIn
             klass.new(node.attribute, new_value)
           else
-            old_value == new_value ? node : node.class.new(new_value, node.attribute, node.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, &block)
+          new_value = each_transposable_predicate_value_cb(old_value, original_block, &)
 
-          old_value == new_value ? node : node.class.new(node.left, new_value)
+          (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
           yield(node)
         end
@@ -343,6 +443,8 @@ module Switchman
               Shard.current(klass.connection_class_for_self)
             elsif type == :foreign
               source_shard_for_foreign_key(relation, column)
+            else
+              primary_shard
             end
 
           transpose_predicate_value(value, current_source_shard, target_shard, type)
@@ -350,7 +452,9 @@ module Switchman
       end
 
       def transpose_predicate_value(value, current_shard, target_shard, attribute_type)
-        if value.is_a?(::ActiveRecord::StatementCache::Substitute)
+        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

data/lib/switchman/active_record/reflection.rb

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

data/lib/switchman/active_record/relation.rb

@@ -28,15 +28,15 @@ module Switchman
         relation
       end
 
-      def new(*, &block)
+      def new(*, &)
         primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
-      def create(*, &block)
+      def create(*, &)
         primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
-      def create!(*, &block)
+      def create!(*, &)
         primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
@@ -44,12 +44,18 @@ module Switchman
         primary_shard.activate(klass.connection_class_for_self) { super }
       end
 
-      def explain
-        activate { |relation| relation.call_super(:explain, Relation) }
+      if ::Rails.version > "7.1.2"
+        def transaction(...)
+          primary_shard.activate(klass.connection_class_for_self) { super }
+        end
+      end
+
+      def explain(*args)
+        activate { |relation| relation.call_super(:explain, Relation, *args) }
       end
 
       def load(&block)
-        if !loaded? || (::Rails.version >= '7.0' && scheduled?)
+        if !loaded? || scheduled?
           @records = activate { |relation| relation.send(:exec_queries, &block) }
           @loaded = true
         end
@@ -58,27 +64,83 @@ module Switchman
       end
 
       %I[update_all delete_all].each do |method|
-        arg_params = RUBY_VERSION <= '2.8' ? '*args' : '*args, **kwargs'
         class_eval <<-RUBY, __FILE__, __LINE__ + 1
-          def #{method}(#{arg_params})
-            result = self.activate(unordered: true) { |relation| relation.call_super(#{method.inspect}, Relation, #{arg_params}) }
+          def #{method}(*args, **kwargs)
+            result = self.activate(unordered: true) { |relation| relation.call_super(#{method.inspect}, Relation, *args, **kwargs) }
             result = result.sum if result.is_a?(Array)
             result
           end
         RUBY
       end
 
+      # https://github.com/rails/rails/commit/ed2c15b52450ff927a05629f031376f25b670335
+      # once the minimum version is Rails 7.2, we can drop this separate module
+      module InsertUpsertAll
+        %w[insert_all upsert_all].each do |method|
+          class_eval <<-RUBY, __FILE__, __LINE__ + 1
+            def #{method}(attributes, returning: nil, **)
+              scope = self != ::ActiveRecord::Base && current_scope
+              if (target_shard = scope&.primary_shard) == (current_shard = Shard.current(connection_class_for_self))
+                scope = nil
+              end
+              if scope
+                dupped = false
+                attributes.each_with_index do |hash, i|
+                  if dupped || hash.any? { |k, v| sharded_column?(k) }
+                    unless dupped
+                      attributes = attributes.dup
+                      dupped = true
+                    end
+                    attributes[i] = hash.to_h do |k, v|
+                      if sharded_column?(k)
+                        [k, Shard.relative_id_for(v, current_shard, target_shard)]
+                      else
+                        [k, v]
+                      end
+                    end
+                  end
+                end
+              end
+
+              if scope
+                scope.activate do
+                  db = Shard.current(connection_class_for_self).database_server
+                  result = db.unguard { super }
+                  if result&.columns&.any? { |c| sharded_column?(c) }
+                    transposed_rows = result.rows.map do |row|
+                      row.map.with_index do |value, i|
+                        sharded_column?(result.columns[i]) ? Shard.relative_id_for(value, target_shard, current_shard) : value
+                      end
+                    end
+                    result = ::ActiveRecord::Result.new(result.columns, transposed_rows, result.column_types)
+                  end
+
+                  result
+                end
+              else
+                db = Shard.current(connection_class_for_self).database_server
+                db.unguard { super }
+              end
+            end
+          RUBY
+        end
+      end
+
       def find_ids_in_ranges(options = {})
         is_integer = columns_hash[primary_key.to_s].type == :integer
         loose_mode = options[:loose] && is_integer
         # loose_mode: if we don't care about getting exactly batch_size ids in between
         # don't get the max - just get the min and add batch_size so we get that many _at most_
-        values = loose_mode ? 'MIN(id)' : 'MIN(id), MAX(id)'
+        values = loose_mode ? "MIN(id)" : "MIN(id), MAX(id)"
 
         batch_size = options[:batch_size].try(:to_i) || 1000
-        quoted_primary_key = "#{klass.connection.quote_local_table_name(table_name)}.#{klass.connection.quote_column_name(primary_key)}"
-        as_id = ' AS id' unless primary_key == 'id'
-        subquery_scope = except(:select).select("#{quoted_primary_key}#{as_id}").reorder(primary_key.to_sym).limit(loose_mode ? 1 : batch_size)
+        quoted_primary_key =
+          "#{klass.connection.quote_local_table_name(table_name)}.#{klass.connection.quote_column_name(primary_key)}"
+        as_id = " AS id" unless primary_key == "id"
+        subquery_scope = except(:select)
+                         .select("#{quoted_primary_key}#{as_id}")
+                         .reorder(primary_key.to_sym)
+                         .limit(loose_mode ? 1 : batch_size)
         subquery_scope = subquery_scope.where("#{quoted_primary_key} <= ?", options[:end_at]) if options[:end_at]
 
         first_subquery_scope = if options[:start_at]
@@ -101,7 +163,7 @@ module Switchman
         end
       end
 
-      def activate(unordered: false, &block)
+      def activate(count: false, unordered: false, &block)
         shards = all_shards
         if Array === shards && shards.length == 1
           if !loaded? && shard_value != shards.first
@@ -121,11 +183,19 @@ module Switchman
             relation = shard(Shard.current(klass.connection_class_for_self))
             relation.remove_nonlocal_primary_keys!
             # do a minimal query if possible
-            relation = relation.limit(limit_value - result_count) if limit_value && !result_count.zero? && order_values.empty?
+            if limit_value && !result_count.zero? && order_values.empty?
+              relation = relation.limit(limit_value - result_count)
+            end
 
             shard_results = relation.activate(&block)
 
-            if shard_results.present? && !unordered
+            if shard_results.present? && count
+              unless shard_results.is_a?(Integer)
+                raise "expected integer result for count, got #{shard_results.class.name}"
+              end
+
+              result_count += shard_results
+            elsif shard_results.present? && !unordered
               can_order ||= can_order_cross_shard_results? unless order_values.empty?
               raise OrderOnMultiShardQuery if !can_order && !order_values.empty? && result_count.positive?
 

data/lib/switchman/active_record/spawn_methods.rb

@@ -17,7 +17,7 @@ module Switchman
             final_shard_source_value = %i[explicit association].detect do |source_value|
               shard_source_value == source_value || rhs.shard_source_value == source_value
             end
-            raise 'unknown shard_source_value' unless final_shard_source_value
+            raise "unknown shard_source_value" unless final_shard_source_value
 
             # have to merge shard_value
             lhs_shard_value = all_shards
@@ -36,7 +36,7 @@ module Switchman
           final_shard_source_value = %i[explicit association implicit].detect do |source_value|
             shard_source_value == source_value || rhs.shard_source_value == source_value
           end
-          raise 'unknown shard_source_value' unless final_shard_source_value
+          raise "unknown shard_source_value" unless final_shard_source_value
         end
 
         [final_shard_value, final_primary_shard, final_shard_source_value]
@@ -62,16 +62,12 @@ module Switchman
         if primary_shard != final_primary_shard && rhs.primary_shard != final_primary_shard
           shard!(final_primary_shard)
           rhs = rhs.shard(final_primary_shard)
-          super(rhs)
         elsif primary_shard != final_primary_shard
           shard!(final_primary_shard)
-          super(rhs)
         elsif rhs.primary_shard != final_primary_shard
           rhs = rhs.shard(final_primary_shard)
-          super(rhs)
-        else
-          super
         end
+        super
 
         self.shard_value = final_shard_value
         self.shard_source_value = final_shard_source_value

data/lib/switchman/active_record/statement_cache.rb

@@ -4,8 +4,8 @@ module Switchman
   module ActiveRecord
     module StatementCache
       module ClassMethods
-        def create(connection, &block)
-          relation = block.call ::ActiveRecord::StatementCache::Params.new
+        def create(connection)
+          relation = yield ::ActiveRecord::StatementCache::Params.new
 
           _query_builder, binds = connection.cacheable_query(self, relation.arel)
           bind_map = ::ActiveRecord::StatementCache::BindMap.new(binds)
@@ -29,14 +29,14 @@ module Switchman
         params, connection = args
         klass = @klass
         target_shard = nil
-        if (primary_index = bind_map.primary_value_index)
+        if (primary_index = @bind_map.primary_value_index)
           primary_value = params[primary_index]
           target_shard = Shard.local_id_for(primary_value)[1]
         end
         current_shard = Shard.current(klass.connection_class_for_self)
         target_shard ||= current_shard
 
-        bind_values = bind_map.bind(params, current_shard, target_shard)
+        bind_values = @bind_map.bind(params, current_shard, target_shard)
 
         target_shard.activate(klass.connection_class_for_self) do
           sql = qualified_query_builder(target_shard, klass).sql_for(bind_values, connection)

data/lib/switchman/active_record/table_definition.rb

@@ -4,7 +4,7 @@ module Switchman
   module ActiveRecord
     module TableDefinition
       def column(name, type, limit: nil, **)
-        Switchman.foreign_key_check(name, type, limit: limit)
+        Switchman.foreign_key_check(name, type, limit:)
         super
       end
     end

data/lib/switchman/active_record/tasks/database_tasks.rb

@@ -7,9 +7,14 @@ module Switchman
         def drop(*)
           super
           # no really, it's gone
-          Switchman.cache.delete('default_shard')
+          Switchman.cache.delete("default_shard")
           Shard.default(reload: true)
         end
+
+        def raise_for_multi_db(*)
+          # ignore; Switchman doesn't use namespaced tasks for multiple shards; it uses
+          # environment variables to filter which shards you want to target
+        end
       end
     end
   end