switchman 1.14.10 → 2.0.0

Sign up to get free protection for your applications and to get access to all the features.
checksums.yaml CHANGED
@@ -1,7 +1,7 @@
1
1
  ---
2
2
  SHA256:
3
- metadata.gz: 36690ef40f267eaebf781f63dd9e236b5ac3db9a8241e139a4d38e97fc8e7102
4
- data.tar.gz: 83e6084591b9d1a24ade9c3cb0e1a9b92b2d5c2af3fdc18057c9d0922728d53d
3
+ metadata.gz: 7cec081d77daa09b754a98ce7bc335198b2dfbdbcd093df492a43803ff7a2742
4
+ data.tar.gz: d93576725471024b24c4622d33e0eff9c1b61498b6c037e624d9b6ff8c78a442
5
5
  SHA512:
6
- metadata.gz: '0784224cb3eb7bbcf009e2b39ede3c082bc0c990cda3828f9375312f41d874210a655f2cfe7e5de670e30e1c0e84b85378931c3d27eb4ef191249c72775f7973'
7
- data.tar.gz: 23cbea13d056719ab560b0721238a8088a1205174ba8259debbf6165b55601fec180608bc21abaf70409458bcc8886b60310f42b4824d30a83159bbc491dbc8d
6
+ metadata.gz: 6d4c7b1e82460fbdfda1659fc09379583c5e63ecc21d412466563a62c42679db864f8f09c2eaa2bed17b53216c64d18c42caabec76291f25e00bfc1e6a2499a1
7
+ data.tar.gz: e52baf68ce2301dfa26abffd423e9855c875232871539fea59ae88ffd8f3b355299cb4c0bb60e32bf90555e7e935d2f73954ec01919cab16993bee73b39e56b6
@@ -1,11 +1,743 @@
1
- # ActiveRecord::Base _must_ be loaded prior to the file that defines Shard
2
- # Because Switchman hooks into ActiveSupport.on_load(:active_record) and
3
- # reference it again (technically it gets referenced from inside
4
- # ConnectionHandler#establish_connection, when AR's on_load(:active_record)
5
- # hook gets called
6
- ActiveRecord::Base
7
- # in case it was referenced _before_ the on_load hook is added, then
8
- # we need to make sure we define the class someone wanted (in which case
9
- # AR's on_load hook won't be called either yet, and establish_connection
10
- # will be safe)
11
- require_dependency 'switchman/shard_internal'
1
+ require 'switchman/database_server'
2
+ require 'switchman/default_shard'
3
+ require 'switchman/environment'
4
+ require 'switchman/errors'
5
+
6
+ module Switchman
7
+ class Shard < ::ActiveRecord::Base
8
+ # ten trillion possible ids per shard. yup.
9
+ IDS_PER_SHARD = 10_000_000_000_000
10
+
11
+ CATEGORIES =
12
+ {
13
+ # special cased to mean all other models
14
+ :primary => nil,
15
+ # special cased to not allow activating a shard other than the default
16
+ :unsharded => [Shard]
17
+ }
18
+ private_constant :CATEGORIES
19
+ @connection_specification_name = @shard_category = :unsharded
20
+
21
+ if defined?(::ProtectedAttributes)
22
+ attr_accessible :default, :name, :database_server
23
+ end
24
+
25
+ # only allow one default
26
+ validates_uniqueness_of :default, :if => lambda { |s| s.default? }
27
+
28
+ after_save :clear_cache
29
+ after_destroy :clear_cache
30
+
31
+ after_rollback :on_rollback
32
+
33
+ scope :primary, -> { where(name: nil).order(:database_server_id, :id).distinct_on(:database_server_id) }
34
+
35
+ class << self
36
+ def categories
37
+ CATEGORIES.keys
38
+ end
39
+
40
+ def default(reload_deprecated = false, reload: false, with_fallback: false)
41
+ if reload_deprecated
42
+ reload = reload_deprecated
43
+ ActiveSupport::Deprecation.warn("positional reload parameter to Switchman::Shard.default is deprecated; use `reload: true`")
44
+ end
45
+ if !@default || reload
46
+ # Have to create a dummy object so that several key methods still work
47
+ # (it's easier to do this in one place here, and just assume that sharding
48
+ # is up and running everywhere else). This includes for looking up the
49
+ # default shard itself. This also needs to be a local so that this method
50
+ # can be re-entrant
51
+ default = DefaultShard.instance
52
+
53
+ # if we already have a default shard in place, and the caller wants
54
+ # to use it as a fallback, use that instead of the dummy instance
55
+ if with_fallback && @default
56
+ default = @default
57
+ end
58
+
59
+ # the first time we need a dummy dummy for re-entrancy to avoid looping on ourselves
60
+ @default ||= default
61
+
62
+ # Now find the actual record, if it exists; rescue the fake default if the table doesn't exist
63
+ @default = begin
64
+ find_cached("default_shard") { Shard.where(default: true).take } || default
65
+ rescue
66
+ default
67
+ end
68
+
69
+ # rebuild current shard activations - it might have "another" default shard serialized there
70
+ active_shards.replace(active_shards.dup.map do |category, shard|
71
+ shard = Shard.lookup((!shard || shard.default?) ? 'default' : shard.id)
72
+ [category, shard]
73
+ end.to_h)
74
+
75
+ activate!(primary: @default) if active_shards.empty?
76
+
77
+ # make sure this is not erroneously cached
78
+ if @default.database_server.instance_variable_defined?(:@primary_shard)
79
+ @default.database_server.remove_instance_variable(:@primary_shard)
80
+ end
81
+
82
+ # and finally, check for cached references to the default shard on the existing connection
83
+ if ::ActiveRecord::Base.connected? && ::ActiveRecord::Base.connection.shard.default?
84
+ ::ActiveRecord::Base.connection.shard = @default
85
+ end
86
+ end
87
+ @default
88
+ end
89
+
90
+ def current(category = :primary)
91
+ active_shards[category] || Shard.default
92
+ end
93
+
94
+ def activate(shards)
95
+ old_shards = activate!(shards)
96
+ yield
97
+ ensure
98
+ active_shards.merge!(old_shards) if old_shards
99
+ end
100
+
101
+ def activate!(shards)
102
+ old_shards = nil
103
+ currently_active_shards = active_shards
104
+ shards.each do |category, shard|
105
+ next if category == :unsharded
106
+ unless currently_active_shards[category] == shard
107
+ old_shards ||= {}
108
+ old_shards[category] = currently_active_shards[category]
109
+ currently_active_shards[category] = shard
110
+ end
111
+ end
112
+ old_shards
113
+ end
114
+
115
+ def lookup(id)
116
+ id_i = id.to_i
117
+ return current if id_i == current.id || id == 'self'
118
+ return default if id_i == default.id || id.nil? || id == 'default'
119
+ id = id_i
120
+ raise ArgumentError if id == 0
121
+
122
+ unless cached_shards.has_key?(id)
123
+ cached_shards[id] = Shard.default.activate do
124
+ find_cached(['shard', id]) { find_by(id: id) }
125
+ end
126
+ end
127
+ cached_shards[id]
128
+ end
129
+
130
+ def clear_cache
131
+ cached_shards.clear
132
+ end
133
+
134
+ # ==== Parameters
135
+ #
136
+ # * +shards+ - an array or relation of Shards to iterate over
137
+ # * +categories+ - an array of categories to activate
138
+ # * +options+ -
139
+ # :parallel - true/false to execute in parallel, or a integer of how many
140
+ # sub-processes per database server. Note that parallel
141
+ # invocation currently uses forking, so should be used sparingly
142
+ # because errors are not raised, and you cannot get results back
143
+ # :max_procs - only run this many parallel processes at a time
144
+ # :exception - :ignore, :raise, :defer (wait until the end and raise the first
145
+ # error), or a proc
146
+ def with_each_shard(*args)
147
+ raise ArgumentError, "wrong number of arguments (#{args.length} for 0...3)" if args.length > 3
148
+
149
+ unless default.is_a?(Shard)
150
+ return Array.wrap(yield)
151
+ end
152
+
153
+ options = args.extract_options!
154
+ if args.length == 1
155
+ if Array === args.first && args.first.first.is_a?(Symbol)
156
+ categories = args.first
157
+ else
158
+ scope = args.first
159
+ end
160
+ else
161
+ scope, categories = args
162
+ end
163
+
164
+ parallel = case options[:parallel]
165
+ when true
166
+ 1
167
+ when false, nil
168
+ 0
169
+ else
170
+ options[:parallel]
171
+ end
172
+ options.delete(:parallel)
173
+
174
+ scope ||= Shard.all
175
+ if ::ActiveRecord::Relation === scope && scope.order_values.empty?
176
+ scope = scope.order(::Arel.sql("database_server_id IS NOT NULL, database_server_id, id"))
177
+ end
178
+
179
+ if parallel > 0
180
+ max_procs = determine_max_procs(options.delete(:max_procs), parallel)
181
+ if ::ActiveRecord::Relation === scope
182
+ # still need a post-uniq, cause the default database server could be NULL or Rails.env in the db
183
+ database_servers = scope.reorder('database_server_id').select(:database_server_id).distinct.
184
+ map(&:database_server).compact.uniq
185
+ # nothing to do
186
+ return if database_servers.count == 0
187
+ parallel = [(max_procs.to_f / database_servers.count).ceil, parallel].min if max_procs
188
+
189
+ scopes = Hash[database_servers.map do |server|
190
+ server_scope = server.shards.merge(scope)
191
+ if parallel == 1
192
+ subscopes = [server_scope]
193
+ else
194
+ subscopes = []
195
+ total = server_scope.count
196
+ ranges = []
197
+ server_scope.find_ids_in_ranges(:batch_size => (total.to_f / parallel).ceil) do |min, max|
198
+ ranges << [min, max]
199
+ end
200
+ # create a half-open range on the last one
201
+ ranges.last[1] = nil
202
+ ranges.each do |min, max|
203
+ subscope = server_scope.where("id>=?", min)
204
+ subscope = subscope.where("id<=?", max) if max
205
+ subscopes << subscope
206
+ end
207
+ end
208
+ [server, subscopes]
209
+ end]
210
+ else
211
+ scopes = scope.group_by(&:database_server)
212
+ if parallel > 1
213
+ parallel = [(max_procs.to_f / scopes.count).ceil, parallel].min if max_procs
214
+ scopes = Hash[scopes.map do |(server, shards)|
215
+ [server, shards.in_groups(parallel, false).compact]
216
+ end]
217
+ else
218
+ scopes = Hash[scopes.map { |(server, shards)| [server, [shards]] }]
219
+ end
220
+ end
221
+
222
+ exception_pipes = []
223
+ pids = []
224
+ out_fds = []
225
+ err_fds = []
226
+ pid_to_name_map = {}
227
+ fd_to_name_map = {}
228
+ errors = []
229
+
230
+ wait_for_output = lambda do |out_fds, err_fds, fd_to_name_map|
231
+ ready, _ = IO.select(out_fds + err_fds)
232
+ ready.each do |fd|
233
+ if fd.eof?
234
+ fd.close
235
+ out_fds.delete(fd)
236
+ err_fds.delete(fd)
237
+ next
238
+ end
239
+ line = fd.readline
240
+ puts "#{fd_to_name_map[fd]}: #{line}"
241
+ end
242
+ end
243
+
244
+ # only one process; don't bother forking
245
+ if scopes.length == 1 && parallel == 1
246
+ return with_each_shard(scopes.first.last.first, categories, options) { yield }
247
+ end
248
+
249
+ # clear connections prior to forking (no more queries will be executed in the parent,
250
+ # and we want them gone so that we don't accidentally use them post-fork doing something
251
+ # silly like dealloc'ing prepared statements)
252
+ ::ActiveRecord::Base.clear_all_connections!
253
+
254
+ scopes.each do |server, subscopes|
255
+ subscopes.each_with_index do |subscope, idx|
256
+ if subscopes.length > 1
257
+ name = "#{server.id} #{idx + 1}"
258
+ else
259
+ name = server.id
260
+ end
261
+
262
+ exception_pipe = IO.pipe
263
+ exception_pipes << exception_pipe
264
+ pid, io_in, io_out, io_err = Open4.pfork4(lambda do
265
+ begin
266
+ Switchman.config[:on_fork_proc]&.call
267
+
268
+ # set a pretty name for the process title, up to 128 characters
269
+ # (we don't actually know the limit, depending on how the process
270
+ # was started)
271
+ # first, simplify the binary name by stripping directories,
272
+ # then truncate arguments as necessary
273
+ bin = File.basename($0) # Process.argv0 doesn't work on Ruby 2.5 (https://bugs.ruby-lang.org/issues/15887)
274
+ max_length = 128 - bin.length - name.length - 3
275
+ args = ARGV.join(" ")
276
+ if max_length >= 0
277
+ args = args[0..max_length]
278
+ end
279
+ new_title = [bin, args, name].join(" ")
280
+ Process.setproctitle(new_title)
281
+
282
+ with_each_shard(subscope, categories, options) { yield }
283
+ exception_pipe.last.close
284
+ rescue => e
285
+ begin
286
+ dumped = Marshal.dump(e)
287
+ rescue
288
+ # couldn't dump the exception; create a copy with just
289
+ # the message and the backtrace
290
+ e2 = e.class.new(e.message)
291
+ e2.set_backtrace(e.backtrace)
292
+ e2.instance_variable_set(:@active_shards, e.instance_variable_get(:@active_shards))
293
+ dumped = Marshal.dump(e2)
294
+ end
295
+ exception_pipe.last.set_encoding(dumped.encoding)
296
+ exception_pipe.last.write(dumped)
297
+ exception_pipe.last.flush
298
+ exception_pipe.last.close
299
+ exit! 1
300
+ end
301
+ end)
302
+ exception_pipe.last.close
303
+ pids << pid
304
+ io_in.close # don't care about writing to stdin
305
+ out_fds << io_out
306
+ err_fds << io_err
307
+ pid_to_name_map[pid] = name
308
+ fd_to_name_map[io_out] = name
309
+ fd_to_name_map[io_err] = name
310
+
311
+ while max_procs && pids.count >= max_procs
312
+ while max_procs && out_fds.count >= max_procs
313
+ # wait for output if we've hit the max_procs limit
314
+ wait_for_output.call(out_fds, err_fds, fd_to_name_map)
315
+ end
316
+ # we've gotten all the output from one fd so wait for its child process to exit
317
+ found_pid, status = Process.wait2
318
+ pids.delete(found_pid)
319
+ errors << pid_to_name_map[found_pid] if status.exitstatus != 0
320
+ end
321
+ end
322
+ end
323
+
324
+ while out_fds.any? || err_fds.any?
325
+ wait_for_output.call(out_fds, err_fds, fd_to_name_map)
326
+ end
327
+ pids.each do |pid|
328
+ _, status = Process.waitpid2(pid)
329
+ errors << pid_to_name_map[pid] if status.exitstatus != 0
330
+ end
331
+
332
+ # check for an exception; we only re-raise the first one
333
+ exception_pipes.each do |exception_pipe|
334
+ begin
335
+ serialized_exception = exception_pipe.first.read
336
+ next if serialized_exception.empty?
337
+ exception = Marshal.load(serialized_exception)
338
+ raise exception
339
+ ensure
340
+ exception_pipe.first.close
341
+ end
342
+ end
343
+
344
+ unless errors.empty?
345
+ raise ParallelShardExecError.new("The following subprocesses did not exit cleanly: #{errors.sort.join(", ")}")
346
+ end
347
+ return
348
+ end
349
+
350
+ categories ||= []
351
+
352
+ previous_shard = nil
353
+ close_connections_if_needed = lambda do |shard|
354
+ # prune the prior connection unless it happened to be the same
355
+ if previous_shard && shard != previous_shard && !previous_shard.database_server.shareable?
356
+ previous_shard.activate do
357
+ ::GuardRail.activated_environments.each do |env|
358
+ ::GuardRail.activate(env) do
359
+ if ::ActiveRecord::Base.connected? && ::ActiveRecord::Base.connection.open_transactions == 0
360
+ ::ActiveRecord::Base.connection_pool.current_pool.disconnect!
361
+ end
362
+ end
363
+ end
364
+ end
365
+ end
366
+ end
367
+
368
+ result = []
369
+ exception = nil
370
+ scope.each do |shard|
371
+ # shard references a database server that isn't configured in this environment
372
+ next unless shard.database_server
373
+ close_connections_if_needed.call(shard)
374
+ shard.activate(*categories) do
375
+ begin
376
+ result.concat Array.wrap(yield)
377
+ rescue
378
+ case options[:exception]
379
+ when :ignore
380
+ when :defer
381
+ exception ||= $!
382
+ when Proc
383
+ options[:exception].call
384
+ when :raise
385
+ raise
386
+ else
387
+ raise
388
+ end
389
+ end
390
+ end
391
+ previous_shard = shard
392
+ end
393
+ close_connections_if_needed.call(Shard.current)
394
+ raise exception if exception
395
+ result
396
+ end
397
+
398
+ def partition_by_shard(array, partition_proc = nil)
399
+ shard_arrays = {}
400
+ array.each do |object|
401
+ partition_object = partition_proc ? partition_proc.call(object) : object
402
+ case partition_object
403
+ when Shard
404
+ shard = partition_object
405
+ when ::ActiveRecord::Base
406
+ if partition_object.respond_to?(:associated_shards)
407
+ partition_object.associated_shards.each do |a_shard|
408
+ shard_arrays[a_shard] ||= []
409
+ shard_arrays[a_shard] << object
410
+ end
411
+ next
412
+ else
413
+ shard = partition_object.shard
414
+ end
415
+ when Integer, /^\d+$/, /^(\d+)~(\d+)$/
416
+ local_id, shard = Shard.local_id_for(partition_object)
417
+ local_id ||= partition_object
418
+ object = local_id if !partition_proc
419
+ end
420
+ shard ||= Shard.current
421
+ shard_arrays[shard] ||= []
422
+ shard_arrays[shard] << object
423
+ end
424
+ # TODO: use with_each_shard (or vice versa) to get
425
+ # connection management and parallelism benefits
426
+ shard_arrays.inject([]) do |results, (shard, objects)|
427
+ results.concat shard.activate { Array.wrap(yield objects) }
428
+ end
429
+ end
430
+
431
+ # it's tedious to hold onto this same
432
+ # kind of sign state and transform the
433
+ # result in multiple places, so
434
+ # here we can operate on the absolute value
435
+ # in a provided block and trust the sign will
436
+ # stay as provided. This assumes no consumer
437
+ # will return a nil value from the block.
438
+ def signed_id_operation(input_id)
439
+ sign = input_id < 0 ? -1 : 1
440
+ output = yield input_id.abs
441
+ output * sign
442
+ end
443
+
444
+ # converts an AR object, integral id, string id, or string short-global-id to a
445
+ # integral id. nil if it can't be interpreted
446
+ def integral_id_for(any_id)
447
+ if any_id.is_a?(::Arel::Nodes::Casted)
448
+ any_id = any_id.val
449
+ elsif any_id.is_a?(::Arel::Nodes::BindParam) && ::Rails.version >= "5.2"
450
+ any_id = any_id.value.value_before_type_cast
451
+ end
452
+
453
+ case any_id
454
+ when ::ActiveRecord::Base
455
+ any_id.id
456
+ when /^(\d+)~(-?\d+)$/
457
+ local_id = $2.to_i
458
+ signed_id_operation(local_id) do |id|
459
+ return nil if id > IDS_PER_SHARD
460
+ $1.to_i * IDS_PER_SHARD + id
461
+ end
462
+ when Integer, /^-?\d+$/
463
+ any_id.to_i
464
+ else
465
+ nil
466
+ end
467
+ end
468
+
469
+ # takes an id-ish, and returns a local id and the shard it's
470
+ # local to. [nil, nil] if it can't be interpreted. [id, nil]
471
+ # if it's already a local ID. [nil, nil] if it's a well formed
472
+ # id, but the shard it refers to does not exist
473
+ NIL_NIL_ID = [nil, nil].freeze
474
+ def local_id_for(any_id)
475
+ id = integral_id_for(any_id)
476
+ return NIL_NIL_ID unless id
477
+ return_shard = nil
478
+ local_id = signed_id_operation(id) do |abs_id|
479
+ if abs_id < IDS_PER_SHARD
480
+ abs_id
481
+ elsif return_shard = lookup(abs_id / IDS_PER_SHARD)
482
+ abs_id % IDS_PER_SHARD
483
+ else
484
+ return NIL_NIL_ID
485
+ end
486
+ end
487
+ [local_id, return_shard]
488
+ end
489
+
490
+ # takes an id-ish, and returns an integral id relative to
491
+ # target_shard. returns nil if it can't be interpreted,
492
+ # or the integral version of the id if it refers to a shard
493
+ # that does not exist
494
+ def relative_id_for(any_id, source_shard, target_shard)
495
+ integral_id = integral_id_for(any_id)
496
+ local_id, shard = local_id_for(integral_id)
497
+ return integral_id unless local_id
498
+ shard ||= source_shard
499
+ return local_id if shard == target_shard
500
+ shard.global_id_for(local_id)
501
+ end
502
+
503
+ # takes an id-ish, and returns a shortened global
504
+ # string id if global, and itself if local.
505
+ # returns any_id itself if it can't be interpreted
506
+ def short_id_for(any_id)
507
+ local_id, shard = local_id_for(any_id)
508
+ return any_id unless local_id
509
+ return local_id unless shard
510
+ "#{shard.id}~#{local_id}"
511
+ end
512
+
513
+ # takes an id-ish, and returns an integral global id.
514
+ # returns nil if it can't be interpreted
515
+ def global_id_for(any_id, source_shard = nil)
516
+ id = integral_id_for(any_id)
517
+ return any_id unless id
518
+ signed_id_operation(id) do |abs_id|
519
+ if abs_id >= IDS_PER_SHARD
520
+ abs_id
521
+ else
522
+ source_shard ||= Shard.current
523
+ source_shard.global_id_for(abs_id)
524
+ end
525
+ end
526
+ end
527
+
528
+ def shard_for(any_id, source_shard = nil)
529
+ return any_id.shard if any_id.is_a?(::ActiveRecord::Base)
530
+ _, shard = local_id_for(any_id)
531
+ shard || source_shard || Shard.current
532
+ end
533
+
534
+ # given the provided option, determines whether we need to (and whether
535
+ # it's possible) to determine a reasonable default.
536
+ def determine_max_procs(max_procs_input, parallel_input=2)
537
+ max_procs = nil
538
+ if max_procs_input
539
+ max_procs = max_procs_input.to_i
540
+ max_procs = nil if max_procs == 0
541
+ else
542
+ return 1 if parallel_input.nil? || parallel_input < 1
543
+ cpus = Environment.cpu_count
544
+ if cpus && cpus > 0
545
+ max_procs = cpus * parallel_input
546
+ end
547
+ end
548
+
549
+ return max_procs
550
+ end
551
+
552
+ private
553
+ # in-process caching
554
+ def cached_shards
555
+ @cached_shards ||= {}.compare_by_identity
556
+ end
557
+
558
+ def add_to_cache(shard)
559
+ cached_shards[shard.id] = shard
560
+ end
561
+
562
+ def remove_from_cache(shard)
563
+ cached_shards.delete(shard.id)
564
+ end
565
+
566
+ def find_cached(key)
567
+ # can't simply cache the AR object since Shard has a custom serializer
568
+ # that calls this method
569
+ attributes = Switchman.cache.fetch(key) { yield&.attributes }
570
+ return nil unless attributes
571
+
572
+ shard = Shard.new
573
+ attributes.each do |attr, value|
574
+ shard.send(:"#{attr}=", value) if shard.respond_to?(:"#{attr}=")
575
+ end
576
+ shard.clear_changes_information
577
+ shard.instance_variable_set(:@new_record, false)
578
+ # connection info doesn't exist in database.yml;
579
+ # pretend the shard doesn't exist either
580
+ shard = nil unless shard.database_server
581
+ shard
582
+ end
583
+
584
+ def active_shards
585
+ Thread.current[:active_shards] ||= {}.compare_by_identity
586
+ end
587
+ end
588
+
589
+ def name
590
+ unless instance_variable_defined?(:@name)
591
+ # protect against re-entrancy
592
+ @name = nil
593
+ @name = read_attribute(:name) || default_name
594
+ end
595
+ @name
596
+ end
597
+
598
+ def name=(name)
599
+ write_attribute(:name, @name = name)
600
+ remove_instance_variable(:@name) if name == nil
601
+ end
602
+
603
+ def database_server
604
+ @database_server ||= DatabaseServer.find(self.database_server_id)
605
+ end
606
+
607
+ def database_server=(database_server)
608
+ self.database_server_id = database_server.id
609
+ @database_server = database_server
610
+ end
611
+
612
+ def primary?
613
+ self == database_server.primary_shard
614
+ end
615
+
616
+ def description
617
+ [database_server.id, name].compact.join(':')
618
+ end
619
+
620
+ # Shards are always on the default shard
621
+ def shard
622
+ Shard.default
623
+ end
624
+
625
+ def activate(*categories)
626
+ shards = hashify_categories(categories)
627
+ Shard.activate(shards) do
628
+ yield
629
+ end
630
+ end
631
+
632
+ # for use from console ONLY
633
+ def activate!(*categories)
634
+ shards = hashify_categories(categories)
635
+ Shard.activate!(shards)
636
+ nil
637
+ end
638
+
639
+ # custom serialization, since shard is self-referential
640
+ def _dump(depth)
641
+ self.id.to_s
642
+ end
643
+
644
+ def self._load(str)
645
+ lookup(str.to_i)
646
+ end
647
+
648
+ def drop_database
649
+ raise("Cannot drop the database of the default shard") if self.default?
650
+ return unless read_attribute(:name)
651
+
652
+ begin
653
+ adapter = self.database_server.config[:adapter]
654
+ sharding_config = Switchman.config || {}
655
+ drop_statement = sharding_config[adapter]&.[](:drop_statement)
656
+ drop_statement ||= sharding_config[:drop_statement]
657
+ if drop_statement
658
+ drop_statement = Array(drop_statement).dup.
659
+ map { |statement| statement.gsub('%{name}', self.name) }
660
+ end
661
+
662
+ case adapter
663
+ when 'mysql', 'mysql2'
664
+ self.activate do
665
+ ::GuardRail.activate(:deploy) do
666
+ drop_statement ||= "DROP DATABASE #{self.name}"
667
+ Array(drop_statement).each do |stmt|
668
+ ::ActiveRecord::Base.connection.execute(stmt)
669
+ end
670
+ end
671
+ end
672
+ when 'postgresql'
673
+ self.activate do
674
+ ::GuardRail.activate(:deploy) do
675
+ # Shut up, Postgres!
676
+ conn = ::ActiveRecord::Base.connection
677
+ old_proc = conn.raw_connection.set_notice_processor {}
678
+ begin
679
+ drop_statement ||= "DROP SCHEMA #{self.name} CASCADE"
680
+ Array(drop_statement).each do |stmt|
681
+ ::ActiveRecord::Base.connection.execute(stmt)
682
+ end
683
+ ensure
684
+ conn.raw_connection.set_notice_processor(&old_proc) if old_proc
685
+ end
686
+ end
687
+ end
688
+ end
689
+ rescue
690
+ logger.info "Drop failed: #{$!}"
691
+ end
692
+ end
693
+
694
+ # takes an id local to this shard, and returns a global id
695
+ def global_id_for(local_id)
696
+ return nil unless local_id
697
+ self.class.signed_id_operation(local_id) do |abs_id|
698
+ abs_id + self.id * IDS_PER_SHARD
699
+ end
700
+ end
701
+
702
+ # skip global_id.hash
703
+ def hash
704
+ id.hash
705
+ end
706
+
707
+ def destroy
708
+ raise("Cannot destroy the default shard") if self.default?
709
+ super
710
+ end
711
+
712
+ private
713
+
714
+ def clear_cache
715
+ Shard.default.activate do
716
+ Switchman.cache.delete(['shard', id].join('/'))
717
+ Switchman.cache.delete("default_shard") if default?
718
+ end
719
+ end
720
+
721
+ def default_name
722
+ database_server.shard_name(self)
723
+ end
724
+
725
+ def on_rollback
726
+ # make sure all connection pool proxies are referencing valid pools
727
+ ::ActiveRecord::Base.connection_handler.connection_pools.each do |pool|
728
+ next unless pool.is_a?(ConnectionPoolProxy)
729
+
730
+ pool.remove_shard!(self)
731
+ end
732
+ end
733
+
734
+ def hashify_categories(categories)
735
+ if categories.empty?
736
+ { :primary => self }
737
+ else
738
+ categories.inject({}) { |h, category| h[category] = self; h }
739
+ end
740
+ end
741
+
742
+ end
743
+ end