switchman 1.14.8 → 1.15.2

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