switchman 3.2.1 → 4.0.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (46) hide show
  1. checksums.yaml +4 -4
  2. data/Rakefile +15 -14
  3. data/db/migrate/20180828183945_add_default_shard_index.rb +1 -1
  4. data/db/migrate/20190114212900_add_unique_name_indexes.rb +10 -4
  5. data/lib/switchman/active_record/abstract_adapter.rb +4 -2
  6. data/lib/switchman/active_record/associations.rb +89 -49
  7. data/lib/switchman/active_record/attribute_methods.rb +72 -34
  8. data/lib/switchman/active_record/base.rb +145 -27
  9. data/lib/switchman/active_record/calculations.rb +96 -49
  10. data/lib/switchman/active_record/connection_handler.rb +18 -0
  11. data/lib/switchman/active_record/connection_pool.rb +24 -3
  12. data/lib/switchman/active_record/database_configurations.rb +37 -15
  13. data/lib/switchman/active_record/finder_methods.rb +44 -14
  14. data/lib/switchman/active_record/log_subscriber.rb +11 -5
  15. data/lib/switchman/active_record/migration.rb +45 -3
  16. data/lib/switchman/active_record/pending_migration_connection.rb +17 -0
  17. data/lib/switchman/active_record/persistence.rb +30 -0
  18. data/lib/switchman/active_record/postgresql_adapter.rb +11 -10
  19. data/lib/switchman/active_record/predicate_builder.rb +2 -2
  20. data/lib/switchman/active_record/query_cache.rb +49 -20
  21. data/lib/switchman/active_record/query_methods.rb +192 -135
  22. data/lib/switchman/active_record/relation.rb +28 -13
  23. data/lib/switchman/active_record/spawn_methods.rb +2 -2
  24. data/lib/switchman/active_record/statement_cache.rb +2 -2
  25. data/lib/switchman/active_record/tasks/database_tasks.rb +6 -1
  26. data/lib/switchman/active_record/test_fixtures.rb +26 -16
  27. data/lib/switchman/active_support/cache.rb +9 -4
  28. data/lib/switchman/arel.rb +34 -18
  29. data/lib/switchman/call_super.rb +2 -2
  30. data/lib/switchman/database_server.rb +69 -31
  31. data/lib/switchman/default_shard.rb +14 -3
  32. data/lib/switchman/engine.rb +29 -22
  33. data/lib/switchman/environment.rb +2 -2
  34. data/lib/switchman/errors.rb +13 -0
  35. data/lib/switchman/guard_rail/relation.rb +1 -1
  36. data/lib/switchman/parallel.rb +6 -6
  37. data/lib/switchman/r_spec_helper.rb +12 -11
  38. data/lib/switchman/shard.rb +180 -68
  39. data/lib/switchman/sharded_instrumenter.rb +3 -3
  40. data/lib/switchman/shared_schema_cache.rb +11 -0
  41. data/lib/switchman/standard_error.rb +4 -0
  42. data/lib/switchman/test_helper.rb +2 -2
  43. data/lib/switchman/version.rb +1 -1
  44. data/lib/switchman.rb +27 -15
  45. data/lib/tasks/switchman.rake +96 -60
  46. metadata +35 -45
@@ -3,6 +3,19 @@
3
3
  module Switchman
4
4
  module ActiveRecord
5
5
  module QueryMethods
6
+ # Use this class to prevent a value from getting transposed across shards
7
+ class NonTransposingValue < SimpleDelegator
8
+ def class
9
+ __getobj__.class
10
+ end
11
+
12
+ def is_a?(other)
13
+ return true if other == NonTransposingValue
14
+
15
+ __getobj__.is_a?(other)
16
+ end
17
+ end
18
+
6
19
  # shard_value is one of:
7
20
  # A shard
8
21
  # An array or relation of shards
@@ -12,8 +25,6 @@ module Switchman
12
25
  # :explicit - explicit set on the relation
13
26
  # :association - a special value that scopes from associations use to use slightly different logic
14
27
  # for foreign key transposition
15
- # :to_a - a special value that Relation#to_a uses when querying multiple shards to
16
- # remove primary keys from conditions that aren't applicable to the current shard
17
28
  def shard_value
18
29
  @values[:shard]
19
30
  end
@@ -44,10 +55,7 @@ module Switchman
44
55
  old_primary_shard = primary_shard
45
56
  self.shard_value = value
46
57
  self.shard_source_value = source
47
- if old_primary_shard != primary_shard || source == :to_a
48
- transpose_clauses(old_primary_shard, primary_shard,
49
- remove_nonlocal_primary_keys: source == :to_a)
50
- end
58
+ transpose_predicates(nil, old_primary_shard, primary_shard) if old_primary_shard != primary_shard
51
59
  self
52
60
  end
53
61
 
@@ -89,35 +97,37 @@ module Switchman
89
97
  super(other.shard(primary_shard))
90
98
  end
91
99
 
92
- private
100
+ # use a temp variable so that the new where clause is built before self.where_clause is read,
101
+ # since build_where_clause might mutate self.where_clause
102
+ def where!(opts, *rest)
103
+ new_clause = build_where_clause(opts, rest)
104
+ self.where_clause += new_clause
105
+ self
106
+ end
93
107
 
94
- %i[where having].each do |type|
95
- class_eval <<-RUBY, __FILE__, __LINE__ + 1
96
- def transpose_#{type}_clauses(source_shard, target_shard, remove_nonlocal_primary_keys:)
97
- unless (predicates = #{type}_clause.send(:predicates)).empty?
98
- new_predicates = transpose_predicates(predicates, source_shard,
99
- target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
100
- if new_predicates != predicates
101
- self.#{type}_clause = #{type}_clause.dup
102
- if new_predicates != predicates
103
- #{type}_clause.instance_variable_set(:@predicates, new_predicates)
104
- end
105
- end
106
- end
108
+ protected
109
+
110
+ def remove_nonlocal_primary_keys!
111
+ each_transposable_predicate_value do |value, predicate, _relation, _column, type|
112
+ next value unless
113
+ type == :primary &&
114
+ predicate.left.relation.klass == klass &&
115
+ (predicate.is_a?(::Arel::Nodes::Equality) || predicate.is_a?(::Arel::Nodes::HomogeneousIn))
116
+
117
+ (value.is_a?(Integer) && value > Shard::IDS_PER_SHARD) ? [] : value
107
118
  end
108
- RUBY
119
+ self
109
120
  end
110
121
 
111
- def transpose_clauses(source_shard, target_shard, remove_nonlocal_primary_keys: false)
112
- transpose_where_clauses(source_shard, target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
113
- transpose_having_clauses(source_shard, target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
114
- end
122
+ private
115
123
 
116
124
  def infer_shards_from_primary_key(predicates)
117
125
  return unless klass.integral_id?
118
126
 
119
127
  primary_key = predicates.detect do |predicate|
120
- (predicate.is_a?(::Arel::Nodes::Binary) || predicate.is_a?(::Arel::Nodes::HomogeneousIn)) &&
128
+ (predicate.is_a?(::Arel::Nodes::Equality) ||
129
+ predicate.is_a?(::Arel::Nodes::In) ||
130
+ predicate.is_a?(::Arel::Nodes::HomogeneousIn)) &&
121
131
  predicate.left.is_a?(::Arel::Attributes::Attribute) &&
122
132
  predicate.left.relation.is_a?(::Arel::Table) && predicate.left.relation.klass == klass &&
123
133
  klass.primary_key == predicate.left.name
@@ -145,7 +155,7 @@ module Switchman
145
155
  return
146
156
  else
147
157
  id_shards = id_shards.to_a
148
- transpose_clauses(primary_shard, id_shards.first)
158
+ transpose_predicates(nil, primary_shard, id_shards.first)
149
159
  self.shard_value = id_shards
150
160
  return
151
161
  end
@@ -162,7 +172,7 @@ module Switchman
162
172
 
163
173
  return if !id_shard || id_shard == primary_shard
164
174
 
165
- transpose_clauses(primary_shard, id_shard)
175
+ transpose_predicates(nil, primary_shard, id_shard)
166
176
  self.shard_value = id_shard
167
177
  end
168
178
 
@@ -180,19 +190,19 @@ module Switchman
180
190
  end
181
191
 
182
192
  def sharded_foreign_key?(relation, column)
183
- models_for_table(relation.table_name).any? { |m| m.sharded_column?(column) }
193
+ models_for_table(relation.name).any? { |m| m.sharded_column?(column) }
184
194
  end
185
195
 
186
196
  def sharded_primary_key?(relation, column)
187
197
  column = column.to_s
188
- return column == 'id' if relation.klass == ::ActiveRecord::Base
198
+ return column == "id" if relation.klass == ::ActiveRecord::Base
189
199
 
190
200
  relation.klass.primary_key == column && relation.klass.integral_id?
191
201
  end
192
202
 
193
203
  def source_shard_for_foreign_key(relation, column)
194
204
  reflection = nil
195
- models_for_table(relation.table_name).each do |model|
205
+ models_for_table(relation.name).each do |model|
196
206
  reflection = model.send(:reflection_for_integer_attribute, column)
197
207
  break if reflection
198
208
  end
@@ -212,12 +222,11 @@ module Switchman
212
222
 
213
223
  case opts
214
224
  when String, Array
215
- values = Hash === rest.first ? rest.first.values : rest
225
+ values = (Hash === rest.first) ? rest.first.values : rest
216
226
 
217
- values.grep(ActiveRecord::Relation) do |rel|
218
- # serialize subqueries against the same shard as the outer query is currently
219
- # targeted to run against
220
- rel.shard!(primary_shard) if rel.shard_source_value == :implicit && rel.primary_shard != primary_shard
227
+ if shard_source_value != :explicit && values.grep(ActiveRecord::Relation).first
228
+ raise "Sub-queries are not allowed as simple substitutions; " \
229
+ "please build your relation with more structured methods so that Switchman is able to introspect it."
221
230
  end
222
231
 
223
232
  super
@@ -246,122 +255,170 @@ module Switchman
246
255
  connection.with_global_table_name { super }
247
256
  end
248
257
 
249
- def transpose_predicates(predicates,
250
- source_shard,
251
- target_shard,
252
- remove_nonlocal_primary_keys: false)
253
- predicates.map do |predicate|
254
- transpose_single_predicate(predicate, source_shard, target_shard,
255
- remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
256
- end
258
+ def each_predicate(predicates = nil, &block)
259
+ return predicates.map(&block) if predicates
260
+
261
+ each_predicate_cb(:having_clause, :having_clause=, &block)
262
+ each_predicate_cb(:where_clause, :where_clause=, &block)
257
263
  end
258
264
 
259
- def transpose_single_predicate(predicate,
260
- source_shard,
261
- target_shard,
262
- remove_nonlocal_primary_keys: false)
263
- if predicate.is_a?(::Arel::Nodes::Grouping)
264
- return predicate unless predicate.expr.is_a?(::Arel::Nodes::Or)
265
-
266
- # Dang, we have an OR. OK, that means we have other epxressions below this
267
- # level, perhaps many, that may need transposition.
268
- # the left side and right side must each be treated as predicate lists and
269
- # transformed in kind, if neither of them changes we can just return the grouping as is.
270
- # hold on, it's about to get recursive...
271
- or_expr = predicate.expr
272
- left_node = or_expr.left
273
- right_node = or_expr.right
274
- new_left_predicates = transpose_single_predicate(left_node, source_shard,
275
- target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
276
- new_right_predicates = transpose_single_predicate(right_node, source_shard,
277
- target_shard, remove_nonlocal_primary_keys: remove_nonlocal_primary_keys)
278
- return predicate if new_left_predicates == left_node && new_right_predicates == right_node
279
-
280
- return ::Arel::Nodes::Grouping.new ::Arel::Nodes::Or.new(new_left_predicates, new_right_predicates)
281
- end
282
- return predicate unless predicate.is_a?(::Arel::Nodes::Binary) || predicate.is_a?(::Arel::Nodes::HomogeneousIn)
283
- return predicate unless predicate.left.is_a?(::Arel::Attributes::Attribute)
284
-
285
- relation, column = relation_and_column(predicate.left)
286
- return predicate unless (type = transposable_attribute_type(relation, column))
287
-
288
- remove = true if type == :primary &&
289
- remove_nonlocal_primary_keys &&
290
- predicate.left.relation.klass == klass &&
291
- (predicate.is_a?(::Arel::Nodes::Equality) || predicate.is_a?(::Arel::Nodes::HomogeneousIn))
292
-
293
- current_source_shard =
294
- if source_shard
295
- source_shard
296
- elsif type == :primary
297
- Shard.current(klass.connection_class_for_self)
298
- elsif type == :foreign
299
- source_shard_for_foreign_key(relation, column)
300
- end
265
+ def each_predicate_cb(clause_getter, clause_setter, &block)
266
+ old_clause = send(clause_getter)
267
+ old_predicates = old_clause.send(:predicates)
268
+ return if old_predicates.empty?
301
269
 
302
- right = if predicate.is_a?(::Arel::Nodes::HomogeneousIn)
303
- predicate.values
304
- else
305
- predicate.right
306
- end
270
+ new_predicates = old_predicates.map(&block)
271
+ return if new_predicates == old_predicates
307
272
 
308
- new_right_value =
309
- case right
310
- when Array
311
- right.map { |val| transpose_predicate_value(val, current_source_shard, target_shard, type, remove).presence }.compact
312
- else
313
- transpose_predicate_value(right, current_source_shard, target_shard, type, remove)
273
+ new_clause = old_clause.dup
274
+ new_clause.instance_variable_set(:@predicates, new_predicates)
275
+
276
+ send(clause_setter, new_clause)
277
+ end
278
+
279
+ def each_transposable_predicate(predicates, &block)
280
+ each_predicate(predicates) do |predicate|
281
+ case predicate
282
+ when ::Arel::Nodes::Grouping
283
+ next predicate unless predicate.expr.is_a?(::Arel::Nodes::Or)
284
+
285
+ or_expr = predicate.expr
286
+ old_left = or_expr.left
287
+ old_right = or_expr.right
288
+ new_left, new_right = each_transposable_predicate([old_left, old_right], &block)
289
+
290
+ next predicate if new_left == old_left && new_right == old_right
291
+
292
+ next predicate.class.new predicate.expr.class.new(new_left, new_right)
293
+ when ::Arel::Nodes::SelectStatement
294
+ new_cores = predicate.cores.map do |core|
295
+ next core unless core.is_a?(::Arel::Nodes::SelectCore) # just in case something weird is going on
296
+
297
+ new_wheres = each_transposable_predicate(core.wheres, &block)
298
+ new_havings = each_transposable_predicate(core.havings, &block)
299
+
300
+ next core if core.wheres == new_wheres && core.havings == new_havings
301
+
302
+ new_core = core.clone
303
+ new_core.wheres = new_wheres
304
+ new_core.havings = new_havings
305
+ new_core
306
+ end
307
+
308
+ next predicate if predicate.cores == new_cores
309
+
310
+ new_node = predicate.clone
311
+ new_node.instance_variable_set(:@cores, new_cores)
312
+ next new_node
313
+ when ::Arel::Nodes::Not
314
+ old_value = predicate.expr
315
+ new_value = each_transposable_predicate([old_value], &block).first
316
+
317
+ next predicate if old_value == new_value
318
+
319
+ next predicate.class.new(new_value)
320
+ when ::Arel::Nodes::Exists
321
+ old_value = predicate.expressions
322
+ new_value = each_transposable_predicate([old_value], &block).first
323
+
324
+ next predicate if old_value == new_value
325
+
326
+ next predicate.class.new(new_value)
314
327
  end
315
328
 
316
- if new_right_value == right
317
- predicate
318
- elsif predicate.right.is_a?(::Arel::Nodes::Casted)
319
- if new_right_value == right.value
320
- predicate
321
- else
322
- predicate.class.new(predicate.left, right.class.new(new_right_value, right.attribute))
329
+ next predicate unless predicate.is_a?(::Arel::Nodes::Binary) || predicate.is_a?(::Arel::Nodes::HomogeneousIn)
330
+ next predicate unless predicate.left.is_a?(::Arel::Attributes::Attribute)
331
+
332
+ relation, column = relation_and_column(predicate.left)
333
+ next predicate unless (type = transposable_attribute_type(relation, column))
334
+
335
+ yield(predicate, relation, column, type)
336
+ end
337
+ end
338
+
339
+ def each_transposable_predicate_value(predicates = nil, &block)
340
+ each_transposable_predicate(predicates) do |predicate, relation, column, type|
341
+ each_transposable_predicate_value_cb(predicate, block) do |value|
342
+ yield(value, predicate, relation, column, type)
323
343
  end
324
- elsif predicate.is_a?(::Arel::Nodes::HomogeneousIn)
344
+ end
345
+ end
346
+
347
+ def each_transposable_predicate_value_cb(node, original_block, &block)
348
+ case node
349
+ when Array
350
+ node.filter_map { |val| each_transposable_predicate_value_cb(val, original_block, &block).presence }
351
+ when ::ActiveModel::Attribute
352
+ old_value = node.value_before_type_cast
353
+ new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
354
+
355
+ (old_value == new_value) ? node : node.class.new(node.name, new_value, node.type)
356
+ when ::Arel::Nodes::And
357
+ old_value = node.children
358
+ new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
359
+
360
+ (old_value == new_value) ? node : node.class.new(new_value)
361
+ when ::Arel::Nodes::BindParam
362
+ old_value = node.value
363
+ new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
364
+
365
+ (old_value == new_value) ? node : node.class.new(new_value)
366
+ when ::Arel::Nodes::Casted
367
+ old_value = node.value
368
+ new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
369
+
370
+ (old_value == new_value) ? node : node.class.new(new_value, node.attribute)
371
+ when ::Arel::Nodes::HomogeneousIn
372
+ old_value = node.values
373
+ new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
374
+
325
375
  # switch to a regular In, so that Relation::WhereClause#contradiction? knows about it
326
- if new_right_value.empty?
327
- klass = predicate.type == :in ? ::Arel::Nodes::In : ::Arel::Nodes::NotIn
328
- klass.new(predicate.attribute, new_right_value)
376
+ if new_value.empty?
377
+ klass = (node.type == :in) ? ::Arel::Nodes::In : ::Arel::Nodes::NotIn
378
+ klass.new(node.attribute, new_value)
329
379
  else
330
- predicate.class.new(new_right_value, predicate.attribute, predicate.type)
380
+ (old_value == new_value) ? node : node.class.new(new_value, node.attribute, node.type)
331
381
  end
382
+ when ::Arel::Nodes::Binary
383
+ old_value = node.right
384
+ new_value = each_transposable_predicate_value_cb(old_value, original_block, &block)
385
+
386
+ (old_value == new_value) ? node : node.class.new(node.left, new_value)
387
+ when ::Arel::Nodes::SelectStatement
388
+ each_transposable_predicate_value([node], &original_block).first
332
389
  else
333
- predicate.class.new(predicate.left, new_right_value)
390
+ yield(node)
334
391
  end
335
392
  end
336
393
 
337
- def transpose_predicate_value(value, current_shard, target_shard, attribute_type, remove_non_local_ids)
338
- case value
339
- when ::Arel::Nodes::BindParam, ::ActiveModel::Attribute
340
- query_att = value.is_a?(::ActiveModel::Attribute) ? value : value.value
341
- current_id = query_att.value_before_type_cast
342
- if current_id.is_a?(::ActiveRecord::StatementCache::Substitute)
343
- current_id.sharded = true # mark for transposition later
344
- current_id.primary = true if attribute_type == :primary
345
- value
346
- else
347
- local_id = Shard.relative_id_for(current_id, current_shard, target_shard) || current_id
348
- local_id = [] if remove_non_local_ids && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
349
- if current_id == local_id
350
- # make a new bind param
351
- value
394
+ def transpose_predicates(predicates,
395
+ source_shard,
396
+ target_shard)
397
+ each_transposable_predicate_value(predicates) do |value, _predicate, relation, column, type|
398
+ current_source_shard =
399
+ if source_shard
400
+ source_shard
401
+ elsif type == :primary
402
+ Shard.current(klass.connection_class_for_self)
403
+ elsif type == :foreign
404
+ source_shard_for_foreign_key(relation, column)
352
405
  else
353
- new_att = query_att.class.new(query_att.name, local_id, query_att.type)
354
- if value.is_a?(::ActiveModel::Attribute)
355
- new_att
356
- else
357
- ::Arel::Nodes::BindParam.new(new_att)
358
- end
406
+ primary_shard
359
407
  end
360
- end
408
+
409
+ transpose_predicate_value(value, current_source_shard, target_shard, type)
410
+ end
411
+ end
412
+
413
+ def transpose_predicate_value(value, current_shard, target_shard, attribute_type)
414
+ if value.is_a?(NonTransposingValue)
415
+ value
416
+ elsif value.is_a?(::ActiveRecord::StatementCache::Substitute)
417
+ value.sharded = true # mark for transposition later
418
+ value.primary = true if attribute_type == :primary
419
+ value
361
420
  else
362
- local_id = Shard.relative_id_for(value, current_shard, target_shard) || value
363
- local_id = [] if remove_non_local_ids && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
364
- local_id
421
+ Shard.relative_id_for(value, current_shard, target_shard) || value
365
422
  end
366
423
  end
367
424
  end
@@ -4,7 +4,7 @@ module Switchman
4
4
  module ActiveRecord
5
5
  module Relation
6
6
  def self.prepended(klass)
7
- klass::SINGLE_VALUE_METHODS.concat %i[shard shard_source]
7
+ klass::SINGLE_VALUE_METHODS.push(:shard, :shard_source)
8
8
  end
9
9
 
10
10
  def initialize(*, **)
@@ -44,12 +44,18 @@ module Switchman
44
44
  primary_shard.activate(klass.connection_class_for_self) { super }
45
45
  end
46
46
 
47
- def explain
48
- activate { |relation| relation.call_super(:explain, Relation) }
47
+ if ::Rails.version > "7.1.2"
48
+ def transaction(...)
49
+ primary_shard.activate(klass.connection_class_for_self) { super }
50
+ end
51
+ end
52
+
53
+ def explain(*args)
54
+ activate { |relation| relation.call_super(:explain, Relation, *args) }
49
55
  end
50
56
 
51
57
  def load(&block)
52
- if !loaded? || (::Rails.version >= '7.0' && scheduled?)
58
+ if !loaded? || scheduled?
53
59
  @records = activate { |relation| relation.send(:exec_queries, &block) }
54
60
  @loaded = true
55
61
  end
@@ -59,8 +65,8 @@ module Switchman
59
65
 
60
66
  %I[update_all delete_all].each do |method|
61
67
  class_eval <<-RUBY, __FILE__, __LINE__ + 1
62
- def #{method}(*args)
63
- result = self.activate(unordered: true) { |relation| relation.call_super(#{method.inspect}, Relation, *args) }
68
+ def #{method}(*args, **kwargs)
69
+ result = self.activate(unordered: true) { |relation| relation.call_super(#{method.inspect}, Relation, *args, **kwargs) }
64
70
  result = result.sum if result.is_a?(Array)
65
71
  result
66
72
  end
@@ -72,12 +78,16 @@ module Switchman
72
78
  loose_mode = options[:loose] && is_integer
73
79
  # loose_mode: if we don't care about getting exactly batch_size ids in between
74
80
  # don't get the max - just get the min and add batch_size so we get that many _at most_
75
- values = loose_mode ? 'MIN(id)' : 'MIN(id), MAX(id)'
81
+ values = loose_mode ? "MIN(id)" : "MIN(id), MAX(id)"
76
82
 
77
83
  batch_size = options[:batch_size].try(:to_i) || 1000
78
- quoted_primary_key = "#{klass.connection.quote_local_table_name(table_name)}.#{klass.connection.quote_column_name(primary_key)}"
79
- as_id = ' AS id' unless primary_key == 'id'
80
- subquery_scope = except(:select).select("#{quoted_primary_key}#{as_id}").reorder(primary_key.to_sym).limit(loose_mode ? 1 : batch_size)
84
+ quoted_primary_key =
85
+ "#{klass.connection.quote_local_table_name(table_name)}.#{klass.connection.quote_column_name(primary_key)}"
86
+ as_id = " AS id" unless primary_key == "id"
87
+ subquery_scope = except(:select)
88
+ .select("#{quoted_primary_key}#{as_id}")
89
+ .reorder(primary_key.to_sym)
90
+ .limit(loose_mode ? 1 : batch_size)
81
91
  subquery_scope = subquery_scope.where("#{quoted_primary_key} <= ?", options[:end_at]) if options[:end_at]
82
92
 
83
93
  first_subquery_scope = if options[:start_at]
@@ -103,7 +113,9 @@ module Switchman
103
113
  def activate(unordered: false, &block)
104
114
  shards = all_shards
105
115
  if Array === shards && shards.length == 1
106
- if shards.first == DefaultShard || shards.first == Shard.current(klass.connection_class_for_self)
116
+ if !loaded? && shard_value != shards.first
117
+ shard(shards.first).activate(&block)
118
+ elsif shards.first == DefaultShard || shards.first == Shard.current(klass.connection_class_for_self)
107
119
  yield(self, shards.first)
108
120
  else
109
121
  shards.first.activate(klass.connection_class_for_self) { yield(self, shards.first) }
@@ -115,9 +127,12 @@ module Switchman
115
127
  # don't even query other shards if we're already past the limit
116
128
  next if limit_value && result_count >= limit_value && order_values.empty?
117
129
 
118
- relation = shard(Shard.current(klass.connection_class_for_self), :to_a)
130
+ relation = shard(Shard.current(klass.connection_class_for_self))
131
+ relation.remove_nonlocal_primary_keys!
119
132
  # do a minimal query if possible
120
- relation = relation.limit(limit_value - result_count) if limit_value && !result_count.zero? && order_values.empty?
133
+ if limit_value && !result_count.zero? && order_values.empty?
134
+ relation = relation.limit(limit_value - result_count)
135
+ end
121
136
 
122
137
  shard_results = relation.activate(&block)
123
138
 
@@ -17,7 +17,7 @@ module Switchman
17
17
  final_shard_source_value = %i[explicit association].detect do |source_value|
18
18
  shard_source_value == source_value || rhs.shard_source_value == source_value
19
19
  end
20
- raise 'unknown shard_source_value' unless final_shard_source_value
20
+ raise "unknown shard_source_value" unless final_shard_source_value
21
21
 
22
22
  # have to merge shard_value
23
23
  lhs_shard_value = all_shards
@@ -36,7 +36,7 @@ module Switchman
36
36
  final_shard_source_value = %i[explicit association implicit].detect do |source_value|
37
37
  shard_source_value == source_value || rhs.shard_source_value == source_value
38
38
  end
39
- raise 'unknown shard_source_value' unless final_shard_source_value
39
+ raise "unknown shard_source_value" unless final_shard_source_value
40
40
  end
41
41
 
42
42
  [final_shard_value, final_primary_shard, final_shard_source_value]
@@ -4,8 +4,8 @@ module Switchman
4
4
  module ActiveRecord
5
5
  module StatementCache
6
6
  module ClassMethods
7
- def create(connection, &block)
8
- relation = block.call ::ActiveRecord::StatementCache::Params.new
7
+ def create(connection)
8
+ relation = yield ::ActiveRecord::StatementCache::Params.new
9
9
 
10
10
  _query_builder, binds = connection.cacheable_query(self, relation.arel)
11
11
  bind_map = ::ActiveRecord::StatementCache::BindMap.new(binds)
@@ -7,9 +7,14 @@ module Switchman
7
7
  def drop(*)
8
8
  super
9
9
  # no really, it's gone
10
- Switchman.cache.delete('default_shard')
10
+ Switchman.cache.delete("default_shard")
11
11
  Shard.default(reload: true)
12
12
  end
13
+
14
+ def raise_for_multi_db(*)
15
+ # ignore; Switchman doesn't use namespaced tasks for multiple shards; it uses
16
+ # environment variables to filter which shards you want to target
17
+ end
13
18
  end
14
19
  end
15
20
  end
@@ -12,31 +12,41 @@ module Switchman
12
12
  # Replace the one that activerecord natively uses with a switchman-optimized one
13
13
  ::ActiveSupport::Notifications.unsubscribe(@connection_subscriber)
14
14
  # Code adapted from the code in rails proper
15
- @connection_subscriber = ::ActiveSupport::Notifications.subscribe('!connection.active_record') do |_, _, _, _, payload|
16
- spec_name = payload[:spec_name] if payload.key?(:spec_name)
17
- shard = payload[:shard] if payload.key?(:shard)
18
- setup_shared_connection_pool
15
+ @connection_subscriber =
16
+ ::ActiveSupport::Notifications.subscribe("!connection.active_record") do |_, _, _, _, payload|
17
+ spec_name = if ::Rails.version < "7.1"
18
+ payload[:spec_name] if payload.key?(:spec_name)
19
+ elsif payload.key?(:connection_name)
20
+ payload[:connection_name]
21
+ end
22
+ shard = payload[:shard] if payload.key?(:shard)
19
23
 
20
- if spec_name && !FORBIDDEN_DB_ENVS.include?(shard)
21
- begin
22
- connection = ::ActiveRecord::Base.connection_handler.retrieve_connection(spec_name, shard: shard)
23
- rescue ::ActiveRecord::ConnectionNotEstablished, ::ActiveRecord::NoDatabaseError
24
- connection = nil
25
- end
24
+ if spec_name && !FORBIDDEN_DB_ENVS.include?(shard)
25
+ begin
26
+ connection = ::ActiveRecord::Base.connection_handler.retrieve_connection(spec_name, shard: shard)
27
+ connection.connect! if ::Rails.version >= "7.1" # eagerly validate the connection
28
+ rescue ::ActiveRecord::ConnectionNotEstablished, ::ActiveRecord::NoDatabaseError
29
+ connection = nil
30
+ end
26
31
 
27
- if connection && !@fixture_connections.include?(connection)
28
- connection.begin_transaction joinable: false, _lazy: false
29
- connection.pool.lock_thread = true if lock_threads
30
- @fixture_connections << connection
32
+ if connection
33
+ setup_shared_connection_pool
34
+ unless @fixture_connections.include?(connection)
35
+ connection.begin_transaction joinable: false, _lazy: false
36
+ connection.pool.lock_thread = true if lock_threads
37
+ @fixture_connections << connection
38
+ end
39
+ end
31
40
  end
32
41
  end
33
- end
34
42
  end
35
43
 
36
44
  def enlist_fixture_connections
37
45
  setup_shared_connection_pool
38
46
 
39
- ::ActiveRecord::Base.connection_handler.connection_pool_list.reject { |cp| FORBIDDEN_DB_ENVS.include?(cp.db_config.env_name.to_sym) }.map(&:connection)
47
+ ::ActiveRecord::Base.connection_handler.connection_pool_list(:primary).reject do |cp|
48
+ FORBIDDEN_DB_ENVS.include?(cp.db_config.env_name.to_sym)
49
+ end.map(&:connection)
40
50
  end
41
51
  end
42
52
  end
@@ -22,9 +22,14 @@ module Switchman
22
22
 
23
23
  def lookup_store(*store_options)
24
24
  store = super
25
- # can't use defined?, because it's a _ruby_ autoloaded constant,
26
- # so just checking that will cause it to get required
27
- ::ActiveSupport::Cache::RedisCacheStore.prepend(RedisCacheStore) if store.instance_of?(ActiveSupport::Cache::RedisCacheStore) && !::ActiveSupport::Cache::RedisCacheStore.ancestors.include?(RedisCacheStore)
25
+ # must use the string name, otherwise it will try to auto-load the constant
26
+ # and we don't want to require redis in this file (since it's not a hard dependency)
27
+ # rubocop:disable Style/ClassEqualityComparison
28
+ if store.class.name == "ActiveSupport::Cache::RedisCacheStore" &&
29
+ !(::ActiveSupport::Cache::RedisCacheStore <= RedisCacheStore)
30
+ ::ActiveSupport::Cache::RedisCacheStore.prepend(RedisCacheStore)
31
+ end
32
+ # rubocop:enable Style/ClassEqualityComparison
28
33
  store.options[:namespace] ||= -> { Shard.current.default? ? nil : "shard_#{Shard.current.id}" }
29
34
  store
30
35
  end
@@ -33,7 +38,7 @@ module Switchman
33
38
  module RedisCacheStore
34
39
  def clear(namespace: nil, **)
35
40
  # RedisCacheStore tries to be smart and only clear the cache under your namespace, if you have one set
36
- # unfortunately, it uses the keys command, which is extraordinarily inefficient in a large redis instance
41
+ # unfortunately, it doesn't work using redis clustering because of the way redis keys are distributed
37
42
  # fortunately, we can assume we control the entire instance, because we set up the namespacing, so just
38
43
  # always unset it temporarily for clear calls
39
44
  namespace = nil # rubocop:disable Lint/ShadowedArgument