switchman 1.5.21 → 2.1.5

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 (56) hide show
  1. checksums.yaml +5 -5
  2. data/app/models/switchman/shard.rb +757 -11
  3. data/db/migrate/20130328212039_create_switchman_shards.rb +3 -1
  4. data/db/migrate/20130328224244_create_default_shard.rb +4 -2
  5. data/db/migrate/20161206323434_add_back_default_string_limits_switchman.rb +13 -0
  6. data/db/migrate/20180828183945_add_default_shard_index.rb +15 -0
  7. data/db/migrate/20180828192111_add_timestamps_to_shards.rb +17 -0
  8. data/db/migrate/20190114212900_add_unique_name_indexes.rb +9 -0
  9. data/lib/switchman/action_controller/caching.rb +2 -0
  10. data/lib/switchman/active_record/abstract_adapter.rb +14 -4
  11. data/lib/switchman/active_record/association.rb +64 -37
  12. data/lib/switchman/active_record/attribute_methods.rb +54 -22
  13. data/lib/switchman/active_record/base.rb +76 -31
  14. data/lib/switchman/active_record/batches.rb +3 -1
  15. data/lib/switchman/active_record/calculations.rb +17 -22
  16. data/lib/switchman/active_record/connection_handler.rb +88 -78
  17. data/lib/switchman/active_record/connection_pool.rb +28 -23
  18. data/lib/switchman/active_record/finder_methods.rb +37 -28
  19. data/lib/switchman/active_record/log_subscriber.rb +14 -19
  20. data/lib/switchman/active_record/migration.rb +80 -0
  21. data/lib/switchman/active_record/model_schema.rb +3 -1
  22. data/lib/switchman/active_record/persistence.rb +9 -1
  23. data/lib/switchman/active_record/postgresql_adapter.rb +170 -126
  24. data/lib/switchman/active_record/predicate_builder.rb +3 -1
  25. data/lib/switchman/active_record/query_cache.rb +22 -87
  26. data/lib/switchman/active_record/query_methods.rb +139 -125
  27. data/lib/switchman/active_record/reflection.rb +42 -14
  28. data/lib/switchman/active_record/relation.rb +108 -33
  29. data/lib/switchman/active_record/spawn_methods.rb +2 -0
  30. data/lib/switchman/active_record/statement_cache.rb +44 -52
  31. data/lib/switchman/active_record/table_definition.rb +4 -2
  32. data/lib/switchman/active_record/type_caster.rb +2 -0
  33. data/lib/switchman/active_record/where_clause_factory.rb +5 -2
  34. data/lib/switchman/active_support/cache.rb +18 -0
  35. data/lib/switchman/arel.rb +8 -25
  36. data/lib/switchman/call_super.rb +19 -0
  37. data/lib/switchman/connection_pool_proxy.rb +70 -24
  38. data/lib/switchman/database_server.rb +69 -59
  39. data/lib/switchman/default_shard.rb +3 -0
  40. data/lib/switchman/engine.rb +44 -41
  41. data/lib/switchman/environment.rb +2 -0
  42. data/lib/switchman/errors.rb +2 -0
  43. data/lib/switchman/{shackles → guard_rail}/relation.rb +7 -5
  44. data/lib/switchman/{shackles.rb → guard_rail.rb} +6 -4
  45. data/lib/switchman/open4.rb +2 -0
  46. data/lib/switchman/r_spec_helper.rb +14 -8
  47. data/lib/switchman/rails.rb +2 -0
  48. data/lib/switchman/schema_cache.rb +17 -0
  49. data/lib/switchman/sharded_instrumenter.rb +4 -2
  50. data/lib/switchman/standard_error.rb +4 -2
  51. data/lib/switchman/test_helper.rb +7 -10
  52. data/lib/switchman/version.rb +3 -1
  53. data/lib/switchman.rb +5 -1
  54. data/lib/tasks/switchman.rake +53 -72
  55. metadata +84 -38
  56. data/app/models/switchman/shard_internal.rb +0 -692
@@ -1,3 +1,5 @@
1
+ # frozen_string_literal: true
2
+
1
3
  module Switchman
2
4
  module ActiveRecord
3
5
  module QueryMethods
@@ -19,11 +21,11 @@ module Switchman
19
21
  @values[:shard_source]
20
22
  end
21
23
  def shard_value=(value)
22
- raise ImmutableRelation if @loaded
24
+ raise ::ActiveRecord::ImmutableRelation if @loaded
23
25
  @values[:shard] = value
24
26
  end
25
27
  def shard_source_value=(value)
26
- raise ImmutableRelation if @loaded
28
+ raise ::ActiveRecord::ImmutableRelation if @loaded
27
29
  @values[:shard_source] = value
28
30
  end
29
31
 
@@ -42,34 +44,6 @@ module Switchman
42
44
  self
43
45
  end
44
46
 
45
- if ::Rails.version < '5'
46
- # moved to WhereClauseFactory#build in Rails 5
47
- def build_where(opts, other = [])
48
- case opts
49
- when String, Array
50
- values = Hash === other.first ? other.first.values : other
51
-
52
- values.grep(ActiveRecord::Relation) do |rel|
53
- # serialize subqueries against the same shard as the outer query is currently
54
- # targeted to run against
55
- if rel.shard_source_value == :implicit && rel.primary_shard != primary_shard
56
- rel.shard!(primary_shard)
57
- end
58
- self.bind_values += rel.bind_values if ::Rails.version < '4.2'
59
- end
60
-
61
- [@klass.send(:sanitize_sql, other.empty? ? opts : ([opts] + other))]
62
- when Hash, ::Arel::Nodes::Node
63
- predicates = super
64
- infer_shards_from_primary_key(predicates) if shard_source_value == :implicit && shard_value.is_a?(Shard)
65
- predicates = transpose_predicates(predicates, nil, primary_shard) if shard_source_value != :explicit
66
- predicates
67
- else
68
- super
69
- end
70
- end
71
- end
72
-
73
47
  # the shard that where_values are relative to. if it's multiple shards, they're stored
74
48
  # relative to the first shard
75
49
  def primary_shard
@@ -104,44 +78,47 @@ module Switchman
104
78
  end
105
79
  end
106
80
 
107
- if ::Rails.version >= '4.2' && ::Rails.version < '5'
108
- # fixes an issue in Rails 4.2 with `reverse_sql_order` and qualified names
109
- # where quoted_table_name is called before shard(s) have been activated
110
- # if there's no ordering
111
- def reverse_order!
112
- orders = order_values.uniq
113
- orders.reject!(&:blank?)
114
- if orders.empty?
115
- self.order_values = [arel_table[primary_key].desc]
116
- else
117
- self.order_values = reverse_sql_order(orders)
118
- end
119
- self
120
- end
121
- end
122
-
123
81
  private
124
82
 
125
- [:where, :having].each do |type|
126
- class_eval <<-RUBY, __FILE__, __LINE__ + 1
83
+ if ::Rails.version >= '5.2'
84
+ [:where, :having].each do |type|
85
+ class_eval <<-RUBY, __FILE__, __LINE__ + 1
127
86
  def transpose_#{type}_clauses(source_shard, target_shard, remove_nonlocal_primary_keys)
128
- if ::Rails.version >= '5'
129
- unless (predicates = #{type}_clause.send(:predicates)).empty?
130
- new_predicates = transpose_predicates(predicates, source_shard,
131
- target_shard, remove_nonlocal_primary_keys)
87
+ unless (predicates = #{type}_clause.send(:predicates)).empty?
88
+ new_predicates, _binds = transpose_predicates(predicates, source_shard,
89
+ target_shard, remove_nonlocal_primary_keys)
90
+ if new_predicates != predicates
91
+ self.#{type}_clause = #{type}_clause.dup
132
92
  if new_predicates != predicates
133
- self.#{type}_clause = #{type}_clause.dup
134
93
  #{type}_clause.instance_variable_set(:@predicates, new_predicates)
135
94
  end
136
95
  end
137
- else
138
- unless #{type}_values.empty?
139
- self.#{type}_values = transpose_predicates(#{type}_values,
140
- source_shard, target_shard, remove_nonlocal_primary_keys)
141
- end
142
96
  end
143
97
  end
144
- RUBY
98
+ RUBY
99
+ end
100
+ else
101
+ [:where, :having].each do |type|
102
+ class_eval <<-RUBY, __FILE__, __LINE__ + 1
103
+ def transpose_#{type}_clauses(source_shard, target_shard, remove_nonlocal_primary_keys)
104
+ unless (predicates = #{type}_clause.send(:predicates)).empty?
105
+ new_predicates, new_binds = transpose_predicates(predicates, source_shard,
106
+ target_shard, remove_nonlocal_primary_keys,
107
+ binds: #{type}_clause.binds,
108
+ dup_binds_on_mutation: true)
109
+ if new_predicates != predicates || !new_binds.equal?(#{type}_clause.binds)
110
+ self.#{type}_clause = #{type}_clause.dup
111
+ if new_predicates != predicates
112
+ #{type}_clause.instance_variable_set(:@predicates, new_predicates)
113
+ end
114
+ if !new_binds.equal?(#{type}_clause.binds)
115
+ #{type}_clause.instance_variable_set(:@binds, new_binds)
116
+ end
117
+ end
118
+ end
119
+ end
120
+ RUBY
121
+ end
145
122
  end
146
123
 
147
124
  def transpose_clauses(source_shard, target_shard, remove_nonlocal_primary_keys = false)
@@ -150,6 +127,8 @@ module Switchman
150
127
  end
151
128
 
152
129
  def infer_shards_from_primary_key(predicates, binds = nil)
130
+ return unless klass.integral_id?
131
+
153
132
  primary_key = predicates.detect do |predicate|
154
133
  predicate.is_a?(::Arel::Nodes::Binary) && predicate.left.is_a?(::Arel::Attributes::Attribute) &&
155
134
  predicate.left.relation.is_a?(::Arel::Table) && predicate.left.relation.model == klass &&
@@ -180,28 +159,23 @@ module Switchman
180
159
  return
181
160
  end
182
161
  when ::Arel::Nodes::BindParam
183
- # look for a bind param with a matching column name
184
- if ::Rails.version >= "5"
185
- binds ||= where_clause.binds + having_clause.binds
186
- if binds && bind = binds.detect{|b| b.try(:name).to_s == klass.primary_key.to_s}
162
+ if ::Rails.version >= "5.2"
163
+ local_id, id_shard = Shard.local_id_for(primary_key.right.value.value_before_type_cast)
164
+ id_shard ||= Shard.current(klass.shard_category) if local_id
165
+ else
166
+ # look for a bind param with a matching column name
167
+ if binds && bind = binds.detect{|b| b&.name.to_s == klass.primary_key.to_s}
187
168
  unless bind.value.is_a?(::ActiveRecord::StatementCache::Substitute)
188
169
  local_id, id_shard = Shard.local_id_for(bind.value)
189
170
  id_shard ||= Shard.current(klass.shard_category) if local_id
190
171
  end
191
172
  end
192
- else
193
- if bind_values && idx = bind_values.find_index{|b| b.is_a?(Array) && b.first.try(:name).to_s == klass.primary_key.to_s}
194
- column, value = bind_values[idx]
195
- unless ::Rails.version >= '4.2' && value.is_a?(::ActiveRecord::StatementCache::Substitute)
196
- local_id, id_shard = Shard.local_id_for(value)
197
- id_shard ||= Shard.current(klass.shard_category) if local_id
198
- end
199
- end
200
173
  end
201
174
  else
202
175
  local_id, id_shard = Shard.local_id_for(primary_key.right)
203
176
  id_shard ||= Shard.current(klass.shard_category) if local_id
204
177
  end
178
+
205
179
  return if !id_shard || id_shard == primary_shard
206
180
  transpose_clauses(primary_shard, id_shard)
207
181
  self.shard_value = id_shard
@@ -252,25 +226,38 @@ module Switchman
252
226
  end
253
227
 
254
228
  def arel_columns(columns)
255
- columns.map do |field|
256
- if (Symbol === field || String === field) && ::Rails.version >= '5' && (klass.has_attribute?(field) || klass.attribute_alias?(field)) && !from_clause.value
257
- klass.arel_attribute(field, table)
258
- elsif (Symbol === field || String === field) && ::Rails.version < '5' && columns_hash.key?(field.to_s) && !from_value
259
- arel_table[field]
260
- elsif Symbol === field
261
- # the rest of this is pulled from AR - the only change is from quote_table_name to quote_column_name here
262
- # otherwise qualified names will add the schema to a column
263
- connection.quote_column_name(field.to_s)
264
- else
265
- field
266
- end
267
- end
229
+ connection.with_local_table_name { super }
230
+ end
231
+
232
+ def arel_column(columns)
233
+ connection.with_local_table_name { super }
234
+ end
235
+
236
+ def table_name_matches?(from)
237
+ connection.with_global_table_name { super }
268
238
  end
269
239
 
270
240
  # semi-private
271
241
  public
272
- def transpose_predicates(predicates, source_shard, target_shard, remove_nonlocal_primary_keys = false, binds = nil)
273
- predicates.map do |predicate|
242
+ def transpose_predicates(predicates,
243
+ source_shard,
244
+ target_shard,
245
+ remove_nonlocal_primary_keys = false,
246
+ binds: nil,
247
+ dup_binds_on_mutation: false)
248
+ result = predicates.map do |predicate|
249
+ if ::Rails.version >= '5.2' && predicate.is_a?(::Arel::Nodes::And)
250
+ new_predicates, _binds = transpose_predicates(predicate.children, source_shard, target_shard,
251
+ remove_nonlocal_primary_keys,
252
+ binds: binds,
253
+ dup_binds_on_mutation: dup_binds_on_mutation)
254
+ next (if new_predicates == predicate.children
255
+ predicate
256
+ else
257
+ ::Arel::Nodes::And.new(new_predicates)
258
+ end)
259
+ end
260
+
274
261
  next predicate unless predicate.is_a?(::Arel::Nodes::Binary)
275
262
  next predicate unless predicate.left.is_a?(::Arel::Attributes::Attribute)
276
263
  relation, column = relation_and_column(predicate.left)
@@ -279,7 +266,7 @@ module Switchman
279
266
  remove = true if type == :primary &&
280
267
  remove_nonlocal_primary_keys &&
281
268
  predicate.left.relation.model == klass &&
282
- predicate.is_a?(::Arel::Nodes::Equality)
269
+ (predicate.is_a?(::Arel::Nodes::Equality) || predicate.is_a?(::Arel::Nodes::In))
283
270
 
284
271
  current_source_shard =
285
272
  if source_shard
@@ -290,29 +277,42 @@ module Switchman
290
277
  source_shard_for_foreign_key(relation, column)
291
278
  end
292
279
 
293
- new_right_value = case predicate.right
294
- when Array
295
- local_ids = []
296
- predicate.right.each do |value|
297
- local_id = Shard.relative_id_for(value, current_source_shard, target_shard)
298
- next unless local_id
299
- unless remove && local_id > Shard::IDS_PER_SHARD
300
- if ::Rails.version > "4.2" && value.is_a?(::Arel::Nodes::Casted)
301
- if local_id == value.val
302
- local_id = value
303
- elsif local_id != value
304
- local_id = value.class.new(local_id, value.attribute)
280
+ if ::Rails.version >= "5.2"
281
+ new_right_value =
282
+ case predicate.right
283
+ when Array
284
+ predicate.right.map {|val| transpose_predicate_value(val, current_source_shard, target_shard, type, remove).presence }.compact
285
+ else
286
+ transpose_predicate_value(predicate.right, current_source_shard, target_shard, type, remove)
287
+ end
288
+ else
289
+ new_right_value = case predicate.right
290
+ when Array
291
+ local_ids = []
292
+ predicate.right.each do |value|
293
+ local_id = Shard.relative_id_for(value, current_source_shard, target_shard)
294
+ next unless local_id
295
+ unless remove && local_id > Shard::IDS_PER_SHARD
296
+ if value.is_a?(::Arel::Nodes::Casted)
297
+ if local_id == value.val
298
+ local_id = value
299
+ elsif local_id != value
300
+ local_id = value.class.new(local_id, value.attribute)
301
+ end
305
302
  end
303
+ local_ids << local_id
306
304
  end
307
- local_ids << local_id
308
305
  end
309
- end
310
- local_ids
311
- when ::Arel::Nodes::BindParam
312
- # look for a bind param with a matching column name
313
- if ::Rails.version >= "5"
314
- binds ||= where_clause.binds + having_clause.binds
315
- if binds && bind = binds.detect{|b| b.try(:name).to_s == predicate.left.name.to_s}
306
+ local_ids
307
+ when ::Arel::Nodes::BindParam
308
+ # look for a bind param with a matching column name
309
+ if binds && bind = binds.detect{|b| b&.name.to_s == predicate.left.name.to_s}
310
+ # before we mutate, dup
311
+ if dup_binds_on_mutation
312
+ binds = binds.map(&:dup)
313
+ dup_binds_on_mutation = false
314
+ bind = binds.find { |b| b&.name.to_s == predicate.left.name.to_s }
315
+ end
316
316
  if bind.value.is_a?(::ActiveRecord::StatementCache::Substitute)
317
317
  bind.value.sharded = true # mark for transposition later
318
318
  bind.value.primary = true if type == :primary
@@ -323,29 +323,16 @@ module Switchman
323
323
  bind.instance_variable_set(:@value_for_database, nil)
324
324
  end
325
325
  end
326
+ predicate.right
326
327
  else
327
- if bind_values && idx = bind_values.find_index{|b| b.is_a?(Array) && b.first.try(:name).to_s == predicate.left.name.to_s}
328
- column, value = bind_values[idx]
329
- if ::Rails.version >= '4.2' && value.is_a?(::ActiveRecord::StatementCache::Substitute)
330
- value.sharded = true # mark for transposition later
331
- value.primary = true if type == :primary
332
- else
333
- local_id = Shard.relative_id_for(value, current_source_shard, target_shard)
334
- local_id = [] if remove && local_id > Shard::IDS_PER_SHARD
335
- bind_values[idx] = [column, local_id]
336
- end
337
- end
328
+ local_id = Shard.relative_id_for(predicate.right, current_source_shard, target_shard) || predicate.right
329
+ local_id = [] if remove && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
330
+ local_id
338
331
  end
339
- predicate.right
340
- else
341
- local_id = Shard.relative_id_for(predicate.right, current_source_shard, target_shard) || predicate.right
342
- local_id = [] if remove && local_id.is_a?(Fixnum) && local_id > Shard::IDS_PER_SHARD
343
- local_id
344
332
  end
345
-
346
333
  if new_right_value == predicate.right
347
334
  predicate
348
- elsif ::Rails.version >= "4.2" && predicate.right.is_a?(::Arel::Nodes::Casted)
335
+ elsif predicate.right.is_a?(::Arel::Nodes::Casted)
349
336
  if new_right_value == predicate.right.val
350
337
  predicate
351
338
  else
@@ -355,6 +342,33 @@ module Switchman
355
342
  predicate.class.new(predicate.left, new_right_value)
356
343
  end
357
344
  end
345
+ result = [result, binds]
346
+ result
347
+ end
348
+
349
+ def transpose_predicate_value(value, current_shard, target_shard, attribute_type, remove_non_local_ids)
350
+ if value.is_a?(::Arel::Nodes::BindParam)
351
+ query_att = value.value
352
+ current_id = query_att.value_before_type_cast
353
+ if current_id.is_a?(::ActiveRecord::StatementCache::Substitute)
354
+ current_id.sharded = true # mark for transposition later
355
+ current_id.primary = true if attribute_type == :primary
356
+ value
357
+ else
358
+ local_id = Shard.relative_id_for(current_id, current_shard, target_shard) || current_id
359
+ return nil if remove_non_local_ids && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
360
+ if current_id != local_id
361
+ # make a new bind param
362
+ ::Arel::Nodes::BindParam.new(query_att.class.new(query_att.name, local_id, query_att.type))
363
+ else
364
+ value
365
+ end
366
+ end
367
+ else
368
+ local_id = Shard.relative_id_for(value, current_shard, target_shard) || value
369
+ local_id = [] if remove_non_local_ids && local_id.is_a?(Integer) && local_id > Shard::IDS_PER_SHARD
370
+ local_id
371
+ end
358
372
  end
359
373
  end
360
374
  end
@@ -1,31 +1,59 @@
1
+ # frozen_string_literal: true
2
+
1
3
  module Switchman
2
4
  module ActiveRecord
3
5
  module Reflection
4
- module AssociationReflection
5
- # removes memoization - ActiveRecord::ModelSchema does that anyway;
6
- # and in fact this is the exact change AR makes in 4.2+
7
- if ::Rails.version < '4.2'
8
- def quoted_table_name
9
- klass.quoted_table_name
10
- end
11
- else
12
- def join_id_for(owner)
13
- owner.send(active_record_primary_key) # use sharded id values in association binds
6
+ module AbstractReflection
7
+ def shard(owner)
8
+ if polymorphic? || klass.shard_category == owner.class.shard_category
9
+ # polymorphic associations assume the same shard as the owning item
10
+ owner.shard
11
+ else
12
+ Shard.default
14
13
  end
15
14
  end
15
+ end
16
+
17
+ module AssociationScopeCache
18
+ def initialize(*args)
19
+ super
20
+ # on ThroughReflection, these won't be initialized (cause it doesn't
21
+ # inherit from AssociationReflection), so make sure they're
22
+ # initialized here
23
+ @association_scope_cache ||= {}
24
+ @scope_lock ||= Mutex.new
25
+ end
16
26
 
17
27
  # cache association scopes by shard.
18
- if ::Rails.version >= '4.2'
19
- def association_scope_cache(conn, owner)
28
+ # this technically belongs on AssociationReflection, but we put it on
29
+ # ThroughReflection as well, instead of delegating to its internal
30
+ # HasManyAssociation, losing its proper `klass`
31
+ if ::Rails.version < '6.0.4'
32
+ def association_scope_cache(conn, owner, &block)
20
33
  key = conn.prepared_statements
21
34
  if polymorphic?
22
35
  key = [key, owner._read_attribute(@foreign_type)]
23
36
  end
24
- key = [key, owner.shard.id].flatten
37
+ key = [key, shard(owner).id].flatten
25
38
  @association_scope_cache[key] ||= @scope_lock.synchronize {
26
- @association_scope_cache[key] ||= yield
39
+ @association_scope_cache[key] ||= (::Rails.version >= "5.2" ? ::ActiveRecord::StatementCache.create(conn, &block) : block.call)
27
40
  }
28
41
  end
42
+ else
43
+ def association_scope_cache(klass, owner, &block)
44
+ key = self
45
+ if polymorphic?
46
+ key = [key, owner._read_attribute(@foreign_type)]
47
+ end
48
+ key = [key, shard(owner).id].flatten
49
+ klass.cached_find_by_statement(key, &block)
50
+ end
51
+ end
52
+ end
53
+
54
+ module AssociationReflection
55
+ def join_id_for(owner)
56
+ owner.send(::Rails.version >= "5.2" ? join_foreign_key : active_record_primary_key) # use sharded id values in association binds
29
57
  end
30
58
  end
31
59
  end
@@ -1,3 +1,5 @@
1
+ # frozen_string_literal: true
2
+
1
3
  module Switchman
2
4
  module ActiveRecord
3
5
  module Relation
@@ -5,19 +7,19 @@ module Switchman
5
7
  klass::SINGLE_VALUE_METHODS.concat [ :shard, :shard_source ]
6
8
  end
7
9
 
8
- def initialize(*args)
10
+ def initialize(*, **)
9
11
  super
10
- self.shard_value = Shard.current(klass ? klass.shard_category : :default) unless shard_value
12
+ self.shard_value = Shard.current(klass ? klass.shard_category : :primary) unless shard_value
11
13
  self.shard_source_value = :implicit unless shard_source_value
12
14
  end
13
15
 
14
16
  def clone
15
17
  result = super
16
- result.shard_value = Shard.current(klass ? klass.shard_category : :default) unless shard_value
18
+ result.shard_value = Shard.current(klass ? klass.shard_category : :primary) unless shard_value
17
19
  result
18
20
  end
19
21
 
20
- def merge(*args)
22
+ def merge(*)
21
23
  relation = super
22
24
  if relation.shard_value != self.shard_value && relation.shard_source_value == :implicit
23
25
  relation.shard_value = self.shard_value
@@ -26,15 +28,15 @@ module Switchman
26
28
  relation
27
29
  end
28
30
 
29
- def new(*args, &block)
31
+ def new(*, &block)
30
32
  primary_shard.activate(klass.shard_category) { super }
31
33
  end
32
34
 
33
- def create(*args, &block)
35
+ def create(*, &block)
34
36
  primary_shard.activate(klass.shard_category) { super }
35
37
  end
36
38
 
37
- def create!(*args, &block)
39
+ def create!(*, &block)
38
40
  primary_shard.activate(klass.shard_category) { super }
39
41
  end
40
42
 
@@ -42,41 +44,60 @@ module Switchman
42
44
  primary_shard.activate(klass.shard_category) { super }
43
45
  end
44
46
 
45
- def explain(super_method: false)
46
- return super() if super_method
47
- self.activate { |relation| relation.explain(super_method: true) }
47
+ def explain
48
+ self.activate { |relation| relation.call_super(:explain, Relation) }
48
49
  end
49
50
 
50
- to_a_method = ::Rails.version > '5.0.0.beta2' ? :records : :to_a
51
- class_eval <<-RUBY, __FILE__, __LINE__ + 1
52
- def #{to_a_method}(super_method: false)
53
- return super() if super_method
54
- return @records if loaded?
55
- results = self.activate { |relation| relation.#{to_a_method}(super_method: true) }
56
- case shard_value
57
- when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
58
- @records = results
59
- @loaded = true
60
- end
61
- results
51
+ def records
52
+ return @records if loaded?
53
+ results = self.activate { |relation| relation.call_super(:records, Relation) }
54
+ case shard_value
55
+ when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
56
+ @records = results
57
+ @loaded = true
62
58
  end
63
- RUBY
64
-
65
- CALL_SUPER = Object.new.freeze
66
- private_constant :CALL_SUPER
59
+ results
60
+ end
67
61
 
68
- %w{update_all delete_all}.each do |method|
62
+ %I{update_all delete_all}.each do |method|
69
63
  class_eval <<-RUBY, __FILE__, __LINE__ + 1
70
64
  def #{method}(*args)
71
- return super(*args[1..-1]) if args.first.equal?(CALL_SUPER)
72
- result = self.activate { |relation| relation.#{method}(CALL_SUPER, *args) }
65
+ result = self.activate(unordered: true) { |relation| relation.call_super(#{method.inspect}, Relation, *args) }
73
66
  result = result.sum if result.is_a?(Array)
74
67
  result
75
68
  end
76
69
  RUBY
77
70
  end
78
71
 
79
- def activate(&block)
72
+ def find_ids_in_ranges(options = {})
73
+ is_integer = columns_hash[primary_key.to_s].type == :integer
74
+ loose_mode = options[:loose] && is_integer
75
+ # loose_mode: if we don't care about getting exactly batch_size ids in between
76
+ # don't get the max - just get the min and add batch_size so we get that many _at most_
77
+ values = loose_mode ? "MIN(id)" : "MIN(id), MAX(id)"
78
+
79
+ batch_size = options[:batch_size].try(:to_i) || 1000
80
+ quoted_primary_key = "#{klass.connection.quote_local_table_name(table_name)}.#{klass.connection.quote_column_name(primary_key)}"
81
+ as_id = " AS id" unless primary_key == 'id'
82
+ subquery_scope = except(:select).select("#{quoted_primary_key}#{as_id}").reorder(primary_key.to_sym).limit(loose_mode ? 1 : batch_size)
83
+ subquery_scope = subquery_scope.where("#{quoted_primary_key} <= ?", options[:end_at]) if options[:end_at]
84
+
85
+ first_subquery_scope = options[:start_at] ? subquery_scope.where("#{quoted_primary_key} >= ?", options[:start_at]) : subquery_scope
86
+
87
+ ids = connection.select_rows("SELECT #{values} FROM (#{first_subquery_scope.to_sql}) AS subquery").first
88
+
89
+ while ids.first.present?
90
+ ids.map!(&:to_i) if is_integer
91
+ ids << ids.first + batch_size if loose_mode
92
+
93
+ yield(*ids)
94
+ last_value = ids.last
95
+ next_subquery_scope = subquery_scope.where(["#{quoted_primary_key}>?", last_value])
96
+ ids = connection.select_rows("SELECT #{values} FROM (#{next_subquery_scope.to_sql}) AS subquery").first
97
+ end
98
+ end
99
+
100
+ def activate(unordered: false, &block)
80
101
  shards = all_shards
81
102
  if (Array === shards && shards.length == 1)
82
103
  if shards.first == DefaultShard || shards.first == Shard.current(klass.shard_category)
@@ -85,10 +106,64 @@ module Switchman
85
106
  shards.first.activate(klass.shard_category) { yield(self, shards.first) }
86
107
  end
87
108
  else
88
- # TODO: implement local limit to avoid querying extra shards
89
- Shard.with_each_shard(shards, [klass.shard_category]) do
90
- shard(Shard.current(klass.shard_category), :to_a).activate(&block)
109
+ result_count = 0
110
+ can_order = false
111
+ result = Shard.with_each_shard(shards, [klass.shard_category]) do
112
+ # don't even query other shards if we're already past the limit
113
+ next if limit_value && result_count >= limit_value && order_values.empty?
114
+
115
+ relation = shard(Shard.current(klass.shard_category), :to_a)
116
+ # do a minimal query if possible
117
+ relation = relation.limit(limit_value - result_count) if limit_value && !result_count.zero? && order_values.empty?
118
+
119
+ shard_results = relation.activate(&block)
120
+
121
+ if shard_results.present? && !unordered
122
+ can_order ||= can_order_cross_shard_results? unless order_values.empty?
123
+ raise OrderOnMultiShardQuery if !can_order && !order_values.empty? && result_count.positive?
124
+
125
+ result_count += shard_results.is_a?(Array) ? shard_results.length : 1
126
+ end
127
+ shard_results
128
+ end
129
+
130
+ result = reorder_cross_shard_results(result) if can_order
131
+ result.slice!(limit_value..-1) if limit_value
132
+ result
133
+ end
134
+ end
135
+
136
+ def can_order_cross_shard_results?
137
+ # we only presume to be able to post-sort the most basic of orderings
138
+ order_values.all? { |ov| ov.is_a?(::Arel::Nodes::Ordering) && ov.expr.is_a?(::Arel::Attributes::Attribute) }
139
+ end
140
+
141
+ def reorder_cross_shard_results(results)
142
+ results.sort! do |l, r|
143
+ result = 0
144
+ order_values.each do |ov|
145
+ if !l.is_a?(::ActiveRecord::Base)
146
+ a, b = l, r
147
+ elsif l.respond_to?(ov.expr.name)
148
+ a = l.send(ov.expr.name)
149
+ b = r.send(ov.expr.name)
150
+ else
151
+ a = l.attributes[ov.expr.name]
152
+ b = r.attributes[ov.expr.name]
153
+ end
154
+ next if a == b
155
+
156
+ if a.nil? || b.nil?
157
+ result = 1 if a.nil?
158
+ result *= -1 if ov.is_a?(::Arel::Nodes::Descending)
159
+ else
160
+ result = a <=> b
161
+ end
162
+
163
+ result *= -1 if ov.is_a?(::Arel::Nodes::Descending)
164
+ break unless result.zero?
91
165
  end
166
+ result
92
167
  end
93
168
  end
94
169
  end
@@ -1,3 +1,5 @@
1
+ # frozen_string_literal: true
2
+
1
3
  module Switchman
2
4
  module ActiveRecord
3
5
  module SpawnMethods