switchman 0.0.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (122) hide show
  1. checksums.yaml +7 -0
  2. data/Rakefile +30 -0
  3. data/app/models/switchman/shard.rb +502 -0
  4. data/db/migrate/20130328212039_create_switchman_shards.rb +9 -0
  5. data/db/migrate/20130328224244_create_default_shard.rb +9 -0
  6. data/lib/switchman.rb +9 -0
  7. data/lib/switchman/active_record/abstract_adapter.rb +11 -0
  8. data/lib/switchman/active_record/association.rb +108 -0
  9. data/lib/switchman/active_record/attribute_methods.rb +104 -0
  10. data/lib/switchman/active_record/base.rb +95 -0
  11. data/lib/switchman/active_record/calculations.rb +63 -0
  12. data/lib/switchman/active_record/connection_handler.rb +147 -0
  13. data/lib/switchman/active_record/connection_pool.rb +117 -0
  14. data/lib/switchman/active_record/finder_methods.rb +25 -0
  15. data/lib/switchman/active_record/log_subscriber.rb +43 -0
  16. data/lib/switchman/active_record/postgresql_adapter.rb +13 -0
  17. data/lib/switchman/active_record/query_cache.rb +12 -0
  18. data/lib/switchman/active_record/query_methods.rb +184 -0
  19. data/lib/switchman/active_record/relation.rb +69 -0
  20. data/lib/switchman/cache_extensions.rb +12 -0
  21. data/lib/switchman/connection_pool_proxy.rb +62 -0
  22. data/lib/switchman/database_server.rb +197 -0
  23. data/lib/switchman/default_shard.rb +28 -0
  24. data/lib/switchman/engine.rb +91 -0
  25. data/lib/switchman/r_spec_helper.rb +124 -0
  26. data/lib/switchman/shackles.rb +34 -0
  27. data/lib/switchman/test_helper.rb +65 -0
  28. data/lib/switchman/version.rb +3 -0
  29. data/spec/dummy/Rakefile +7 -0
  30. data/spec/dummy/app/models/appendage.rb +24 -0
  31. data/spec/dummy/app/models/digit.rb +9 -0
  32. data/spec/dummy/app/models/feature.rb +5 -0
  33. data/spec/dummy/app/models/mirror_user.rb +5 -0
  34. data/spec/dummy/app/models/user.rb +23 -0
  35. data/spec/dummy/config.ru +4 -0
  36. data/spec/dummy/config/application.rb +59 -0
  37. data/spec/dummy/config/boot.rb +10 -0
  38. data/spec/dummy/config/database.yml +17 -0
  39. data/spec/dummy/config/database.yml.example +25 -0
  40. data/spec/dummy/config/environment.rb +5 -0
  41. data/spec/dummy/config/environments/development.rb +37 -0
  42. data/spec/dummy/config/environments/production.rb +67 -0
  43. data/spec/dummy/config/environments/test.rb +37 -0
  44. data/spec/dummy/config/initializers/backtrace_silencers.rb +7 -0
  45. data/spec/dummy/config/initializers/secret_token.rb +7 -0
  46. data/spec/dummy/config/initializers/session_store.rb +8 -0
  47. data/spec/dummy/config/initializers/wrap_parameters.rb +14 -0
  48. data/spec/dummy/config/routes.rb +8 -0
  49. data/spec/dummy/db/migrate/20130403132607_create_users.rb +10 -0
  50. data/spec/dummy/db/migrate/20130411202442_create_appendages.rb +10 -0
  51. data/spec/dummy/db/migrate/20130411202551_create_mirror_users.rb +9 -0
  52. data/spec/dummy/db/migrate/20131022202028_create_digits.rb +10 -0
  53. data/spec/dummy/db/migrate/20131206172923_create_features.rb +12 -0
  54. data/spec/dummy/db/schema.rb +57 -0
  55. data/spec/dummy/log/development.log +504 -0
  56. data/spec/dummy/log/test.log +29907 -0
  57. data/spec/dummy/script/rails +6 -0
  58. data/spec/dummy/tmp/cache/2E2/830/shard%2F2 +0 -0
  59. data/spec/dummy/tmp/cache/2E3/840/shard%2F3 +0 -0
  60. data/spec/dummy/tmp/cache/313/970/shard%2F30 +0 -0
  61. data/spec/dummy/tmp/cache/314/980/shard%2F31 +0 -0
  62. data/spec/dummy/tmp/cache/316/980/shard%2F15 +1 -0
  63. data/spec/dummy/tmp/cache/316/9D0/shard%2F60 +0 -0
  64. data/spec/dummy/tmp/cache/317/990/shard%2F16 +0 -0
  65. data/spec/dummy/tmp/cache/317/9C0/shard%2F43 +1 -0
  66. data/spec/dummy/tmp/cache/317/9E0/shard%2F61 +0 -0
  67. data/spec/dummy/tmp/cache/318/9A0/shard%2F17 +0 -0
  68. data/spec/dummy/tmp/cache/318/9D0/shard%2F44 +0 -0
  69. data/spec/dummy/tmp/cache/318/9F0/shard%2F62 +1 -0
  70. data/spec/dummy/tmp/cache/319/9E0/shard%2F45 +0 -0
  71. data/spec/dummy/tmp/cache/319/9F0/shard%2F54 +1 -0
  72. data/spec/dummy/tmp/cache/319/A10/shard%2F72 +1 -0
  73. data/spec/dummy/tmp/cache/319/A30/shard%2F90 +0 -0
  74. data/spec/dummy/tmp/cache/31B/9E0/shard%2F29 +1 -0
  75. data/spec/dummy/tmp/cache/321/AA0/shard%2F89 +0 -0
  76. data/spec/dummy/tmp/cache/322/AC0/shard%2F99 +1 -0
  77. data/spec/dummy/tmp/cache/344/D70/shard%2F103 +1 -0
  78. data/spec/dummy/tmp/cache/345/D80/shard%2F104 +0 -0
  79. data/spec/dummy/tmp/cache/345/DB0/shard%2F131 +1 -0
  80. data/spec/dummy/tmp/cache/345/DC0/shard%2F140 +0 -0
  81. data/spec/dummy/tmp/cache/346/D90/shard%2F105 +0 -0
  82. data/spec/dummy/tmp/cache/346/DB0/shard%2F123 +0 -0
  83. data/spec/dummy/tmp/cache/346/DD0/shard%2F222 +1 -0
  84. data/spec/dummy/tmp/cache/346/DE0/shard%2F150 +0 -0
  85. data/spec/dummy/tmp/cache/346/DF0/shard%2F240 +1 -0
  86. data/spec/dummy/tmp/cache/347/DA0/shard%2F106 +1 -0
  87. data/spec/dummy/tmp/cache/347/DC0/shard%2F124 +0 -0
  88. data/spec/dummy/tmp/cache/347/DC0/shard%2F205 +1 -0
  89. data/spec/dummy/tmp/cache/347/E10/shard%2F250 +1 -0
  90. data/spec/dummy/tmp/cache/348/DF0/shard%2F143 +1 -0
  91. data/spec/dummy/tmp/cache/348/DF0/shard%2F224 +1 -0
  92. data/spec/dummy/tmp/cache/348/E10/shard%2F161 +1 -0
  93. data/spec/dummy/tmp/cache/349/DD0/shard%2F117 +1 -0
  94. data/spec/dummy/tmp/cache/349/E40/shard%2F180 +1 -0
  95. data/spec/dummy/tmp/cache/34A/DF0/shard%2F127 +1 -0
  96. data/spec/dummy/tmp/cache/34A/DF0/shard%2F208 +1 -0
  97. data/spec/dummy/tmp/cache/34A/E10/shard%2F145 +1 -0
  98. data/spec/dummy/tmp/cache/34A/E60/shard%2F190 +1 -0
  99. data/spec/dummy/tmp/cache/34B/E30/shard%2F155 +1 -0
  100. data/spec/dummy/tmp/cache/34D/E30/shard%2F139 +0 -0
  101. data/spec/dummy/tmp/cache/34E/E50/shard%2F149 +0 -0
  102. data/spec/dummy/tmp/cache/353/EF0/shard%2F199 +1 -0
  103. data/spec/dummy/tmp/cache/3A4/E90/shard%2F10003 +1 -0
  104. data/spec/dummy/tmp/cache/3A5/ED0/shard%2F10031 +1 -0
  105. data/spec/dummy/tmp/cache/3A9/EF0/shard%2F10017 +1 -0
  106. data/spec/lib/active_record/association_spec.rb +305 -0
  107. data/spec/lib/active_record/attribute_methods_spec.rb +108 -0
  108. data/spec/lib/active_record/base_spec.rb +66 -0
  109. data/spec/lib/active_record/calculations_spec.rb +119 -0
  110. data/spec/lib/active_record/connection_handler_spec.rb +45 -0
  111. data/spec/lib/active_record/connection_pool_spec.rb +23 -0
  112. data/spec/lib/active_record/finder_methods_spec.rb +29 -0
  113. data/spec/lib/active_record/query_cache_spec.rb +20 -0
  114. data/spec/lib/active_record/query_methods_spec.rb +130 -0
  115. data/spec/lib/active_record/relation_spec.rb +38 -0
  116. data/spec/lib/cache_extensions_spec.rb +27 -0
  117. data/spec/lib/connection_pool_proxy_spec.rb +13 -0
  118. data/spec/lib/database_server_spec.rb +154 -0
  119. data/spec/lib/shackles_spec.rb +147 -0
  120. data/spec/models/shard_spec.rb +382 -0
  121. data/spec/spec_helper.rb +32 -0
  122. metadata +344 -0
@@ -0,0 +1,117 @@
1
+ module Switchman
2
+ module ActiveRecord
3
+ module ConnectionPool
4
+ def self.included(klass)
5
+ klass.alias_method_chain(:checkout_new_connection, :sharding)
6
+ klass.send(:remove_method, :connection)
7
+ klass.send(:remove_method, :active_connection?)
8
+ klass.send(:remove_method, :release_connection)
9
+ klass.send(:remove_method, :clear_stale_cached_connections!)
10
+ end
11
+
12
+ attr_writer :shard
13
+
14
+ def shard
15
+ @shard || Shard.default
16
+ end
17
+
18
+ def default_schema
19
+ raise "Not postgres!" unless self.spec.config[:adapter] == 'postgresql'
20
+ connection unless @schemas
21
+ # default shard will not switch databases immediately, so it won't be set yet
22
+ @schemas ||= connection.schemas
23
+ @schemas.first
24
+ end
25
+
26
+ def checkout_new_connection_with_sharding
27
+ # TODO: this might be a threading issue
28
+ spec.config[:shard_name] = self.shard.name
29
+
30
+ conn = checkout_new_connection_without_sharding
31
+ conn.shard = self.shard
32
+ conn
33
+ end
34
+
35
+ def connection
36
+ conns = synchronize { @reserved_connections[current_connection_id] ||= [] }
37
+ conn = conns.find { |conn| !conn.open_transactions || conn.shard == self.shard || conn.adapter_name == 'PostgreSQL' }
38
+ unless conn
39
+ conn = checkout
40
+ yield conn if block_given?
41
+ conns << conn
42
+ end
43
+ switch_database(conn) if conn.shard != self.shard
44
+ conn
45
+ end
46
+
47
+ # Is there an open connection that is being used for the current thread?
48
+ def active_connection?
49
+ synchronize do
50
+ @reserved_connections.fetch(current_connection_id) {
51
+ return false
52
+ }.any? { |conn| conn.in_use? }
53
+ end
54
+ end
55
+
56
+ def release_connection
57
+ conns = synchronize { @reserved_connections.delete(current_connection_id) }
58
+ conns.each { |conn| checkin conn } if conns
59
+ end
60
+
61
+ def clear_stale_cached_connections!
62
+ keys = @reserved_connections.keys - Thread.list.find_all { |t|
63
+ t.alive?
64
+ }.map { |thread| thread.object_id }
65
+ keys.each do |key|
66
+ conns = @reserved_connections[key]
67
+ ActiveSupport::Deprecation.warn(<<-eowarn) if conn.in_use?
68
+ Database connections will not be closed automatically, please close your
69
+ database connection at the end of the thread by calling `close` on your
70
+ connection. For example: ActiveRecord::Base.connection.close
71
+ eowarn
72
+ conns.each { |conn| checkin conn }
73
+ @reserved_connections.delete(key)
74
+ end
75
+ end
76
+
77
+ def release(conn)
78
+ synchronize do
79
+ thread_id = nil
80
+
81
+ if @reserved_connections[current_connection_id].include?(conn)
82
+ thread_id = current_connection_id
83
+ @reserved_connections[thread_id].delete(conn)
84
+ else
85
+ thread_id = @reserved_connections.keys.find { |k|
86
+ @reserved_connections[k].include?(conn)
87
+ }
88
+ if thread_id
89
+ @reserved_connections[thread_id].delete(conn)
90
+ end
91
+ end
92
+
93
+ @reserved_connections.delete thread_id if thread_id && @reserved_connections[thread_id].empty?
94
+ end
95
+ end
96
+
97
+ def switch_database(conn)
98
+ if !@schemas && conn.adapter_name == 'PostgreSQL' && !self.shard.database_server.config[:shard_name]
99
+ @schemas = conn.schemas
100
+ end
101
+
102
+ spec.config[:shard_name] = self.shard.name
103
+ case conn.adapter_name
104
+ when 'MySQL', 'Mysql2'
105
+ conn.execute("USE #{spec.config[:database]}")
106
+ when 'PostgreSQL'
107
+ conn.schema_search_path = spec.config[:schema_search_path]
108
+ when 'SQLite'
109
+ # This is an artifact of the adapter modifying the path to be an absolute path when it is instantiated; just let it slide
110
+ else
111
+ raise("Cannot switch databases on same DatabaseServer with adapter type: #{conn.adapter_name}. Limit one Shard per DatabaseServer.")
112
+ end
113
+ conn.shard = shard
114
+ end
115
+ end
116
+ end
117
+ end
@@ -0,0 +1,25 @@
1
+ module Switchman
2
+ module ActiveRecord
3
+ module FinderMethods
4
+ # find_one uses binds, so we can't depend on QueryMethods
5
+ # catching it
6
+ def find_one(id)
7
+ local_id, shard = Shard.local_id_for(id)
8
+
9
+ return super(local_id) if shard_source_value != :implicit
10
+
11
+ if shard
12
+ begin
13
+ old_shard_value = shard_value
14
+ self.shard_value = shard
15
+ super(local_id)
16
+ ensure
17
+ self.shard_value = old_shard_value
18
+ end
19
+ else
20
+ super
21
+ end
22
+ end
23
+ end
24
+ end
25
+ end
@@ -0,0 +1,43 @@
1
+ module Switchman
2
+ module ActiveRecord
3
+ module LogSubscriber
4
+ def self.included(klass)
5
+ klass.send(:remove_method, :sql)
6
+ end
7
+
8
+ # sadly, have to completely replace this
9
+ def sql(event)
10
+ self.class.runtime += event.duration
11
+ return unless logger.debug?
12
+
13
+ payload = event.payload
14
+
15
+ return if 'SCHEMA' == payload[:name]
16
+
17
+ name = '%s (%.1fms)' % [payload[:name], event.duration]
18
+ sql = payload[:sql].squeeze(' ')
19
+ binds = nil
20
+ connection = ObjectSpace._id2ref(payload[:connection_id])
21
+
22
+ unless (payload[:binds] || []).empty?
23
+ binds = " " + payload[:binds].map { |col,v|
24
+ if col
25
+ [col.name, v]
26
+ else
27
+ [nil, v]
28
+ end
29
+ }.inspect
30
+ end
31
+
32
+ if odd?
33
+ name = color(name, self.class::CYAN, true)
34
+ sql = color(sql, nil, true)
35
+ else
36
+ name = color(name, self.class::MAGENTA, true)
37
+ end
38
+
39
+ debug " #{name} #{sql}#{binds} [shard #{connection.shard.id} #{::Shackles.environment}]"
40
+ end
41
+ end
42
+ end
43
+ end
@@ -0,0 +1,13 @@
1
+ module Switchman
2
+ module ActiveRecord
3
+ module PostgreSQLAdapter
4
+ def self.included(klass)
5
+ klass::NATIVE_DATABASE_TYPES[:primary_key] = "bigserial primary key".freeze
6
+ end
7
+
8
+ def schemas
9
+ select_values("SELECT * FROM unnest(current_schemas(false))")
10
+ end
11
+ end
12
+ end
13
+ end
@@ -0,0 +1,12 @@
1
+ module Switchman
2
+ module ActiveRecord
3
+ # needs to be included in the same class as ::ActiveRecord::ConnectionAdapters::QueryCache
4
+ # *after* that module is included
5
+ module QueryCache
6
+ def cache_sql(sql, *args, &block)
7
+ # have to include the shard id in the cache key because of switching dbs on the same connection
8
+ super("#{self.shard.id}::#{sql}", *args, &block)
9
+ end
10
+ end
11
+ end
12
+ end
@@ -0,0 +1,184 @@
1
+ module Switchman
2
+ module ActiveRecord
3
+ module QueryMethods
4
+ # shard_value is one of:
5
+ # A shard
6
+ # An array or relation of shards
7
+ # An AR object (query runs against that object's associated_shards)
8
+ # shard_source_value is one of:
9
+ # :implicit - inferred from current shard when relation was created, or primary key where clause
10
+ # :explicit - explicit set on the relation
11
+ # :to_a - a special value that Relation#to_a uses when querying multiple shards to
12
+ # remove primary keys from conditions that aren't applicable to the current shard
13
+ attr_accessor :shard_value, :shard_source_value
14
+
15
+ def shard(value, source = :explicit)
16
+ relation = clone
17
+ relation.shard_value = value
18
+ relation.shard_source_value = source
19
+ if (primary_shard != relation.primary_shard || source == :to_a)
20
+ relation.where_values = relation.transpose_predicates(relation.where_values, primary_shard, relation.primary_shard, source == :to_a) if !relation.where_values.empty?
21
+ relation.having_values = relation.transpose_predicates(relation.having_values, primary_shard, relation.primary_shard, source == :to_a) if !relation.having_values.empty?
22
+ end
23
+ relation
24
+ end
25
+
26
+ # replace these with versions that call build_where on the
27
+ # result relation, not the source relation (so build_where
28
+ # is able to implicitly change the shard_value)
29
+ def where(opts, *rest)
30
+ return self if opts.blank?
31
+
32
+ relation = clone
33
+ relation.where_values += relation.build_where(opts, rest)
34
+ relation
35
+ end
36
+
37
+ def having(opts, *rest)
38
+ return self if opts.blank?
39
+
40
+ relation = clone
41
+ relation.having_values += relation.build_where(opts, rest)
42
+ relation
43
+ end
44
+
45
+ def build_where(opts, other = [])
46
+ source_shard = Shard.current(klass.shard_category)
47
+ case opts
48
+ when Hash, Arel::Nodes::Node
49
+ predicates = super
50
+ infer_shards_from_primary_key(predicates) if shard_source_value == :implicit && shard_value.is_a?(Shard)
51
+ predicates = transpose_predicates(predicates, source_shard, primary_shard) if shard_source_value != :explicit
52
+ predicates
53
+ else
54
+ super
55
+ end
56
+ end
57
+
58
+ # the shard that where_values are relative to. if it's multiple shards, they're stored
59
+ # relative to the first shard
60
+ def primary_shard
61
+ case shard_value
62
+ when Shard, DefaultShard
63
+ shard_value
64
+ # associated_shards
65
+ when ::ActiveRecord::Base
66
+ shard_value.shard
67
+ when Array
68
+ shard_value.first
69
+ when ::ActiveRecord::Relation
70
+ Shard.default
71
+ else
72
+ raise ArgumentError("invalid shard value #{shard_value}")
73
+ end
74
+ end
75
+
76
+ private
77
+ def infer_shards_from_primary_key(predicates)
78
+ primary_key = predicates.detect do |predicate|
79
+ predicate.is_a?(Arel::Nodes::Binary) && predicate.left.is_a?(Arel::Attributes::Attribute) &&
80
+ predicate.left.relation.engine == klass && klass.primary_key == predicate.left.name
81
+ end
82
+ if primary_key
83
+ case primary_key.right
84
+ when Array
85
+ id_shards = Set.new
86
+ primary_key.right.each do |value|
87
+ local_id, id_shard = Shard.local_id_for(value)
88
+ id_shard ||= Shard.current(klass.shard_category) if local_id
89
+ id_shards << id_shard if id_shard
90
+ end
91
+ if id_shards.empty?
92
+ return
93
+ elsif id_shards.length == 1
94
+ id_shard = id_shards.first
95
+ # prefer to not change the shard
96
+ elsif id_shards.include?(primary_shard)
97
+ id_shards.delete(primary_shard)
98
+ self.shard_value = [primary_shard] + id_shards.to_a
99
+ return
100
+ else
101
+ id_shards = id_shards.to_a
102
+ self.where_values = transpose_predicates(where_values, primary_shard, id_shards.first) if !where_values.empty?
103
+ self.having_values = transpose_predicates(having_values, primary_shard, id_shards.first) if !having_values.empty?
104
+ self.shard_value = id_shards
105
+ return
106
+ end
107
+ else
108
+ local_id, id_shard = Shard.local_id_for(primary_key.right)
109
+ id_shard ||= Shard.current(klass.shard_category) if local_id
110
+ end
111
+ return if !id_shard || id_shard == primary_shard
112
+ self.where_values = transpose_predicates(where_values, primary_shard, id_shard) if !where_values.empty?
113
+ self.having_values = transpose_predicates(having_values, primary_shard, id_shard) if !having_values.empty?
114
+ self.shard_value = id_shard
115
+ end
116
+ end
117
+
118
+ def transposable_attribute_type(attribute)
119
+ return false unless attribute.is_a?(Arel::Attributes::Attribute)
120
+ if sharded_primary_key?(attribute)
121
+ return :primary
122
+ elsif sharded_foreign_key?(attribute)
123
+ return :foreign
124
+ end
125
+ end
126
+
127
+ def sharded_foreign_key?(attribute)
128
+ @@foreign_keys ||= {}
129
+ @@foreign_keys[attribute.relation.table_name] ||= {}
130
+ if @@foreign_keys[attribute.relation.table_name].has_key?(attribute.name)
131
+ @@foreign_keys[attribute.relation.table_name][attribute.name]
132
+ else
133
+ models = attribute.relation.engine.descendants.select{|d| d.table_name == attribute.relation.table_name}
134
+ models << attribute.relation.engine unless attribute.relation.engine == ::ActiveRecord::Base
135
+
136
+ @@foreign_keys[attribute.relation.table_name][attribute.name] = models.any?{|m| m.sharded_column?(attribute.name)}
137
+ end
138
+ end
139
+
140
+ def sharded_primary_key?(attribute)
141
+ attribute.relation.engine.primary_key == attribute.name
142
+ end
143
+
144
+ # semi-private
145
+ public
146
+ def transpose_predicates(predicates, source_shard, target_shard, remove_nonlocal_primary_keys = false)
147
+ predicates.map do |predicate|
148
+ next predicate unless predicate.is_a?(Arel::Nodes::Binary) && type = transposable_attribute_type(predicate.left)
149
+
150
+ remove = true if type == :primary && remove_nonlocal_primary_keys && predicate.left.relation.engine == klass
151
+
152
+ new_right_value = case predicate.right
153
+ when Array
154
+ local_ids = []
155
+ predicate.right.each do |value|
156
+ local_id = Shard.relative_id_for(value, source_shard, target_shard)
157
+ local_ids << local_id unless remove && local_id > Shard::IDS_PER_SHARD
158
+ end
159
+ local_ids
160
+ when Arel::Nodes::BindParam
161
+ # look for a bind param with a matching column name
162
+ if @bind_params && idx = @bind_params.find_index{|b| b.is_a?(Array) && b.first.try(:name) == predicate.left}
163
+ column, value = @bind_params[idx]
164
+ local_id = Shard.relative_id_for(value, source_shard, target_shard)
165
+ local_id = [] if remove && local_id > Shard::IDS_PER_SHARD
166
+ @bind_params[idx] = [column, local_id]
167
+ end
168
+ predicate.right
169
+ else
170
+ local_id = Shard.relative_id_for(predicate.right, source_shard, target_shard)
171
+ local_id = [] if remove && local_id > Shard::IDS_PER_SHARD
172
+ local_id
173
+ end
174
+
175
+ if new_right_value == predicate.right
176
+ predicate
177
+ else
178
+ predicate.class.new(predicate.left, new_right_value)
179
+ end
180
+ end
181
+ end
182
+ end
183
+ end
184
+ end
@@ -0,0 +1,69 @@
1
+ module Switchman
2
+ module ActiveRecord
3
+ module Relation
4
+ def self.included(klass)
5
+ klass::SINGLE_VALUE_METHODS.concat [ :shard, :shard_source ]
6
+ %w{initialize exec_queries update_all delete_all}.each do |method|
7
+ klass.alias_method_chain(method, :sharding)
8
+ end
9
+ end
10
+
11
+ def initialize_with_sharding(klass, table)
12
+ initialize_without_sharding(klass, table)
13
+ self.shard_value = Shard.current(klass.shard_category)
14
+ self.shard_source_value = :implicit
15
+ end
16
+
17
+ def merge(*args)
18
+ relation = super
19
+ if relation.shard_value != self.shard_value && relation.shard_source_value == :implicit
20
+ relation.shard_value = self.shard_value
21
+ relation.shard_source_value = self.shard_source_value
22
+ end
23
+ relation
24
+ end
25
+
26
+ def exec_queries_with_sharding
27
+ return @records if loaded?
28
+ results = self.activate{|relation| relation.send(:exec_queries_without_sharding) }
29
+ case shard_value
30
+ when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
31
+ @records = results
32
+ @loaded = true
33
+ end
34
+ results
35
+ end
36
+
37
+ %w{update_all delete_all}.each do |method|
38
+ class_eval <<-RUBY
39
+ def #{method}_with_sharding(*args)
40
+ self.activate{|relation| relation.#{method}_without_sharding(*args)}
41
+ end
42
+ RUBY
43
+ end
44
+
45
+ def activate(&block)
46
+ case shard_value
47
+ when DefaultShard, Shard.current(klass.shard_category)
48
+ yield(self, shard_value)
49
+ when Shard
50
+ shard_value.activate(klass.shard_category) { yield(self, shard_value) }
51
+ when Array, ::ActiveRecord::Relation, ::ActiveRecord::Base
52
+ # TODO: implement local limit to avoid querying extra shards
53
+ if shard_value.is_a?(::ActiveRecord::Base)
54
+ if shard_value.respond_to?(:associated_shards)
55
+ shards = shard_value.associated_shards
56
+ else
57
+ shards = [shard_value.shard]
58
+ end
59
+ else
60
+ shards = shard_value
61
+ end
62
+ Shard.with_each_shard(shards, [klass.shard_category]) do
63
+ shard(Shard.current(klass.shard_category), :to_a).activate(&block)
64
+ end
65
+ end
66
+ end
67
+ end
68
+ end
69
+ end