switchman 2.0.5 → 3.0.1

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,118 +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(self.primary_shard))
+        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)
@@ -205,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)
@@ -215,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)
@@ -225,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)
@@ -237,144 +239,86 @@ module Switchman
         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,
-                               dup_binds_on_mutation: false)
-        result = predicates.map do |predicate|
-          transposed, binds = transpose_single_predicate(predicate, source_shard, target_shard, remove_nonlocal_primary_keys,
-                                     binds: binds, dup_binds_on_mutation: dup_binds_on_mutation)
-          transposed
+                               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
-        result = [result, binds]
-        result
       end
 
       def transpose_single_predicate(predicate,
                                      source_shard,
                                      target_shard,
-                                     remove_nonlocal_primary_keys = false,
-                                     binds: nil,
-                                     dup_binds_on_mutation: false)
+                                     remove_nonlocal_primary_keys: false)
         if predicate.is_a?(::Arel::Nodes::Grouping)
-          return predicate, binds unless predicate.expr.is_a?(::Arel::Nodes::Or)
+          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...
-          #
-          # TODO: "binds" is getting passed up and down
-          # this stack purely because of the necessary handling for rails <5.2
-          #  Dropping support for 5.2 means we can remove the "binds" argument from
-          # all of this and yank the conditional below where we monkey with their instance state.
           or_expr = predicate.expr
           left_node = or_expr.left
           right_node = or_expr.right
-          left_predicates = left_node.children
-          right_predicates = right_node.children
-          new_left_predicates, binds = transpose_predicates(left_predicates, source_shard,
-                                                               target_shard, remove_nonlocal_primary_keys,
-                                                               binds: binds, dup_binds_on_mutation: dup_binds_on_mutation)
-          new_right_predicates, binds = transpose_predicates(right_predicates, source_shard,
-                                                               target_shard, remove_nonlocal_primary_keys,
-                                                               binds: binds, dup_binds_on_mutation: dup_binds_on_mutation)
-          if new_left_predicates != left_predicates
-            left_node.instance_variable_set(:@children, new_left_predicates)
-          end
-          if new_right_predicates != right_predicates
-            right_node.instance_variable_set(:@children, new_right_predicates)
-          end
-          return predicate, binds
+          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, binds unless predicate.is_a?(::Arel::Nodes::Binary)
-        return predicate, binds unless predicate.left.is_a?(::Arel::Attributes::Attribute)
+        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, binds unless (type = transposable_attribute_type(relation, column))
+        return 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)
+                         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.shard_category)
-            elsif type == :foreign
-              source_shard_for_foreign_key(relation, column)
-            end
+          if source_shard
+            source_shard
+          elsif type == :primary
+            Shard.current(klass.connection_classes)
+          elsif type == :foreign
+            source_shard_for_foreign_key(relation, column)
+          end
 
-        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
+        right = if predicate.is_a?(::Arel::Nodes::HomogeneousIn)
+                  predicate.values
+                else
+                  predicate.right
                 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
-              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)
-              end
-            end
-            predicate.right
+
+        new_right_value =
+          case right
+          when Array
+            right.map { |val| transpose_predicate_value(val, current_source_shard, target_shard, type, remove) }
           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
+            transpose_predicate_value(right, current_source_shard, target_shard, type, remove)
           end
-        end
-        out_predicate = if new_right_value == predicate.right
+
+        if new_right_value == right
           predicate
         elsif predicate.right.is_a?(::Arel::Nodes::Casted)
-          if new_right_value == predicate.right.val
+          if new_right_value == right.value
             predicate
           else
-            predicate.class.new(predicate.left, predicate.right.class.new(new_right_value, predicate.right.attribute))
+            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
-        return out_predicate, binds
       end
 
       def transpose_predicate_value(value, current_shard, target_shard, attribute_type, remove_non_local_ids)
@@ -388,11 +332,11 @@ module Switchman
           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
+            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,21 +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`
-        def association_scope_cache(conn, owner, &block)
-          key = conn.prepared_statements
-          if polymorphic?
-            key = [key, owner._read_attribute(@foreign_type)]
-          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
-          @association_scope_cache[key] ||= @scope_lock.synchronize {
-            @association_scope_cache[key] ||= (::Rails.version >= "5.2" ? ::ActiveRecord::StatementCache.create(conn, &block) : block.call)
-          }
+          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

data/lib/switchman/active_record/relation.rb

@@ -4,53 +4,54 @@ module Switchman
   module ActiveRecord
     module Relation
       def self.prepended(klass)
-        klass::SINGLE_VALUE_METHODS.concat [ :shard, :shard_source ]
+        klass::SINGLE_VALUE_METHODS.concat %i[shard shard_source]
       end
 
       def initialize(*, **)
         super
-        self.shard_value = Shard.current(klass ? klass.shard_category : :primary) unless shard_value
+        self.shard_value = Shard.current(klass ? klass.connection_classes : :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 : :primary) unless shard_value
+        result.shard_value = Shard.current(klass ? klass.connection_classes : :primary) unless shard_value
         result
       end
 
       def merge(*)
         relation = super
-        if relation.shard_value != self.shard_value && relation.shard_source_value == :implicit
-          relation.shard_value = self.shard_value
-          relation.shard_source_value = self.shard_source_value
+        if relation.shard_value != shard_value && relation.shard_source_value == :implicit
+          relation.shard_value = shard_value
+          relation.shard_source_value = shard_source_value
         end
         relation
       end
 
       def new(*, &block)
-        primary_shard.activate(klass.shard_category) { super }
+        primary_shard.activate(klass.connection_classes) { super }
       end
 
       def create(*, &block)
-        primary_shard.activate(klass.shard_category) { super }
+        primary_shard.activate(klass.connection_classes) { super }
       end
 
       def create!(*, &block)
-        primary_shard.activate(klass.shard_category) { super }
+        primary_shard.activate(klass.connection_classes) { super }
       end
 
       def to_sql
-        primary_shard.activate(klass.shard_category) { super }
+        primary_shard.activate(klass.connection_classes) { super }
       end
 
       def explain
-        self.activate { |relation| relation.call_super(:explain, Relation) }
+        activate { |relation| relation.call_super(:explain, Relation) }
       end
 
       def records
         return @records if loaded?
-        results = self.activate { |relation| relation.call_super(:records, Relation) }
+
+        results = activate { |relation| relation.call_super(:records, Relation) }
         case shard_value
         when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
           @records = results
@@ -59,7 +60,7 @@ module Switchman
         results
       end
 
-      %I{update_all delete_all}.each do |method|
+      %I[update_all delete_all].each do |method|
         class_eval <<-RUBY, __FILE__, __LINE__ + 1
           def #{method}(*args)
             result = self.activate { |relation| relation.call_super(#{method.inspect}, Relation, *args) }
@@ -74,15 +75,20 @@ module Switchman
         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'
+        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
+        first_subquery_scope = if options[:start_at]
+                                 subquery_scope.where("#{quoted_primary_key} >= ?",
+                                                      options[:start_at])
+                               else
+                                 subquery_scope
+                               end
 
         ids = connection.select_rows("SELECT #{values} FROM (#{first_subquery_scope.to_sql}) AS subquery").first
 
@@ -99,16 +105,16 @@ module Switchman
 
       def activate(&block)
         shards = all_shards
-        if (Array === shards && shards.length == 1)
-          if shards.first == DefaultShard || shards.first == Shard.current(klass.shard_category)
+        if Array === shards && shards.length == 1
+          if shards.first == DefaultShard || shards.first == Shard.current(klass.connection_classes)
             yield(self, shards.first)
           else
-            shards.first.activate(klass.shard_category) { yield(self, shards.first) }
+            shards.first.activate(klass.connection_classes) { yield(self, shards.first) }
           end
         else
           # TODO: implement local limit to avoid querying extra shards
-          Shard.with_each_shard(shards, [klass.shard_category]) do
-            shard(Shard.current(klass.shard_category), :to_a).activate(&block)
+          Shard.with_each_shard(shards, [klass.connection_classes]) do
+            shard(Shard.current(klass.connection_classes), :to_a).activate(&block)
           end
         end
       end