switchman 2.0.2 → 3.0.0

Sign up to get free protection for your applications and to get access to all the features.
Files changed (56) hide show
  1. checksums.yaml +4 -4
  2. data/Rakefile +10 -2
  3. data/app/models/switchman/shard.rb +234 -270
  4. data/app/models/switchman/unsharded_record.rb +7 -0
  5. data/db/migrate/20130328212039_create_switchman_shards.rb +1 -1
  6. data/db/migrate/20130328224244_create_default_shard.rb +5 -5
  7. data/db/migrate/20161206323434_add_back_default_string_limits_switchman.rb +1 -0
  8. data/db/migrate/20180828183945_add_default_shard_index.rb +1 -1
  9. data/db/migrate/20180828192111_add_timestamps_to_shards.rb +7 -5
  10. data/db/migrate/20190114212900_add_unique_name_indexes.rb +5 -3
  11. data/lib/switchman.rb +3 -3
  12. data/lib/switchman/action_controller/caching.rb +2 -2
  13. data/lib/switchman/active_record/abstract_adapter.rb +1 -0
  14. data/lib/switchman/active_record/association.rb +78 -85
  15. data/lib/switchman/active_record/attribute_methods.rb +58 -52
  16. data/lib/switchman/active_record/base.rb +58 -59
  17. data/lib/switchman/active_record/calculations.rb +73 -66
  18. data/lib/switchman/active_record/connection_pool.rb +14 -41
  19. data/lib/switchman/active_record/database_configurations.rb +34 -0
  20. data/lib/switchman/active_record/database_configurations/database_config.rb +13 -0
  21. data/lib/switchman/active_record/finder_methods.rb +11 -16
  22. data/lib/switchman/active_record/log_subscriber.rb +4 -8
  23. data/lib/switchman/active_record/migration.rb +6 -15
  24. data/lib/switchman/active_record/model_schema.rb +1 -1
  25. data/lib/switchman/active_record/persistence.rb +4 -6
  26. data/lib/switchman/active_record/postgresql_adapter.rb +42 -53
  27. data/lib/switchman/active_record/predicate_builder.rb +1 -1
  28. data/lib/switchman/active_record/query_cache.rb +18 -19
  29. data/lib/switchman/active_record/query_methods.rb +172 -181
  30. data/lib/switchman/active_record/reflection.rb +6 -10
  31. data/lib/switchman/active_record/relation.rb +27 -21
  32. data/lib/switchman/active_record/spawn_methods.rb +27 -29
  33. data/lib/switchman/active_record/statement_cache.rb +18 -35
  34. data/lib/switchman/active_record/tasks/database_tasks.rb +16 -0
  35. data/lib/switchman/active_support/cache.rb +3 -5
  36. data/lib/switchman/arel.rb +13 -8
  37. data/lib/switchman/database_server.rb +121 -142
  38. data/lib/switchman/default_shard.rb +52 -16
  39. data/lib/switchman/engine.rb +61 -57
  40. data/lib/switchman/environment.rb +4 -8
  41. data/lib/switchman/errors.rb +1 -0
  42. data/lib/switchman/guard_rail.rb +6 -19
  43. data/lib/switchman/guard_rail/relation.rb +5 -7
  44. data/lib/switchman/r_spec_helper.rb +29 -37
  45. data/lib/switchman/rails.rb +14 -12
  46. data/lib/switchman/schema_cache.rb +1 -1
  47. data/lib/switchman/sharded_instrumenter.rb +1 -1
  48. data/lib/switchman/standard_error.rb +15 -3
  49. data/lib/switchman/test_helper.rb +6 -10
  50. data/lib/switchman/version.rb +1 -1
  51. data/lib/tasks/switchman.rake +54 -69
  52. metadata +100 -44
  53. data/lib/switchman/active_record/batches.rb +0 -11
  54. data/lib/switchman/active_record/connection_handler.rb +0 -172
  55. data/lib/switchman/active_record/where_clause_factory.rb +0 -36
  56. data/lib/switchman/connection_pool_proxy.rb +0 -169
@@ -5,7 +5,7 @@ module Switchman
5
5
  module Reflection
6
6
  module AbstractReflection
7
7
  def shard(owner)
8
- if polymorphic? || klass.shard_category == owner.class.shard_category
8
+ if polymorphic? || klass.connection_classes == owner.class.connection_classes
9
9
  # polymorphic associations assume the same shard as the owning item
10
10
  owner.shard
11
11
  else
@@ -28,21 +28,17 @@ module Switchman
28
28
  # this technically belongs on AssociationReflection, but we put it on
29
29
  # ThroughReflection as well, instead of delegating to its internal
30
30
  # HasManyAssociation, losing its proper `klass`
31
- def association_scope_cache(conn, owner, &block)
32
- key = conn.prepared_statements
33
- if polymorphic?
34
- key = [key, owner._read_attribute(@foreign_type)]
35
- end
31
+ def association_scope_cache(klass, owner, &block)
32
+ key = self
33
+ key = [key, owner._read_attribute(@foreign_type)] if polymorphic?
36
34
  key = [key, shard(owner).id].flatten
37
- @association_scope_cache[key] ||= @scope_lock.synchronize {
38
- @association_scope_cache[key] ||= (::Rails.version >= "5.2" ? ::ActiveRecord::StatementCache.create(conn, &block) : block.call)
39
- }
35
+ klass.cached_find_by_statement(key, &block)
40
36
  end
41
37
  end
42
38
 
43
39
  module AssociationReflection
44
40
  def join_id_for(owner)
45
- owner.send(::Rails.version >= "5.2" ? join_foreign_key : active_record_primary_key) # use sharded id values in association binds
41
+ owner.send(join_foreign_key) # use sharded id values in association binds
46
42
  end
47
43
  end
48
44
  end
@@ -4,53 +4,54 @@ module Switchman
4
4
  module ActiveRecord
5
5
  module Relation
6
6
  def self.prepended(klass)
7
- klass::SINGLE_VALUE_METHODS.concat [ :shard, :shard_source ]
7
+ klass::SINGLE_VALUE_METHODS.concat %i[shard shard_source]
8
8
  end
9
9
 
10
10
  def initialize(*, **)
11
11
  super
12
- self.shard_value = Shard.current(klass ? klass.shard_category : :primary) unless shard_value
12
+ self.shard_value = Shard.current(klass ? klass.connection_classes : :primary) unless shard_value
13
13
  self.shard_source_value = :implicit unless shard_source_value
14
14
  end
15
15
 
16
16
  def clone
17
17
  result = super
18
- result.shard_value = Shard.current(klass ? klass.shard_category : :primary) unless shard_value
18
+ result.shard_value = Shard.current(klass ? klass.connection_classes : :primary) unless shard_value
19
19
  result
20
20
  end
21
21
 
22
22
  def merge(*)
23
23
  relation = super
24
- if relation.shard_value != self.shard_value && relation.shard_source_value == :implicit
25
- relation.shard_value = self.shard_value
26
- relation.shard_source_value = self.shard_source_value
24
+ if relation.shard_value != shard_value && relation.shard_source_value == :implicit
25
+ relation.shard_value = shard_value
26
+ relation.shard_source_value = shard_source_value
27
27
  end
28
28
  relation
29
29
  end
30
30
 
31
31
  def new(*, &block)
32
- primary_shard.activate(klass.shard_category) { super }
32
+ primary_shard.activate(klass.connection_classes) { super }
33
33
  end
34
34
 
35
35
  def create(*, &block)
36
- primary_shard.activate(klass.shard_category) { super }
36
+ primary_shard.activate(klass.connection_classes) { super }
37
37
  end
38
38
 
39
39
  def create!(*, &block)
40
- primary_shard.activate(klass.shard_category) { super }
40
+ primary_shard.activate(klass.connection_classes) { super }
41
41
  end
42
42
 
43
43
  def to_sql
44
- primary_shard.activate(klass.shard_category) { super }
44
+ primary_shard.activate(klass.connection_classes) { super }
45
45
  end
46
46
 
47
47
  def explain
48
- self.activate { |relation| relation.call_super(:explain, Relation) }
48
+ activate { |relation| relation.call_super(:explain, Relation) }
49
49
  end
50
50
 
51
51
  def records
52
52
  return @records if loaded?
53
- results = self.activate { |relation| relation.call_super(:records, Relation) }
53
+
54
+ results = activate { |relation| relation.call_super(:records, Relation) }
54
55
  case shard_value
55
56
  when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
56
57
  @records = results
@@ -59,7 +60,7 @@ module Switchman
59
60
  results
60
61
  end
61
62
 
62
- %I{update_all delete_all}.each do |method|
63
+ %I[update_all delete_all].each do |method|
63
64
  class_eval <<-RUBY, __FILE__, __LINE__ + 1
64
65
  def #{method}(*args)
65
66
  result = self.activate { |relation| relation.call_super(#{method.inspect}, Relation, *args) }
@@ -74,15 +75,20 @@ module Switchman
74
75
  loose_mode = options[:loose] && is_integer
75
76
  # loose_mode: if we don't care about getting exactly batch_size ids in between
76
77
  # 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
+ values = loose_mode ? 'MIN(id)' : 'MIN(id), MAX(id)'
78
79
 
79
80
  batch_size = options[:batch_size].try(:to_i) || 1000
80
81
  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
+ as_id = ' AS id' unless primary_key == 'id'
82
83
  subquery_scope = except(:select).select("#{quoted_primary_key}#{as_id}").reorder(primary_key.to_sym).limit(loose_mode ? 1 : batch_size)
83
84
  subquery_scope = subquery_scope.where("#{quoted_primary_key} <= ?", options[:end_at]) if options[:end_at]
84
85
 
85
- first_subquery_scope = options[:start_at] ? subquery_scope.where("#{quoted_primary_key} >= ?", options[:start_at]) : subquery_scope
86
+ first_subquery_scope = if options[:start_at]
87
+ subquery_scope.where("#{quoted_primary_key} >= ?",
88
+ options[:start_at])
89
+ else
90
+ subquery_scope
91
+ end
86
92
 
87
93
  ids = connection.select_rows("SELECT #{values} FROM (#{first_subquery_scope.to_sql}) AS subquery").first
88
94
 
@@ -99,16 +105,16 @@ module Switchman
99
105
 
100
106
  def activate(&block)
101
107
  shards = all_shards
102
- if (Array === shards && shards.length == 1)
103
- if shards.first == DefaultShard || shards.first == Shard.current(klass.shard_category)
108
+ if Array === shards && shards.length == 1
109
+ if shards.first == DefaultShard || shards.first == Shard.current(klass.connection_classes)
104
110
  yield(self, shards.first)
105
111
  else
106
- shards.first.activate(klass.shard_category) { yield(self, shards.first) }
112
+ shards.first.activate(klass.connection_classes) { yield(self, shards.first) }
107
113
  end
108
114
  else
109
115
  # TODO: implement local limit to avoid querying extra shards
110
- Shard.with_each_shard(shards, [klass.shard_category]) do
111
- shard(Shard.current(klass.shard_category), :to_a).activate(&block)
116
+ Shard.with_each_shard(shards, [klass.connection_classes]) do
117
+ shard(Shard.current(klass.connection_classes), :to_a).activate(&block)
112
118
  end
113
119
  end
114
120
  end
@@ -3,27 +3,27 @@
3
3
  module Switchman
4
4
  module ActiveRecord
5
5
  module SpawnMethods
6
- def shard_values_for_merge(r)
7
- if shard_value != r.shard_value
8
- if r.shard_source_value == :implicit
6
+ def shard_values_for_merge(rhs)
7
+ if shard_value != rhs.shard_value
8
+ if rhs.shard_source_value == :implicit
9
9
  final_shard_value = shard_value
10
10
  final_primary_shard = primary_shard
11
11
  final_shard_source_value = shard_source_value
12
12
  elsif shard_source_value == :implicit
13
- final_shard_value = r.shard_value
14
- final_primary_shard = r.primary_shard
15
- final_shard_source_value = r.shard_source_value
13
+ final_shard_value = rhs.shard_value
14
+ final_primary_shard = rhs.primary_shard
15
+ final_shard_source_value = rhs.shard_source_value
16
16
  else
17
- final_shard_source_value = [:explicit, :association].detect do |source_value|
18
- shard_source_value == source_value || r.shard_source_value == source_value
17
+ final_shard_source_value = %i[explicit association].detect do |source_value|
18
+ shard_source_value == source_value || rhs.shard_source_value == source_value
19
19
  end
20
- raise "unknown shard_source_value" unless final_shard_source_value
20
+ raise 'unknown shard_source_value' unless final_shard_source_value
21
21
 
22
22
  # have to merge shard_value
23
23
  lhs_shard_value = all_shards
24
- rhs_shard_value = r.all_shards
25
- if (::ActiveRecord::Relation === lhs_shard_value &&
26
- ::ActiveRecord::Relation === rhs_shard_value)
24
+ rhs_shard_value = rhs.all_shards
25
+ if ::ActiveRecord::Relation === lhs_shard_value &&
26
+ ::ActiveRecord::Relation === rhs_shard_value
27
27
  final_shard_value = lhs_shard_value.merge(rhs_shard_value)
28
28
  final_primary_shard = Shard.default
29
29
  else
@@ -32,27 +32,25 @@ module Switchman
32
32
  final_shard_value = final_shard_value.first if final_shard_value.length == 1
33
33
  end
34
34
  end
35
- elsif shard_source_value != r.shard_source_value
36
- final_shard_source_value = [:explicit, :association, :implicit].detect do |source_value|
37
- shard_source_value == source_value || r.shard_source_value == source_value
35
+ elsif shard_source_value != rhs.shard_source_value
36
+ final_shard_source_value = %i[explicit association implicit].detect do |source_value|
37
+ shard_source_value == source_value || rhs.shard_source_value == source_value
38
38
  end
39
- raise "unknown shard_source_value" unless final_shard_source_value
40
- else
41
- # nothing fancy
39
+ raise 'unknown shard_source_value' unless final_shard_source_value
42
40
  end
43
41
 
44
42
  [final_shard_value, final_primary_shard, final_shard_source_value]
45
43
  end
46
44
 
47
- def merge!(r)
48
- return super unless ::ActiveRecord::Relation === r
45
+ def merge!(rhs)
46
+ return super unless ::ActiveRecord::Relation === rhs
49
47
 
50
48
  # have to figure out shard stuff *before* conditions are merged
51
- final_shard_value, final_primary_shard, final_shard_source_value = shard_values_for_merge(r)
49
+ final_shard_value, final_primary_shard, final_shard_source_value = shard_values_for_merge(rhs)
52
50
 
53
51
  return super unless final_shard_source_value
54
52
 
55
- if !final_shard_value
53
+ unless final_shard_value
56
54
  super
57
55
  self.shard_source_value = final_shard_source_value
58
56
  return self
@@ -61,16 +59,16 @@ module Switchman
61
59
  return none! if final_shard_value == []
62
60
 
63
61
  # change the primary shard if necessary before merging
64
- if primary_shard != final_primary_shard && r.primary_shard != final_primary_shard
62
+ if primary_shard != final_primary_shard && rhs.primary_shard != final_primary_shard
65
63
  shard!(final_primary_shard)
66
- r = r.shard(final_primary_shard)
67
- super(r)
64
+ rhs = rhs.shard(final_primary_shard)
65
+ super(rhs)
68
66
  elsif primary_shard != final_primary_shard
69
67
  shard!(final_primary_shard)
70
- super(r)
71
- elsif r.primary_shard != final_primary_shard
72
- r = r.shard(final_primary_shard)
73
- super(r)
68
+ super(rhs)
69
+ elsif rhs.primary_shard != final_primary_shard
70
+ rhs = rhs.shard(final_primary_shard)
71
+ super(rhs)
74
72
  else
75
73
  super
76
74
  end
@@ -7,18 +7,13 @@ module Switchman
7
7
  def create(connection, &block)
8
8
  relation = block.call ::ActiveRecord::StatementCache::Params.new
9
9
 
10
- if ::Rails.version >= "5.2"
11
- query_builder, binds = connection.cacheable_query(self, relation.arel)
12
- bind_map = ::ActiveRecord::StatementCache::BindMap.new(binds)
13
- new(relation.arel, bind_map, relation.klass)
14
- else
15
- bind_map = ::ActiveRecord::StatementCache::BindMap.new(relation.bound_attributes)
16
- new relation.arel, bind_map
17
- end
10
+ _query_builder, binds = connection.cacheable_query(self, relation.arel)
11
+ bind_map = ::ActiveRecord::StatementCache::BindMap.new(binds)
12
+ new(relation.arel, bind_map, relation.klass)
18
13
  end
19
14
  end
20
15
 
21
- def initialize(arel, bind_map, klass=nil)
16
+ def initialize(arel, bind_map, klass = nil)
22
17
  @arel = arel
23
18
  @bind_map = bind_map
24
19
  @klass = klass
@@ -31,49 +26,39 @@ module Switchman
31
26
  # (e.g. infer from the primary key or use the current shard)
32
27
 
33
28
  def execute(*args)
34
- if ::Rails.version >= '5.2'
35
- params, connection = args
36
- klass = @klass
37
- else
38
- params, klass, connection = args
39
- end
29
+ params, connection = args
30
+ klass = @klass
40
31
  target_shard = nil
41
- if primary_index = bind_map.primary_value_index
32
+ if (primary_index = bind_map.primary_value_index)
42
33
  primary_value = params[primary_index]
43
34
  target_shard = Shard.local_id_for(primary_value)[1]
44
35
  end
45
- current_shard = Shard.current(klass.shard_category)
36
+ current_shard = Shard.current(klass.connection_classes)
46
37
  target_shard ||= current_shard
47
38
 
48
39
  bind_values = bind_map.bind(params, current_shard, target_shard)
49
40
 
50
- target_shard.activate(klass.shard_category) do
41
+ target_shard.activate(klass.connection_classes) do
51
42
  sql = qualified_query_builder(target_shard, klass).sql_for(bind_values, connection)
52
43
  klass.find_by_sql(sql, bind_values)
53
44
  end
54
45
  end
55
46
 
56
- if ::Rails.version < '5.2'
57
- def qualified_query_builder(shard, klass)
58
- @qualified_query_builders[shard.id] ||= klass.connection.cacheable_query(self.class, @arel)
59
- end
60
- else
61
- def qualified_query_builder(shard, klass)
62
- @qualified_query_builders[shard.id] ||= klass.connection.cacheable_query(self.class, @arel).first
63
- end
47
+ def qualified_query_builder(shard, klass)
48
+ @qualified_query_builders[shard.id] ||= klass.connection.cacheable_query(self.class, @arel).first
64
49
  end
65
50
 
66
51
  module BindMap
67
52
  # performs id transposition here instead of query_methods.rb
68
53
  def bind(values, current_shard, target_shard)
69
54
  bas = @bound_attributes.dup
70
- @indexes.each_with_index do |offset,i|
55
+ @indexes.each_with_index do |offset, i|
71
56
  ba = bas[offset]
72
- if ba.is_a?(::ActiveRecord::Relation::QueryAttribute) && ba.value.sharded
73
- new_value = Shard.relative_id_for(values[i], current_shard, target_shard || current_shard)
74
- else
75
- new_value = values[i]
76
- end
57
+ new_value = if ba.is_a?(::ActiveRecord::Relation::QueryAttribute) && ba.value.sharded
58
+ Shard.relative_id_for(values[i], current_shard, target_shard || current_shard)
59
+ else
60
+ values[i]
61
+ end
77
62
  bas[offset] = ba.with_cast_value(new_value)
78
63
  end
79
64
  bas
@@ -83,9 +68,7 @@ module Switchman
83
68
  primary_ba_index = @bound_attributes.index do |ba|
84
69
  ba.is_a?(::ActiveRecord::Relation::QueryAttribute) && ba.value.primary
85
70
  end
86
- if primary_ba_index
87
- @indexes.index(primary_ba_index)
88
- end
71
+ @indexes.index(primary_ba_index) if primary_ba_index
89
72
  end
90
73
  end
91
74
  end
@@ -0,0 +1,16 @@
1
+ # frozen_string_literal: true
2
+
3
+ module Switchman
4
+ module ActiveRecord
5
+ module Tasks
6
+ module DatabaseTasks
7
+ def drop(*)
8
+ super
9
+ # no really, it's gone
10
+ Switchman.cache.delete('default_shard')
11
+ Shard.default(reload: true)
12
+ end
13
+ end
14
+ end
15
+ end
16
+ end
@@ -8,10 +8,8 @@ module Switchman
8
8
  store = super
9
9
  # can't use defined?, because it's a _ruby_ autoloaded constant,
10
10
  # so just checking that will cause it to get required
11
- if store.class.name == "ActiveSupport::Cache::RedisCacheStore" && !::ActiveSupport::Cache::RedisCacheStore.ancestors.include?(RedisCacheStore)
12
- ::ActiveSupport::Cache::RedisCacheStore.prepend(RedisCacheStore)
13
- end
14
- store.options[:namespace] ||= lambda { Shard.current.default? ? nil : "shard_#{Shard.current.id}" }
11
+ ::ActiveSupport::Cache::RedisCacheStore.prepend(RedisCacheStore) if store.instance_of?(ActiveSupport::Cache::RedisCacheStore) && !::ActiveSupport::Cache::RedisCacheStore.ancestors.include?(RedisCacheStore)
12
+ store.options[:namespace] ||= -> { Shard.current.default? ? nil : "shard_#{Shard.current.id}" }
15
13
  store
16
14
  end
17
15
  end
@@ -22,7 +20,7 @@ module Switchman
22
20
  # unfortunately, it uses the keys command, which is extraordinarily inefficient in a large redis instance
23
21
  # fortunately, we can assume we control the entire instance, because we set up the namespacing, so just
24
22
  # always unset it temporarily for clear calls
25
- namespace = nil
23
+ namespace = nil # rubocop:disable Lint/ShadowedArgument
26
24
  super
27
25
  end
28
26
  end
@@ -3,29 +3,34 @@
3
3
  module Switchman
4
4
  module Arel
5
5
  module Table
6
- def model
7
- type_caster.model
6
+ def klass
7
+ @klass || ::ActiveRecord::Base
8
8
  end
9
9
  end
10
+
10
11
  module Visitors
11
12
  module ToSql
12
- def visit_Arel_Nodes_TableAlias *args
13
+ # rubocop:disable Naming/MethodName
14
+
15
+ def visit_Arel_Nodes_TableAlias(*args)
13
16
  o, collector = args
14
17
  collector = visit o.relation, collector
15
- collector << " "
18
+ collector << ' '
16
19
  collector << quote_local_table_name(o.name)
17
20
  end
18
21
 
19
- def visit_Arel_Attributes_Attribute *args
22
+ def visit_Arel_Attributes_Attribute(*args)
20
23
  o = args.first
21
24
  join_name = o.relation.table_alias || o.relation.name
22
25
  result = "#{quote_local_table_name join_name}.#{quote_column_name o.name}"
23
- result = args.last << result
24
- result
26
+ args.last << result
25
27
  end
26
28
 
27
- def quote_local_table_name name
29
+ # rubocop:enable Naming/MethodName
30
+
31
+ def quote_local_table_name(name)
28
32
  return name if ::Arel::Nodes::SqlLiteral === name
33
+
29
34
  @connection.quote_local_table_name(name)
30
35
  end
31
36
  end
@@ -1,6 +1,6 @@
1
1
  # frozen_string_literal: true
2
2
 
3
- require "securerandom"
3
+ require 'securerandom'
4
4
 
5
5
  module Switchman
6
6
  class DatabaseServer
@@ -13,16 +13,22 @@ module Switchman
13
13
  database_servers.values
14
14
  end
15
15
 
16
+ def all_roles
17
+ @all_roles ||= all.map(&:roles).flatten.uniq
18
+ end
19
+
16
20
  def find(id_or_all)
17
- return self.all if id_or_all == :all
18
- return id_or_all.map { |id| self.database_servers[id || ::Rails.env] }.compact.uniq if id_or_all.is_a?(Array)
21
+ return all if id_or_all == :all
22
+ return id_or_all.map { |id| database_servers[id || ::Rails.env] }.compact.uniq if id_or_all.is_a?(Array)
23
+
19
24
  database_servers[id_or_all || ::Rails.env]
20
25
  end
21
26
 
22
27
  def create(settings = {})
23
- raise "database servers should be set up in database.yml" unless ::Rails.env.test?
28
+ raise 'database servers should be set up in database.yml' unless ::Rails.env.test?
29
+
24
30
  id = settings[:id]
25
- if !id
31
+ unless id
26
32
  @id ||= 0
27
33
  @id += 1
28
34
  id = @id
@@ -30,25 +36,42 @@ module Switchman
30
36
  server = DatabaseServer.new(id.to_s, settings)
31
37
  server.instance_variable_set(:@fake, true)
32
38
  database_servers[server.id] = server
39
+ ::ActiveRecord::Base.configurations.configurations <<
40
+ ::ActiveRecord::DatabaseConfigurations::HashConfig.new(::Rails.env, "#{server.id}/primary", settings)
41
+ Shard.send(:initialize_sharding)
42
+ server
33
43
  end
34
44
 
35
45
  def server_for_new_shard
36
46
  servers = all.select { |s| s.config[:open] }
37
47
  return find(nil) if servers.empty?
48
+
38
49
  servers[rand(servers.length)]
39
50
  end
40
51
 
41
52
  private
53
+
54
+ def reference_role(role)
55
+ return if all_roles.include?(role)
56
+
57
+ @all_roles << role
58
+ Shard.send(:initialize_sharding)
59
+ end
60
+
42
61
  def database_servers
43
62
  unless @database_servers
44
63
  @database_servers = {}.with_indifferent_access
45
- if ::Rails.version >= '6.0'
46
- ::ActiveRecord::Base.configurations.configurations.each do |config|
47
- @database_servers[config.env_name] = DatabaseServer.new(config.env_name, config.config)
64
+ ::ActiveRecord::Base.configurations.configurations.each do |config|
65
+ if config.name.include?('/')
66
+ name, role = config.name.split('/')
67
+ else
68
+ name, role = config.env_name, config.name
48
69
  end
49
- else
50
- ::ActiveRecord::Base.configurations.each do |(id, config)|
51
- @database_servers[id] = DatabaseServer.new(id, config)
70
+
71
+ if role == 'primary'
72
+ @database_servers[name] = DatabaseServer.new(config.env_name, config.configuration_hash)
73
+ else
74
+ @database_servers[name].roles << role
52
75
  end
53
76
  end
54
77
  end
@@ -56,15 +79,33 @@ module Switchman
56
79
  end
57
80
  end
58
81
 
82
+ attr_reader :roles
83
+
59
84
  def initialize(id = nil, config = {})
60
85
  @id = id
61
86
  @config = config.deep_symbolize_keys
62
87
  @configs = {}
88
+ @roles = [:primary]
89
+ end
90
+
91
+ def connects_to_hash
92
+ self.class.all_roles.map do |role|
93
+ config_role = role
94
+ config_role = :primary unless roles.include?(role)
95
+ config_name = :"#{id}/#{config_role}"
96
+ config_name = :primary if id == ::Rails.env && config_role == :primary
97
+ [role.to_sym, config_name]
98
+ end.to_h
63
99
  end
64
100
 
65
101
  def destroy
66
- raise "database servers should be set up in database.yml" unless ::Rails.env.test?
67
- self.class.send(:database_servers).delete(self.id) if self.id
102
+ self.class.send(:database_servers).delete(id) if id
103
+ Shard.sharded_models.each do |klass|
104
+ self.class.all_roles.each do |role|
105
+ klass.connection_handler.remove_connection_pool(klass.connection_specification_name, role: role,
106
+ shard: id.to_sym)
107
+ end
108
+ end
68
109
  end
69
110
 
70
111
  def fake?
@@ -73,14 +114,15 @@ module Switchman
73
114
 
74
115
  def config(environment = :primary)
75
116
  @configs[environment] ||= begin
76
- if @config[environment].is_a?(Array)
117
+ case @config[environment]
118
+ when Array
77
119
  @config[environment].map do |config|
78
120
  config = @config.merge((config || {}).symbolize_keys)
79
121
  # make sure GuardRail doesn't get any brilliant ideas about choosing the first possible server
80
122
  config.delete(environment)
81
123
  config
82
124
  end
83
- elsif @config[environment].is_a?(Hash)
125
+ when Hash
84
126
  @config.merge(@config[environment])
85
127
  else
86
128
  @config
@@ -111,172 +153,109 @@ module Switchman
111
153
  guard!(old_env)
112
154
  end
113
155
 
114
- def shareable?
115
- @shareable_environment_key ||= []
116
- environment = guard_rail_environment
117
- explicit_user = ::GuardRail.global_config[:username]
118
- return @shareable if @shareable_environment_key == [environment, explicit_user]
119
- @shareable_environment_key = [environment, explicit_user]
120
- if explicit_user
121
- username = explicit_user
122
- else
123
- config = self.config(environment)
124
- config = config.first if config.is_a?(Array)
125
- username = config[:username]
126
- end
127
- @shareable = username !~ /%?\{[a-zA-Z0-9_]+\}/
128
- end
129
-
130
156
  def shards
131
- if self.id == ::Rails.env
132
- Shard.where("database_server_id IS NULL OR database_server_id=?", self.id)
157
+ if id == ::Rails.env
158
+ Shard.where('database_server_id IS NULL OR database_server_id=?', id)
133
159
  else
134
- Shard.where(:database_server_id => self.id)
160
+ Shard.where(database_server_id: id)
135
161
  end
136
162
  end
137
163
 
138
- def pool_key
139
- self.id == ::Rails.env ? nil : self.id
140
- end
141
-
142
- def create_new_shard(options = {})
143
- raise NotImplementedError.new("Cannot create new shards when sharding isn't initialized") unless Shard.default.is_a?(Shard)
144
-
145
- name = options[:name]
146
- create_schema = options[:schema]
147
- # look for another shard associated with this db
148
- other_shard = self.shards.where("name<>':memory:' OR name IS NULL").order(:id).first
149
-
150
- case config[:adapter]
151
- when 'postgresql'
152
- create_statement = lambda { "CREATE SCHEMA #{name}" }
153
- password = " PASSWORD #{::ActiveRecord::Base.connection.quote(config[:password])}" if config[:password]
154
- else
155
- create_statement = lambda { "CREATE DATABASE #{name}" }
164
+ def create_new_shard(id: nil, name: nil, schema: true)
165
+ unless Shard.default.is_a?(Shard)
166
+ raise NotImplementedError,
167
+ "Cannot create new shards when sharding isn't initialized"
156
168
  end
169
+
170
+ create_statement = -> { "CREATE SCHEMA #{name}" }
171
+ password = " PASSWORD #{::ActiveRecord::Base.connection.quote(config[:password])}" if config[:password]
157
172
  sharding_config = Switchman.config
158
173
  config_create_statement = sharding_config[config[:adapter]]&.[](:create_statement)
159
174
  config_create_statement ||= sharding_config[:create_statement]
160
175
  if config_create_statement
161
176
  create_commands = Array(config_create_statement).dup
162
177
  create_statement = lambda {
163
- create_commands.map { |statement| statement.gsub('%{name}', name).gsub('%{password}', password || '') }
178
+ create_commands.map { |statement| format(statement, name: name, password: password) }
164
179
  }
165
180
  end
166
181
 
167
- create_shard = lambda do
168
- shard_id = options.fetch(:id) do
169
- case config[:adapter]
170
- when 'postgresql'
171
- id_seq = Shard.connection.quote(Shard.connection.quote_table_name('switchman_shards_id_seq'))
172
- next_id = Shard.connection.select_value("SELECT nextval(#{id_seq})")
173
- next_id.to_i
174
- else
175
- nil
176
- end
177
- end
182
+ id ||= begin
183
+ id_seq = Shard.connection.quote(Shard.connection.quote_table_name('switchman_shards_id_seq'))
184
+ next_id = Shard.connection.select_value("SELECT nextval(#{id_seq})")
185
+ next_id.to_i
186
+ end
178
187
 
179
- if name.nil?
180
- base_name = self.config[:database].to_s % self.config
181
- base_name = nil if base_name == ':memory:'
182
- base_name << '_' if base_name
183
- base_id = shard_id || SecureRandom.uuid
184
- name = "#{base_name}shard_#{base_id}"
185
- end
188
+ name ||= "#{config[:database]}_shard_#{id}"
186
189
 
187
- shard = Shard.create!(:id => shard_id,
188
- :name => name,
189
- :database_server => self)
190
+ Shard.connection.transaction do
191
+ shard = Shard.create!(id: id,
192
+ name: name,
193
+ database_server_id: self.id)
190
194
  schema_already_existed = false
191
195
 
192
196
  begin
193
197
  self.class.creating_new_shard = true
194
- shard.activate(*Shard.categories) do
195
- ::GuardRail.activate(:deploy) do
196
- begin
197
- if create_statement
198
- if (::ActiveRecord::Base.connection.select_value("SELECT 1 FROM pg_namespace WHERE nspname=#{::ActiveRecord::Base.connection.quote(name)}"))
199
- schema_already_existed = true
200
- raise "This schema already exists; cannot overwrite"
201
- end
202
- Array(create_statement.call).each do |stmt|
203
- ::ActiveRecord::Base.connection.execute(stmt)
204
- end
205
- # have to disconnect and reconnect to the correct db
206
- if self.shareable? && other_shard
207
- other_shard.activate { ::ActiveRecord::Base.connection }
208
- else
209
- ::ActiveRecord::Base.connection_pool.current_pool.disconnect!
210
- end
211
- end
212
- old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor {} if config[:adapter] == 'postgresql'
213
- old_verbose = ::ActiveRecord::Migration.verbose
214
- ::ActiveRecord::Migration.verbose = false
215
-
216
- unless create_schema == false
217
- reset_column_information
218
-
219
- migrate = ::Rails.version >= '5.2' ?
220
- -> { ::ActiveRecord::Base.connection.migration_context.migrate } :
221
- -> { ::ActiveRecord::Migrator.migrate(::ActiveRecord::Migrator.migrations_paths) }
222
- if ::ActiveRecord::Base.connection.supports_ddl_transactions?
223
- ::ActiveRecord::Base.connection.transaction(requires_new: true, &migrate)
224
- else
225
- migrate.call
226
- end
227
- reset_column_information
228
- ::ActiveRecord::Base.descendants.reject { |m| m == Shard || !m.table_exists? }.each(&:define_attribute_methods)
198
+ DatabaseServer.send(:reference_role, :deploy)
199
+ ::ActiveRecord::Base.connected_to(shard: self.id.to_sym, role: :deploy) do
200
+ if create_statement
201
+ if ::ActiveRecord::Base.connection.select_value("SELECT 1 FROM pg_namespace WHERE nspname=#{::ActiveRecord::Base.connection.quote(name)}")
202
+ schema_already_existed = true
203
+ raise 'This schema already exists; cannot overwrite'
204
+ end
205
+ Array(create_statement.call).each do |stmt|
206
+ ::ActiveRecord::Base.connection.execute(stmt)
207
+ end
208
+ end
209
+ if config[:adapter] == 'postgresql'
210
+ old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor do
211
+ end
212
+ end
213
+ old_verbose = ::ActiveRecord::Migration.verbose
214
+ ::ActiveRecord::Migration.verbose = false
215
+
216
+ unless schema == false
217
+ shard.activate do
218
+ reset_column_information
219
+
220
+ ::ActiveRecord::Base.connection.transaction(requires_new: true) do
221
+ ::ActiveRecord::Base.connection.migration_context.migrate
229
222
  end
230
- ensure
231
- ::ActiveRecord::Migration.verbose = old_verbose
232
- ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
223
+ reset_column_information
224
+ ::ActiveRecord::Base.descendants.reject do |m|
225
+ m <= UnshardedRecord || !m.table_exists?
226
+ end.each(&:define_attribute_methods)
233
227
  end
234
228
  end
229
+ ensure
230
+ ::ActiveRecord::Migration.verbose = old_verbose
231
+ ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
235
232
  end
236
233
  shard
237
234
  rescue
238
235
  shard.destroy
239
- unless schema_already_existed
240
- shard.drop_database rescue nil
241
- end
242
- reset_column_information unless create_schema == false rescue nil
236
+ shard.drop_database rescue nil unless schema_already_existed
237
+ reset_column_information unless schema == false rescue nil
243
238
  raise
244
239
  ensure
245
240
  self.class.creating_new_shard = false
246
241
  end
247
242
  end
248
-
249
- if Shard.connection.supports_ddl_transactions? && self.shareable? && other_shard
250
- Shard.transaction do
251
- other_shard.activate do
252
- ::ActiveRecord::Base.connection.transaction(&create_shard)
253
- end
254
- end
255
- else
256
- create_shard.call
257
- end
258
243
  end
259
244
 
260
245
  def cache_store
261
- unless @cache_store
262
- @cache_store = Switchman.config[:cache_map][self.id] || Switchman.config[:cache_map][::Rails.env]
263
- end
246
+ @cache_store ||= Switchman.config[:cache_map][id] || Switchman.config[:cache_map][::Rails.env]
264
247
  @cache_store
265
248
  end
266
249
 
267
250
  def shard_name(shard)
268
- if config[:shard_name]
269
- config[:shard_name]
270
- elsif config[:adapter] == 'postgresql'
271
- if shard == :bootstrap
272
- # rescue nil because the database may not exist yet; if it doesn't,
273
- # it will shortly, and this will be re-invoked
274
- ::ActiveRecord::Base.connection.current_schemas.first rescue nil
275
- else
276
- shard.activate { ::ActiveRecord::Base.connection_pool.default_schema }
277
- end
251
+ return config[:shard_name] if config[:shard_name]
252
+
253
+ if shard == :bootstrap
254
+ # rescue nil because the database may not exist yet; if it doesn't,
255
+ # it will shortly, and this will be re-invoked
256
+ ::ActiveRecord::Base.connection.current_schemas.first rescue nil
278
257
  else
279
- config[:database]
258
+ shard.activate { ::ActiveRecord::Base.connection_pool.default_schema }
280
259
  end
281
260
  end
282
261
 
@@ -290,9 +269,9 @@ module Switchman
290
269
  end
291
270
 
292
271
  private
272
+
293
273
  def reset_column_information
294
- ::ActiveRecord::Base.descendants.reject { |m| m == Shard }.each(&:reset_column_information)
295
- ::ActiveRecord::Base.connection_handler.switchman_connection_pool_proxies.each { |pool| pool.schema_cache.clear! }
274
+ ::ActiveRecord::Base.descendants.reject { |m| m <= UnshardedRecord }.each(&:reset_column_information)
296
275
  end
297
276
  end
298
277
  end