switchman 3.0.0 → 3.0.7

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.
checksums.yaml CHANGED
@@ -1,7 +1,7 @@
1
1
  ---
2
2
  SHA256:
3
- metadata.gz: bb89193fead8142e30280a6ab2706f6e9af798ff18820fa144be3dd9ac77591f
4
- data.tar.gz: e1157f5656f06ae9bdb7c62a2c9ae1611e2073cf47ca4278bae8605296603d21
3
+ metadata.gz: 326fba116e3a3801cdd40b47ede332a372dab86522cbeaab1a3b6b361eccce1f
4
+ data.tar.gz: 190457646d607aa444dbff867c67c5cd854817e8105efab44dd00fc782d1f21e
5
5
  SHA512:
6
- metadata.gz: fbb9b4d9385dad87ba827400dc9245085c31e133d2445daf64b3d3e81f3102c189f8feafb328d085649146c8ccd2782fd7ba3cb410c262d06a31e014175cac99
7
- data.tar.gz: 55c8010cbc73d112262d674d131b1fe0acbd39b5b482e73b9b7165aac23a25a93d5bbca4320a56b97450f68af8526045aede39c13cfb3426189073acba8f6a3f
6
+ metadata.gz: 3843b4407a36b8f7f3aff2381dc6ee2e5f886169349794e679759a6d13d96b52e96e832caf78728f20df7d1d017701d7e7d8847c8a4620cd6b669ba8f25e9cc6
7
+ data.tar.gz: 8e667cd250473f5b885fea29888454a186bafa9f694138bf0cf9f6761962ef4bf2a85b595a9cd7d26fb3f6532b4b8b8be9d6362c01329a63f94272f63ee55b4b
@@ -39,12 +39,17 @@ module Switchman
39
39
  # the first time we need a dummy dummy for re-entrancy to avoid looping on ourselves
40
40
  @default ||= default
41
41
 
42
- # Now find the actual record, if it exists; rescue the fake default if the table doesn't exist
42
+ # Now find the actual record, if it exists
43
43
  @default = begin
44
44
  find_cached('default_shard') { Shard.where(default: true).take } || default
45
+ # If we are *super* early in boot, the connection pool won't exist; we don't want to fill in the default shard yet
46
+ rescue ::ActiveRecord::ConnectionNotEstablished
47
+ nil
48
+ # rescue the fake default if the table doesn't exist
45
49
  rescue
46
50
  default
47
51
  end
52
+ return default unless @default
48
53
 
49
54
  # make sure this is not erroneously cached
50
55
  @default.database_server.remove_instance_variable(:@primary_shard) if @default.database_server.instance_variable_defined?(:@primary_shard)
@@ -59,34 +64,38 @@ module Switchman
59
64
 
60
65
  def current(klass = ::ActiveRecord::Base)
61
66
  klass ||= ::ActiveRecord::Base
62
- klass.connection_pool.shard
67
+ klass.current_switchman_shard
63
68
  end
64
69
 
65
70
  def activate(shards)
66
71
  activated_classes = activate!(shards)
67
72
  yield
68
73
  ensure
69
- activated_classes.each do |klass|
70
- klass.connection_pool.shard_stack.pop
74
+ activated_classes&.each do |klass|
71
75
  klass.connected_to_stack.pop
72
76
  end
73
77
  end
74
78
 
75
79
  def activate!(shards)
76
- activated_classes = []
80
+ activated_classes = nil
77
81
  shards.each do |klass, shard|
78
82
  next if klass == UnshardedRecord
79
83
 
80
84
  next unless klass.current_shard != shard.database_server.id.to_sym ||
81
- klass.connection_pool.shard != shard
85
+ klass.current_switchman_shard != shard
82
86
 
83
- activated_classes << klass
84
- klass.connected_to_stack << { shard: shard.database_server.id.to_sym, klasses: [klass] }
85
- klass.connection_pool.shard_stack << shard
87
+ (activated_classes ||= []) << klass
88
+ klass.connected_to_stack << { shard: shard.database_server.id.to_sym, klasses: [klass], switchman_shard: shard }
86
89
  end
87
90
  activated_classes
88
91
  end
89
92
 
93
+ def active_shards
94
+ sharded_models.map do |klass|
95
+ [klass, current(klass)]
96
+ end.compact.to_h
97
+ end
98
+
90
99
  def lookup(id)
91
100
  id_i = id.to_i
92
101
  return current if id_i == current.id || id == 'self'
@@ -103,6 +112,11 @@ module Switchman
103
112
  cached_shards[id]
104
113
  end
105
114
 
115
+ def preload_cache
116
+ cached_shards.reverse_merge!(active_shards.values.index_by(&:id))
117
+ cached_shards.reverse_merge!(all.index_by(&:id))
118
+ end
119
+
106
120
  def clear_cache
107
121
  cached_shards.clear
108
122
  end
@@ -111,14 +125,13 @@ module Switchman
111
125
  #
112
126
  # * +shards+ - an array or relation of Shards to iterate over
113
127
  # * +classes+ - an array of classes to activate
114
- # parallel: - true/false to execute in parallel, or a integer of how many
115
- # sub-processes per database server. Note that parallel
116
- # invocation currently uses forking, so should be used sparingly
117
- # because errors are not raised, and you cannot get results back
118
- # max_procs: - only run this many parallel processes at a time
128
+ # parallel: - true/false to execute in parallel, or an integer of how many
129
+ # sub-processes. Note that parallel invocation currently uses
130
+ # forking, so should be used sparingly because you cannot get
131
+ # results back
119
132
  # exception: - :ignore, :raise, :defer (wait until the end and raise the first
120
133
  # error), or a proc
121
- def with_each_shard(*args, parallel: false, max_procs: nil, exception: :raise, &block)
134
+ def with_each_shard(*args, parallel: false, exception: :raise, &block)
122
135
  raise ArgumentError, "wrong number of arguments (#{args.length} for 0...2)" if args.length > 2
123
136
 
124
137
  return Array.wrap(yield) unless default.is_a?(Shard)
@@ -133,14 +146,13 @@ module Switchman
133
146
  scope, classes = args
134
147
  end
135
148
 
136
- parallel = 1 if parallel == true
149
+ parallel = [Environment.cpu_count || 2, 2].min if parallel == true
137
150
  parallel = 0 if parallel == false || parallel.nil?
138
151
 
139
152
  scope ||= Shard.all
140
153
  scope = scope.order(::Arel.sql('database_server_id IS NOT NULL, database_server_id, id')) if ::ActiveRecord::Relation === scope && scope.order_values.empty?
141
154
 
142
- if parallel.positive?
143
- max_procs = determine_max_procs(max_procs, parallel)
155
+ if parallel > 1
144
156
  if ::ActiveRecord::Relation === scope
145
157
  # still need a post-uniq, cause the default database server could be NULL or Rails.env in the db
146
158
  database_servers = scope.reorder('database_server_id').select(:database_server_id).distinct.
@@ -148,39 +160,11 @@ module Switchman
148
160
  # nothing to do
149
161
  return if database_servers.count.zero?
150
162
 
151
- parallel = [(max_procs.to_f / database_servers.count).ceil, parallel].min if max_procs
152
-
153
- scopes = database_servers.map do |server|
154
- server_scope = server.shards.merge(scope)
155
- if parallel == 1
156
- subscopes = [server_scope]
157
- else
158
- subscopes = []
159
- total = server_scope.count
160
- ranges = []
161
- server_scope.find_ids_in_ranges(batch_size: (total.to_f / parallel).ceil) do |min, max|
162
- ranges << [min, max]
163
- end
164
- # create a half-open range on the last one
165
- ranges.last[1] = nil
166
- ranges.each do |min, max|
167
- subscope = server_scope.where('id>=?', min)
168
- subscope = subscope.where('id<=?', max) if max
169
- subscopes << subscope
170
- end
171
- end
172
- [server, subscopes]
173
- end.to_h
163
+ scopes = database_servers.to_h do |server|
164
+ [server, server.shards.merge(scope)]
165
+ end
174
166
  else
175
167
  scopes = scope.group_by(&:database_server)
176
- if parallel > 1
177
- parallel = [(max_procs.to_f / scopes.count).ceil, parallel].min if max_procs
178
- scopes = scopes.map do |(server, shards)|
179
- [server, shards.in_groups(parallel, false).compact]
180
- end.to_h
181
- else
182
- scopes = scopes.map { |(server, shards)| [server, [shards]] }.to_h
183
- end
184
168
  end
185
169
 
186
170
  exception_pipes = []
@@ -206,80 +190,83 @@ module Switchman
206
190
  end
207
191
 
208
192
  # only one process; don't bother forking
209
- if scopes.length == 1 && parallel == 1
210
- return with_each_shard(scopes.first.last.first, classes, exception: exception,
211
- &block)
212
- end
193
+ return with_each_shard(scopes.first.last, classes, exception: exception, &block) if scopes.length == 1
213
194
 
214
195
  # clear connections prior to forking (no more queries will be executed in the parent,
215
196
  # and we want them gone so that we don't accidentally use them post-fork doing something
216
197
  # silly like dealloc'ing prepared statements)
217
198
  ::ActiveRecord::Base.clear_all_connections!
218
199
 
219
- scopes.each do |server, subscopes|
220
- subscopes.each_with_index do |subscope, idx|
221
- name = if subscopes.length > 1
222
- "#{server.id} #{idx + 1}"
223
- else
224
- server.id
225
- end
226
-
227
- exception_pipe = IO.pipe
228
- exception_pipes << exception_pipe
229
- pid, io_in, io_out, io_err = Open4.pfork4(lambda do
230
- Switchman.config[:on_fork_proc]&.call
231
-
232
- # set a pretty name for the process title, up to 128 characters
233
- # (we don't actually know the limit, depending on how the process
234
- # was started)
235
- # first, simplify the binary name by stripping directories,
236
- # then truncate arguments as necessary
237
- bin = File.basename($0) # Process.argv0 doesn't work on Ruby 2.5 (https://bugs.ruby-lang.org/issues/15887)
238
- max_length = 128 - bin.length - name.length - 3
239
- args = ARGV.join(' ')
240
- args = args[0..max_length] if max_length >= 0
241
- new_title = [bin, args, name].join(' ')
242
- Process.setproctitle(new_title)
243
-
244
- with_each_shard(subscope, classes, exception: exception, &block)
245
- exception_pipe.last.close
246
- rescue => e
247
- begin
248
- dumped = Marshal.dump(e)
249
- rescue
250
- # couldn't dump the exception; create a copy with just
251
- # the message and the backtrace
252
- e2 = e.class.new(e.message)
253
- e2.set_backtrace(e.backtrace)
254
- e2.instance_variable_set(:@active_shards, e.instance_variable_get(:@active_shards))
255
- dumped = Marshal.dump(e2)
256
- end
257
- exception_pipe.last.set_encoding(dumped.encoding)
258
- exception_pipe.last.write(dumped)
259
- exception_pipe.last.flush
260
- exception_pipe.last.close
261
- exit! 1
262
- end)
200
+ scopes.each do |server, subscope|
201
+ name = server.id
202
+
203
+ exception_pipe = IO.pipe
204
+ exception_pipes << exception_pipe
205
+ pid, io_in, io_out, io_err = Open4.pfork4(lambda do
206
+ Switchman.config[:on_fork_proc]&.call
207
+
208
+ # set a pretty name for the process title, up to 128 characters
209
+ # (we don't actually know the limit, depending on how the process
210
+ # was started)
211
+ # first, simplify the binary name by stripping directories,
212
+ # then truncate arguments as necessary
213
+ bin = File.basename($0) # Process.argv0 doesn't work on Ruby 2.5 (https://bugs.ruby-lang.org/issues/15887)
214
+ max_length = 128 - bin.length - name.length - 3
215
+ args = ARGV.join(' ')
216
+ args = args[0..max_length] if max_length >= 0
217
+ new_title = [bin, args, name].join(' ')
218
+ Process.setproctitle(new_title)
219
+
220
+ with_each_shard(subscope, classes, exception: exception, &block)
263
221
  exception_pipe.last.close
264
- pids << pid
265
- io_in.close # don't care about writing to stdin
266
- out_fds << io_out
267
- err_fds << io_err
268
- pid_to_name_map[pid] = name
269
- fd_to_name_map[io_out] = name
270
- fd_to_name_map[io_err] = name
271
-
272
- while max_procs && pids.count >= max_procs
273
- while max_procs && out_fds.count >= max_procs
274
- # wait for output if we've hit the max_procs limit
275
- wait_for_output.call
276
- end
277
- # we've gotten all the output from one fd so wait for its child process to exit
278
- found_pid, status = Process.wait2
279
- pids.delete(found_pid)
280
- errors << pid_to_name_map[found_pid] if status.exitstatus != 0
222
+ rescue Exception => e # rubocop:disable Lint/RescueException
223
+ begin
224
+ dumped = Marshal.dump(e)
225
+ dumped = nil if dumped.length > 64 * 1024
226
+ rescue
227
+ dumped = nil
228
+ end
229
+
230
+ if dumped.nil?
231
+ # couldn't dump the exception; create a copy with just
232
+ # the message and the backtrace
233
+ e2 = e.class.new(e.message)
234
+ backtrace = e.backtrace
235
+ # truncate excessively long backtraces
236
+ backtrace = backtrace[0...25] + ['...'] + backtrace[-25..] if backtrace.length > 50
237
+ e2.set_backtrace(backtrace)
238
+ e2.instance_variable_set(:@active_shards, e.instance_variable_get(:@active_shards))
239
+ dumped = Marshal.dump(e2)
281
240
  end
241
+ exception_pipe.last.set_encoding(dumped.encoding)
242
+ exception_pipe.last.write(dumped)
243
+ exception_pipe.last.flush
244
+ exception_pipe.last.close
245
+ exit! 1
246
+ end)
247
+ exception_pipe.last.close
248
+ pids << pid
249
+ io_in.close # don't care about writing to stdin
250
+ out_fds << io_out
251
+ err_fds << io_err
252
+ pid_to_name_map[pid] = name
253
+ fd_to_name_map[io_out] = name
254
+ fd_to_name_map[io_err] = name
255
+
256
+ while pids.count >= parallel
257
+ while out_fds.count >= parallel
258
+ # wait for output if we've hit the parallel limit
259
+ wait_for_output.call
260
+ end
261
+ # we've gotten all the output from one fd so wait for its child process to exit
262
+ found_pid, status = Process.wait2
263
+ pids.delete(found_pid)
264
+ errors << pid_to_name_map[found_pid] if status.exitstatus != 0
282
265
  end
266
+ # we've gotten all the output from one fd so wait for its child process to exit
267
+ found_pid, status = Process.wait2
268
+ pids.delete(found_pid)
269
+ errors << pid_to_name_map[found_pid] if status.exitstatus != 0
283
270
  end
284
271
 
285
272
  wait_for_output.call while out_fds.any? || err_fds.any?
@@ -402,7 +389,7 @@ module Switchman
402
389
  signed_id_operation(local_id) do |id|
403
390
  return nil if id > IDS_PER_SHARD
404
391
 
405
- $1.to_i * IDS_PER_SHARD + id
392
+ ($1.to_i * IDS_PER_SHARD) + id
406
393
  end
407
394
  when Integer, /^-?\d+$/
408
395
  any_id.to_i
@@ -480,23 +467,6 @@ module Switchman
480
467
  shard || source_shard || Shard.current
481
468
  end
482
469
 
483
- # given the provided option, determines whether we need to (and whether
484
- # it's possible) to determine a reasonable default.
485
- def determine_max_procs(max_procs_input, parallel_input = 2)
486
- max_procs = nil
487
- if max_procs_input
488
- max_procs = max_procs_input.to_i
489
- max_procs = nil if max_procs.zero?
490
- else
491
- return 1 if parallel_input.nil? || parallel_input < 1
492
-
493
- cpus = Environment.cpu_count
494
- max_procs = cpus * parallel_input if cpus&.positive?
495
- end
496
-
497
- max_procs
498
- end
499
-
500
470
  private
501
471
 
502
472
  def add_sharded_model(klass)
@@ -505,7 +475,7 @@ module Switchman
505
475
  end
506
476
 
507
477
  def initialize_sharding
508
- full_connects_to_hash = DatabaseServer.all.map { |db| [db.id.to_sym, db.connects_to_hash] }.to_h
478
+ full_connects_to_hash = DatabaseServer.all.to_h { |db| [db.id.to_sym, db.connects_to_hash] }
509
479
  sharded_models.each do |klass|
510
480
  connects_to_hash = full_connects_to_hash.deep_dup
511
481
  if klass == UnshardedRecord
@@ -596,6 +566,10 @@ module Switchman
596
566
  Shard.default
597
567
  end
598
568
 
569
+ def original_id
570
+ id
571
+ end
572
+
599
573
  def activate(*classes, &block)
600
574
  shards = hashify_classes(classes)
601
575
  Shard.activate(shards, &block)
@@ -668,7 +642,7 @@ module Switchman
668
642
  return nil unless local_id
669
643
 
670
644
  self.class.signed_id_operation(local_id) do |abs_id|
671
- abs_id + id * IDS_PER_SHARD
645
+ abs_id + (id * IDS_PER_SHARD)
672
646
  end
673
647
  end
674
648
 
@@ -690,6 +664,7 @@ module Switchman
690
664
  Switchman.cache.delete(['shard', id].join('/'))
691
665
  Switchman.cache.delete('default_shard') if default?
692
666
  end
667
+ self.class.clear_cache
693
668
  end
694
669
 
695
670
  def default_name
@@ -2,7 +2,7 @@
2
2
 
3
3
  class AddDefaultShardIndex < ActiveRecord::Migration[4.2]
4
4
  def change
5
- Switchman::Shard.where(default: nil).update_all(default: false)
5
+ Switchman::Shard.where(default: nil).update_all(default: false) if Switchman::Shard.current.default?
6
6
  change_column_default :switchman_shards, :default, false
7
7
  change_column_null :switchman_shards, :default, false
8
8
  options = if connection.adapter_name == 'PostgreSQL'
@@ -10,6 +10,6 @@ class AddDefaultShardIndex < ActiveRecord::Migration[4.2]
10
10
  else
11
11
  {}
12
12
  end
13
- add_index :switchman_shards, :default, options
13
+ add_index :switchman_shards, :default, **options
14
14
  end
15
15
  end
@@ -4,7 +4,7 @@ class AddTimestampsToShards < ActiveRecord::Migration[4.2]
4
4
  disable_ddl_transaction!
5
5
 
6
6
  def change
7
- add_timestamps :switchman_shards, null: true
7
+ add_timestamps :switchman_shards, null: true, if_not_exists: true
8
8
  now = Time.now.utc
9
9
  Switchman::Shard.update_all(updated_at: now, created_at: now) if Switchman::Shard.current.default?
10
10
  change_column_null :switchman_shards, :updated_at, false
@@ -40,15 +40,6 @@ module Switchman
40
40
  ensure
41
41
  @last_query_at = Time.now
42
42
  end
43
-
44
- private
45
-
46
- def id_value_for_database(value)
47
- return super unless value.class.sharded_primary_key?
48
-
49
- # do this the Rails 4.2 way, so that if Shard.current != self.shard, the id gets transposed
50
- quote(value.id)
51
- end
52
43
  end
53
44
  end
54
45
  end
@@ -85,7 +85,7 @@ module Switchman
85
85
  # Copypasta from Activerecord but with added global_id_for goodness.
86
86
  def records_for(ids)
87
87
  scope.where(association_key_name => ids).load do |record|
88
- global_key = if record.class.connection_classes == UnshardedRecord
88
+ global_key = if model.connection_classes == UnshardedRecord
89
89
  convert_key(record[association_key_name])
90
90
  else
91
91
  Shard.global_id_for(record[association_key_name], record.shard)
@@ -40,7 +40,11 @@ module Switchman
40
40
  if sharded_column?(attr_name)
41
41
  owner << <<-RUBY
42
42
  def global_#{attr_name}
43
- ::Switchman::Shard.global_id_for(original_#{attr_name}, shard)
43
+ raw_value = original_#{attr_name}
44
+ return nil if raw_value.nil?
45
+ return raw_value if raw_value > ::Switchman::Shard::IDS_PER_SHARD
46
+
47
+ ::Switchman::Shard.global_id_for(raw_value, shard)
44
48
  end
45
49
  RUBY
46
50
  else
@@ -52,7 +56,9 @@ module Switchman
52
56
  if sharded_column?(attr_name)
53
57
  owner << <<-RUBY
54
58
  def local_#{attr_name}
55
- ::Switchman::Shard.local_id_for(original_#{attr_name}).first
59
+ raw_value = original_#{attr_name}
60
+ return nil if raw_value.nil?
61
+ return raw_value % ::Switchman::Shard::IDS_PER_SHARD
56
62
  end
57
63
  RUBY
58
64
  else
@@ -104,7 +110,22 @@ module Switchman
104
110
  alias_method 'original_#{attr_name}', '#{attr_name}'
105
111
  # and replace with one that transposes the id
106
112
  def #{attr_name}
107
- ::Switchman::Shard.relative_id_for(original_#{attr_name}, shard, ::Switchman::Shard.current(#{connection_classes_code_for_reflection(reflection)}))
113
+ raw_value = original_#{attr_name}
114
+ return nil if raw_value.nil?
115
+
116
+ abs_raw_value = raw_value.abs
117
+ current_shard = ::Switchman::Shard.current(#{connection_classes_code_for_reflection(reflection)})
118
+ same_shard = shard == current_shard
119
+ return raw_value if same_shard && abs_raw_value < ::Switchman::Shard::IDS_PER_SHARD
120
+
121
+ value_shard_id = abs_raw_value / ::Switchman::Shard::IDS_PER_SHARD
122
+ # this is a stupid case when someone stuffed a global id for the current shard in instead
123
+ # of a local id
124
+ return raw_value % ::Switchman::Shard::IDS_PER_SHARD if value_shard_id == current_shard.id
125
+ return raw_value if !same_shard && abs_raw_value > ::Switchman::Shard::IDS_PER_SHARD
126
+ return shard.global_id_for(raw_value) if !same_shard && abs_raw_value < ::Switchman::Shard::IDS_PER_SHARD
127
+
128
+ ::Switchman::Shard.relative_id_for(raw_value, shard, current_shard)
108
129
  end
109
130
 
110
131
  alias_method 'original_#{attr_name}=', '#{attr_name}='
@@ -141,6 +162,54 @@ module Switchman
141
162
  self.class.define_attribute_methods
142
163
  super
143
164
  end
165
+
166
+ # these are called if the specific methods haven't been defined yet
167
+ def attribute(attr_name)
168
+ return super unless self.class.sharded_column?(attr_name)
169
+
170
+ reflection = self.class.send(:reflection_for_integer_attribute, attr_name)
171
+ ::Switchman::Shard.relative_id_for(super, shard, ::Switchman::Shard.current(connection_classes_for_reflection(reflection)))
172
+ end
173
+
174
+ def attribute=(attr_name, new_value)
175
+ unless self.class.sharded_column?(attr_name)
176
+ super
177
+ return
178
+ end
179
+
180
+ reflection = self.class.send(:reflection_for_integer_attribute, attr_name)
181
+ super(::Switchman::Shard.relative_id_for(new_value, ::Switchman::Shard.current(connection_classes_for_reflection(reflection)), shard))
182
+ end
183
+
184
+ def global_attribute(attr_name)
185
+ if self.class.sharded_column?(attr_name)
186
+ ::Switchman::Shard.global_id_for(attribute(attr_name), shard)
187
+ else
188
+ attribute(attr_name)
189
+ end
190
+ end
191
+
192
+ def local_attribute(attr_name)
193
+ if self.class.sharded_column?(attr_name)
194
+ ::Switchman::Shard.local_id_for(attribute(attr_name), shard).first
195
+ else
196
+ attribute(attr_name)
197
+ end
198
+ end
199
+
200
+ private
201
+
202
+ def connection_classes_for_reflection(reflection)
203
+ if reflection
204
+ if reflection.options[:polymorphic]
205
+ read_attribute(reflection.foreign_type)&.constantize&.connection_classes
206
+ else
207
+ reflection.klass.connection_classes
208
+ end
209
+ else
210
+ self.class.connection_classes
211
+ end
212
+ end
144
213
  end
145
214
  end
146
215
  end
@@ -77,17 +77,27 @@ module Switchman
77
77
 
78
78
  default_shard
79
79
  end
80
+
81
+ def current_switchman_shard
82
+ connected_to_stack.reverse_each do |hash|
83
+ return hash[:switchman_shard] if hash[:switchman_shard] && hash[:klasses].include?(connection_classes)
84
+ end
85
+
86
+ Shard.default
87
+ end
80
88
  end
81
89
 
82
- def self.included(klass)
90
+ def self.prepended(klass)
83
91
  klass.singleton_class.prepend(ClassMethods)
84
- klass.set_callback(:initialize, :before) do
85
- @shard ||= if self.class.sharded_primary_key?
86
- Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_classes))
87
- else
88
- Shard.current(self.class.connection_classes)
89
- end
90
- end
92
+ end
93
+
94
+ def _run_initialize_callbacks
95
+ @shard ||= if self.class.sharded_primary_key?
96
+ Shard.shard_for(self[self.class.primary_key], Shard.current(self.class.connection_classes))
97
+ else
98
+ Shard.current(self.class.connection_classes)
99
+ end
100
+ super
91
101
  end
92
102
 
93
103
  def shard
@@ -107,12 +117,12 @@ module Switchman
107
117
 
108
118
  def save(*, **)
109
119
  @shard_set_in_stone = true
110
- (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
120
+ super
111
121
  end
112
122
 
113
123
  def save!(*, **)
114
124
  @shard_set_in_stone = true
115
- (self.class.current_scope || self.class.default_scoped).shard(shard, :implicit).scoping { super }
125
+ super
116
126
  end
117
127
 
118
128
  def destroy
@@ -134,6 +144,12 @@ module Switchman
134
144
  end
135
145
  end
136
146
 
147
+ def with_transaction_returning_status
148
+ shard.activate(self.class.connection_classes) do
149
+ super
150
+ end
151
+ end
152
+
137
153
  def hash
138
154
  self.class.sharded_primary_key? ? self.class.hash ^ global_id.hash : super
139
155
  end
@@ -149,19 +165,20 @@ module Switchman
149
165
  copy
150
166
  end
151
167
 
152
- def quoted_id
153
- return super unless self.class.sharded_primary_key?
168
+ def update_columns(*)
169
+ db = shard.database_server
170
+ return db.unguard { super } if ::GuardRail.environment != db.guard_rail_environment
154
171
 
155
- # do this the Rails 4.2 way, so that if Shard.current != self.shard, the id gets transposed
156
- self.class.connection.quote(id)
172
+ super
157
173
  end
158
174
 
159
- def update_columns(*)
160
- db = Shard.current(self.class.connection_classes).database_server
161
- if ::GuardRail.environment == db.guard_rail_environment
162
- super
175
+ def id_for_database
176
+ if self.class.sharded_primary_key?
177
+ # It's an int, so so it's safe to just return it without passing it through anything else
178
+ # In theory we should do `@attributes[@primary_key].type.serialize(id)`, but that seems to have surprising side-effects
179
+ id
163
180
  else
164
- db.unguard { super }
181
+ super
165
182
  end
166
183
  end
167
184
 
@@ -172,7 +189,7 @@ module Switchman
172
189
  if reflection
173
190
  if reflection.options[:polymorphic]
174
191
  begin
175
- read_attribute(reflection.foreign_type)&.constantize&.connection_classes
192
+ read_attribute(reflection.foreign_type)&.constantize&.connection_classes || ::ActiveRecord::Base
176
193
  rescue NameError
177
194
  # in case someone is abusing foreign_type to not point to an actual class
178
195
  ::ActiveRecord::Base
@@ -182,7 +199,7 @@ module Switchman
182
199
  reflection.klass.connection_classes
183
200
  end
184
201
  else
185
- connection_classes
202
+ self.class.connection_classes
186
203
  end
187
204
  end
188
205
  end
@@ -50,7 +50,7 @@ module Switchman
50
50
 
51
51
  def calculate_simple_average(column_name, distinct)
52
52
  # See activerecord#execute_simple_calculation
53
- relation = reorder(nil)
53
+ relation = except(:order)
54
54
  column = aggregate_column(column_name)
55
55
  relation.select_values = [operation_over_aggregate_column(column, 'average', distinct).as('average'),
56
56
  operation_over_aggregate_column(column, 'count', distinct).as('count')]
@@ -83,7 +83,7 @@ module Switchman
83
83
  if opts[:association]
84
84
  key_ids = calculated_data.collect { |row| row[opts[:group_aliases].first] }
85
85
  key_records = opts[:association].klass.base_class.where(id: key_ids)
86
- key_records = key_records.map { |r| [Shard.relative_id_for(r, shard, target_shard), r] }.to_h
86
+ key_records = key_records.to_h { |r| [Shard.relative_id_for(r, shard, target_shard), r] }
87
87
  end
88
88
 
89
89
  calculated_data.map do |row|