module Switchman::ActiveRecord::QueryMethods

Public Instance Methods

all_shards() click to toggle source

the shard value as an array or a relation

# File lib/switchman/active_record/query_methods.rb, line 75
def all_shards
  case shard_value
  when Shard, DefaultShard
    [shard_value]
  when ::ActiveRecord::Base
    shard_value.respond_to?(:associated_shards) ? shard_value.associated_shards : [shard_value.shard]
  when nil
    [Shard.current(klass.connection_classes)]
  else
    shard_value
  end
end
or(other) click to toggle source
Calls superclass method
# File lib/switchman/active_record/query_methods.rb, line 88
def or(other)
  super(other.shard(primary_shard))
end
primary_shard() click to toggle source

the shard that where_values are relative to. if it's multiple shards, they're stored relative to the first shard

# File lib/switchman/active_record/query_methods.rb, line 56
def primary_shard
  case shard_value
  when Shard, DefaultShard
    shard_value
  # associated_shards
  when ::ActiveRecord::Base
    shard_value.shard
  when Array
    shard_value.first
  when ::ActiveRecord::Relation
    Shard.default
  when nil
    Shard.current(klass.connection_classes)
  else
    raise ArgumentError, "invalid shard value #{shard_value}"
  end
end
shard(value, source = :explicit) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 37
def shard(value, source = :explicit)
  spawn.shard!(value, source)
end
shard!(value, source = :explicit) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 41
def shard!(value, source = :explicit)
  raise ArgumentError, "shard can't be nil" unless value

  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
  self
end
shard_source_value() click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 21
def shard_source_value
  @values[:shard_source]
end
shard_source_value=(value) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 31
def shard_source_value=(value)
  raise ::ActiveRecord::ImmutableRelation if @loaded

  @values[:shard_source] = value
end
shard_value() click to toggle source

shard_value is one of:

A shard
An array or relation of shards
An AR object (query runs against that object's associated_shards)

shard_source_value is one of:

:implicit    - inferred from current shard when relation was created, or primary key where clause
: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
# File lib/switchman/active_record/query_methods.rb, line 17
def shard_value
  @values[:shard]
end
shard_value=(value) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 25
def shard_value=(value)
  raise ::ActiveRecord::ImmutableRelation if @loaded

  @values[:shard] = value
end

Private Instance Methods

arel_column(columns) click to toggle source
Calls superclass method
# File lib/switchman/active_record/query_methods.rb, line 238
def arel_column(columns)
  connection.with_local_table_name { super }
end
arel_columns(columns) click to toggle source
Calls superclass method
# File lib/switchman/active_record/query_methods.rb, line 234
def arel_columns(columns)
  connection.with_local_table_name { super }
end
build_where_clause(opts, rest = []) click to toggle source
Calls superclass method
# File lib/switchman/active_record/query_methods.rb, line 207
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
infer_shards_from_primary_key(predicates) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 116
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.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
  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?

    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
models_for_table(table_name) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 174
def models_for_table(table_name)
  @@models_for_table ||= {}
  @@models_for_table[table_name] ||= ::ActiveRecord::Base.descendants.select { |d| d.table_name == table_name }
end
relation_and_column(attribute) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 201
def relation_and_column(attribute)
  column = attribute.name
  attribute = attribute.relation if attribute.relation.is_a?(::Arel::Nodes::TableAlias)
  [attribute.relation, column]
end
sharded_foreign_key?(relation, column) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 179
def sharded_foreign_key?(relation, column)
  models_for_table(relation.table_name).any? { |m| m.sharded_column?(column) }
end
sharded_primary_key?(relation, column) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 183
def sharded_primary_key?(relation, column)
  column = column.to_s
  return column == 'id' if relation.klass == ::ActiveRecord::Base

  relation.klass.primary_key == column && relation.klass.integral_id?
end
source_shard_for_foreign_key(relation, column) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 190
def source_shard_for_foreign_key(relation, column)
  reflection = nil
  models_for_table(relation.table_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]

  Shard.current(reflection.klass.connection_classes)
end
table_name_matches?(from) click to toggle source
Calls superclass method
# File lib/switchman/active_record/query_methods.rb, line 242
def table_name_matches?(from)
  connection.with_global_table_name { super }
end
transposable_attribute_type(relation, column) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 166
def transposable_attribute_type(relation, column)
  if sharded_primary_key?(relation, column)
    :primary
  elsif sharded_foreign_key?(relation, column)
    :foreign
  end
end
transpose_clauses(source_shard, target_shard, remove_nonlocal_primary_keys: false) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 111
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
transpose_predicate_value(value, current_shard, target_shard, attribute_type, remove_non_local_ids) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 334
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
      else
        ::Arel::Nodes::BindParam.new(query_att.class.new(query_att.name, local_id, query_att.type))
      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
transpose_predicates(predicates, source_shard, target_shard, remove_nonlocal_primary_keys: false) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 246
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)
  end
end
transpose_single_predicate(predicate, source_shard, target_shard, remove_nonlocal_primary_keys: false) click to toggle source
# File lib/switchman/active_record/query_methods.rb, line 256
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)
    end

  right = if predicate.is_a?(::Arel::Nodes::HomogeneousIn)
            predicate.values
          else
            predicate.right
          end

  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)
    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))
    end
  elsif predicate.is_a?(::Arel::Nodes::HomogeneousIn)
    # 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)
    else
      predicate.class.new(new_right_value, predicate.attribute, predicate.type)
    end
  else
    predicate.class.new(predicate.left, new_right_value)
  end
end