switchman 1.5.13 → 2.0.9

data/lib/switchman/active_record/query_methods.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module QueryMethods
@@ -19,11 +21,11 @@ module Switchman
         @values[:shard_source]
       end
       def shard_value=(value)
-        raise ImmutableRelation if @loaded
+        raise ::ActiveRecord::ImmutableRelation if @loaded
         @values[:shard] = value
       end
       def shard_source_value=(value)
-        raise ImmutableRelation if @loaded
+        raise ::ActiveRecord::ImmutableRelation if @loaded
         @values[:shard_source] = value
       end
 
@@ -42,34 +44,6 @@ module Switchman
         self
       end
 
-      if ::Rails.version < '5'
-        # moved to WhereClauseFactory#build in Rails 5
-        def build_where(opts, other = [])
-          case opts
-          when String, Array
-            values = Hash === other.first ? other.first.values : other
-
-            values.grep(ActiveRecord::Relation) do |rel|
-              # serialize subqueries against the same shard as the outer query is currently
-              # targeted to run against
-              if rel.shard_source_value == :implicit && rel.primary_shard != primary_shard
-                rel.shard!(primary_shard)
-              end
-              self.bind_values += rel.bind_values if ::Rails.version < '4.2'
-            end
-
-            [@klass.send(:sanitize_sql, other.empty? ? opts : ([opts] + other))]
-          when Hash, ::Arel::Nodes::Node
-            predicates = super
-            infer_shards_from_primary_key(predicates) if shard_source_value == :implicit && shard_value.is_a?(Shard)
-            predicates = transpose_predicates(predicates, nil, primary_shard) if shard_source_value != :explicit
-            predicates
-          else
-            super
-          end
-        end
-      end
-
       # the shard that where_values are relative to. if it's multiple shards, they're stored
       # relative to the first shard
       def primary_shard
@@ -104,44 +78,47 @@ module Switchman
         end
       end
 
-      if ::Rails.version >= '4.2' && ::Rails.version < '5'
-        # fixes an issue in Rails 4.2 with `reverse_sql_order` and qualified names
-        # where quoted_table_name is called before shard(s) have been activated
-        # if there's no ordering
-        def reverse_order!
-          orders = order_values.uniq
-          orders.reject!(&:blank?)
-          if orders.empty?
-            self.order_values = [arel_table[primary_key].desc]
-          else
-            self.order_values = reverse_sql_order(orders)
-          end
-          self
-        end
-      end
-
       private
 
-      [:where, :having].each do |type|
-        class_eval <<-RUBY, __FILE__, __LINE__ + 1
+      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)
-            if ::Rails.version >= '5'
-              unless (predicates = #{type}_clause.send(:predicates)).empty?
-                new_predicates = transpose_predicates(predicates, 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)
+              if new_predicates != predicates
+                self.#{type}_clause = #{type}_clause.dup
                 if new_predicates != predicates
-                  self.#{type}_clause = #{type}_clause.dup
                   #{type}_clause.instance_variable_set(:@predicates, new_predicates)
                 end
               end
-            else
-              unless #{type}_values.empty?
-                self.#{type}_values = transpose_predicates(#{type}_values,
-                  source_shard, target_shard, remove_nonlocal_primary_keys)
-              end
             end
           end
-        RUBY
+          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
       end
 
       def transpose_clauses(source_shard, target_shard, remove_nonlocal_primary_keys = false)
@@ -150,6 +127,8 @@ module Switchman
       end
 
       def infer_shards_from_primary_key(predicates, binds = nil)
+        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 &&
@@ -180,28 +159,23 @@ module Switchman
               return
             end
           when ::Arel::Nodes::BindParam
-            # look for a bind param with a matching column name
-            if ::Rails.version >= "5"
-              binds ||= where_clause.binds + having_clause.binds
-              if binds && bind = binds.detect{|b| b.try(:name).to_s == klass.primary_key.to_s}
+            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
-            else
-              if bind_values && idx = bind_values.find_index{|b| b.is_a?(Array) && b.first.try(:name).to_s == klass.primary_key.to_s}
-                column, value = bind_values[idx]
-                unless ::Rails.version >= '4.2' && value.is_a?(::ActiveRecord::StatementCache::Substitute)
-                  local_id, id_shard = Shard.local_id_for(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
           end
+
           return if !id_shard || id_shard == primary_shard
           transpose_clauses(primary_shard, id_shard)
           self.shard_value = id_shard
@@ -228,7 +202,7 @@ 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.primary_key == column && relation.model.integral_id?
       end
 
       def source_shard_for_foreign_key(relation, column)
@@ -252,25 +226,22 @@ module Switchman
       end
 
       def arel_columns(columns)
-        columns.map do |field|
-          if (Symbol === field || String === field) && ::Rails.version >= '5' && (klass.has_attribute?(field) || klass.attribute_alias?(field)) && !from_clause.value
-            klass.arel_attribute(field, table)
-          elsif (Symbol === field || String === field) && ::Rails.version < '5' && columns_hash.key?(field.to_s) && !from_value
-            arel_table[field]
-          elsif Symbol === field
-            # the rest of this is pulled from AR - the only change is from quote_table_name to quote_column_name here
-            # otherwise qualified names will add the schema to a column
-            connection.quote_column_name(field.to_s)
-          else
-            field
-          end
-        end
+        connection.with_local_table_name { super }
+      end
+
+      def arel_column(columns)
+        connection.with_local_table_name { super }
       end
 
       # semi-private
       public
-      def transpose_predicates(predicates, source_shard, target_shard, remove_nonlocal_primary_keys = false, binds = nil)
-        predicates.map do |predicate|
+      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|
           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)
@@ -290,28 +261,42 @@ module Switchman
                 source_shard_for_foreign_key(relation, column)
               end
 
-          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)
-              unless remove && local_id > Shard::IDS_PER_SHARD
-                if ::Rails.version > "4.2" && 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)
+          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) }
+              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
-                local_ids << local_id
               end
-            end
-            local_ids
-          when ::Arel::Nodes::BindParam
-            # look for a bind param with a matching column name
-            if ::Rails.version >= "5"
-              binds ||= where_clause.binds + having_clause.binds
-              if binds && bind = binds.detect{|b| b.try(:name).to_s == predicate.left.name.to_s}
+              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
@@ -322,29 +307,16 @@ module Switchman
                   bind.instance_variable_set(:@value_for_database, nil)
                 end
               end
+              predicate.right
             else
-              if bind_values && idx = bind_values.find_index{|b| b.is_a?(Array) && b.first.try(:name).to_s == predicate.left.name.to_s}
-                column, value = bind_values[idx]
-                if ::Rails.version >= '4.2' && value.is_a?(::ActiveRecord::StatementCache::Substitute)
-                  value.sharded = true # mark for transposition later
-                  value.primary = true if type == :primary
-                else
-                  local_id = Shard.relative_id_for(value, current_source_shard, target_shard)
-                  local_id = [] if remove && local_id > Shard::IDS_PER_SHARD
-                  bind_values[idx] = [column, local_id]
-                end
-              end
+              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
-            predicate.right
-          else
-            local_id = Shard.relative_id_for(predicate.right, current_source_shard, target_shard)
-            local_id = [] if remove && local_id.is_a?(Fixnum) && local_id > Shard::IDS_PER_SHARD
-            local_id
           end
-
           if new_right_value == predicate.right
             predicate
-          elsif ::Rails.version >= "4.2" && predicate.right.is_a?(::Arel::Nodes::Casted)
+          elsif predicate.right.is_a?(::Arel::Nodes::Casted)
             if new_right_value == predicate.right.val
               predicate
             else
@@ -354,6 +326,33 @@ module Switchman
             predicate.class.new(predicate.left, new_right_value)
           end
         end
+        result = [result, binds]
+        result
+      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
+              ::Arel::Nodes::BindParam.new(query_att.class.new(query_att.name, local_id, query_att.type))
+            else
+              value
+            end
+          end
+        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
+        end
       end
     end
   end

data/lib/switchman/active_record/reflection.rb

@@ -1,31 +1,48 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Reflection
-      module AssociationReflection
-        # removes memoization - ActiveRecord::ModelSchema does that anyway;
-        # and in fact this is the exact change AR makes in 4.2+
-        if ::Rails.version < '4.2'
-          def quoted_table_name
-            klass.quoted_table_name
-          end
-        else
-          def join_id_for(owner)
-            owner.send(active_record_primary_key) # use sharded id values in association binds
+      module AbstractReflection
+        def shard(owner)
+          if polymorphic? || klass.shard_category == owner.class.shard_category
+            # polymorphic associations assume the same shard as the owning item
+            owner.shard
+          else
+            Shard.default
           end
         end
+      end
+
+      module AssociationScopeCache
+        def initialize(*args)
+          super
+          # on ThroughReflection, these won't be initialized (cause it doesn't
+          # inherit from AssociationReflection), so make sure they're
+          # initialized here
+          @association_scope_cache ||= {}
+          @scope_lock ||= Mutex.new
+        end
 
         # cache association scopes by shard.
-        if ::Rails.version >= '4.2'
-          def association_scope_cache(conn, owner)
-            key = conn.prepared_statements
-            if polymorphic?
-              key = [key, owner._read_attribute(@foreign_type)]
-            end
-            key = [key, owner.shard.id].flatten
-            @association_scope_cache[key] ||= @scope_lock.synchronize {
-              @association_scope_cache[key] ||= yield
-            }
+        # 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(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
+      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
         end
       end
     end

data/lib/switchman/active_record/relation.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module Relation
@@ -5,19 +7,19 @@ module Switchman
         klass::SINGLE_VALUE_METHODS.concat [ :shard, :shard_source ]
       end
 
-      def initialize(*args)
+      def initialize(*, **)
         super
-        self.shard_value = Shard.current(klass ? klass.shard_category : :default) unless shard_value
+        self.shard_value = Shard.current(klass ? klass.shard_category : :primary) unless shard_value
         self.shard_source_value = :implicit unless shard_source_value
       end
 
       def clone
         result = super
-        result.shard_value = Shard.current(klass ? klass.shard_category : :default) unless shard_value
+        result.shard_value = Shard.current(klass ? klass.shard_category : :primary) unless shard_value
         result
       end
 
-      def merge(*args)
+      def merge(*)
         relation = super
         if relation.shard_value != self.shard_value && relation.shard_source_value == :implicit
           relation.shard_value = self.shard_value
@@ -26,15 +28,15 @@ module Switchman
         relation
       end
 
-      def new(*args, &block)
+      def new(*, &block)
         primary_shard.activate(klass.shard_category) { super }
       end
 
-      def create(*args, &block)
+      def create(*, &block)
         primary_shard.activate(klass.shard_category) { super }
       end
 
-      def create!(*args, &block)
+      def create!(*, &block)
         primary_shard.activate(klass.shard_category) { super }
       end
 
@@ -42,40 +44,59 @@ module Switchman
         primary_shard.activate(klass.shard_category) { super }
       end
 
-      def explain(super_method: false)
-        return super() if super_method
-        self.activate { |relation| relation.explain(super_method: true) }
+      def explain
+        self.activate { |relation| relation.call_super(:explain, Relation) }
       end
 
-      to_a_method = ::Rails.version > '5.0.0.beta2' ? :records : :to_a
-      class_eval <<-RUBY, __FILE__, __LINE__ + 1
-        def #{to_a_method}(super_method: false)
-          return super() if super_method
-          return @records if loaded?
-          results = self.activate { |relation| relation.#{to_a_method}(super_method: true) }
-          case shard_value
-          when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
-            @records = results
-            @loaded = true
-          end
-          results
+      def records
+        return @records if loaded?
+        results = self.activate { |relation| relation.call_super(:records, Relation) }
+        case shard_value
+        when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
+          @records = results
+          @loaded = true
         end
-      RUBY
-
-      CALL_SUPER = Object.new.freeze
-      private_constant :CALL_SUPER
+        results
+      end
 
-      %w{update_all delete_all}.each do |method|
+      %I{update_all delete_all}.each do |method|
         class_eval <<-RUBY, __FILE__, __LINE__ + 1
           def #{method}(*args)
-            return super(*args[1..-1]) if args.first.equal?(CALL_SUPER)
-            result = self.activate { |relation| relation.#{method}(CALL_SUPER, *args) }
+            result = self.activate { |relation| relation.call_super(#{method.inspect}, Relation, *args) }
             result = result.sum if result.is_a?(Array)
             result
           end
         RUBY
       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)"
+
+        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)
+        subquery_scope = subquery_scope.where("#{quoted_primary_key} <= ?", options[:end_at]) if options[:end_at]
+
+        first_subquery_scope = options[:start_at] ? subquery_scope.where("#{quoted_primary_key} >= ?", options[:start_at]) : subquery_scope
+
+        ids = connection.select_rows("SELECT #{values} FROM (#{first_subquery_scope.to_sql}) AS subquery").first
+
+        while ids.first.present?
+          ids.map!(&:to_i) if is_integer
+          ids << ids.first + batch_size if loose_mode
+
+          yield(*ids)
+          last_value = ids.last
+          next_subquery_scope = subquery_scope.where(["#{quoted_primary_key}>?", last_value])
+          ids = connection.select_rows("SELECT #{values} FROM (#{next_subquery_scope.to_sql}) AS subquery").first
+        end
+      end
+
       def activate(&block)
         shards = all_shards
         if (Array === shards && shards.length == 1)

data/lib/switchman/active_record/spawn_methods.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Switchman
   module ActiveRecord
     module SpawnMethods