switchman 0.0.1

Sign up to get free protection for your applications and to get access to all the features.
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,12 @@
1
+ module Switchman::CacheExtensions
2
+ module ClassMethods
3
+ def cache_with_sharding
4
+ Switchman::Shard.current.database_server.cache_store
5
+ end
6
+ end
7
+
8
+ def self.included(klass)
9
+ klass.extend(ClassMethods)
10
+ klass.singleton_class.alias_method_chain :cache, :sharding
11
+ end
12
+ end
@@ -0,0 +1,62 @@
1
+ module Switchman
2
+ class ConnectionPoolProxy
3
+ delegate :spec, :connected?, :default_schema, :with_connection,
4
+ :to => :current_pool
5
+
6
+ attr_reader :category
7
+
8
+ def default_pool
9
+ @default_pool
10
+ end
11
+
12
+ def initialize(category, default_pool, shard_connection_pools)
13
+ @category = category
14
+ @default_pool = default_pool
15
+ @connection_pools = shard_connection_pools
16
+ end
17
+
18
+ def active_shard
19
+ Shard.current(@category)
20
+ end
21
+
22
+ def current_pool
23
+ pool = self.default_pool if active_shard.default?
24
+ pool = @connection_pools[pool_key] ||= create_pool unless pool
25
+ pool.shard = active_shard
26
+ pool
27
+ end
28
+
29
+ def connection
30
+ pool = current_pool
31
+ pool.connection
32
+ end
33
+
34
+ %w{release_connection disconnect! clear_reloadable_connections! verify_active_connections! clear_stale_cached_connections!}.each do |method|
35
+ class_eval(<<-EOS)
36
+ def #{method}
37
+ @connection_pools.values.each(&:#{method})
38
+ end
39
+ EOS
40
+ end
41
+
42
+ protected
43
+
44
+ def pool_key
45
+ active_shard.database_server.shareable? ? active_shard.database_server.id : active_shard
46
+ end
47
+
48
+ def create_pool
49
+ shard = active_shard
50
+ config = shard.database_server.config.dup
51
+ spec = ::ActiveRecord::Base::ConnectionSpecification.new(config, "#{config[:adapter]}_connection")
52
+ # unfortunately the AR code that does this require logic can't really be
53
+ # called in isolation
54
+ require "active_record/connection_adapters/#{config[:adapter]}_adapter"
55
+
56
+ ::ActiveRecord::ConnectionAdapters::ConnectionPool.new(spec).tap do |pool|
57
+ pool.shard = shard
58
+ end
59
+ end
60
+ end
61
+ end
62
+
@@ -0,0 +1,197 @@
1
+ module Switchman
2
+ class DatabaseServer
3
+ attr_accessor :id, :config
4
+
5
+ class << self
6
+ def all
7
+ database_servers.values
8
+ end
9
+
10
+ def find(id_or_all)
11
+ return self.all if id_or_all == :all
12
+ return id_or_all.map { |id| self.database_servers[id || Rails.env] }.compact.uniq if id_or_all.is_a?(Array)
13
+ database_servers[id_or_all || Rails.env]
14
+ end
15
+
16
+ def create(settings = {})
17
+ raise "database servers should be set up in database.yml" unless Rails.env.test?
18
+ id = 1
19
+ while database_servers[id.to_s]; id += 1; end
20
+ server = DatabaseServer.new({ :id => id.to_s }.merge(settings))
21
+ server.instance_variable_set(:@fake, true)
22
+ raise "database server #{server.id} already exists" if database_servers[server.id]
23
+ database_servers[server.id] = server
24
+ end
25
+
26
+ def server_for_new_shard
27
+ servers = all.select { |s| s.config[:open] }
28
+ return find(nil) if servers.empty?
29
+ servers[rand(servers.length)]
30
+ end
31
+
32
+ private
33
+ def database_servers
34
+ unless @database_servers
35
+ @database_servers = {}.with_indifferent_access
36
+ ::ActiveRecord::Base.configurations.each do |(id, config)|
37
+ @database_servers[id] = DatabaseServer.new(:id => id, :config => config)
38
+ end
39
+ end
40
+ @database_servers
41
+ end
42
+ end
43
+
44
+ def initialize(settings = {})
45
+ self.id = settings[:id]
46
+ self.config = (settings[:config] || {}).symbolize_keys
47
+ end
48
+
49
+ def destroy
50
+ raise "database servers should be set up in database.yml" unless Rails.env.test?
51
+ self.class.send(:database_servers).delete(self.id) if self.id
52
+ end
53
+
54
+ def fake?
55
+ @fake
56
+ end
57
+
58
+ def shareable?
59
+ @shareable_environment_key ||= []
60
+ environment = ::Shackles.environment
61
+ explicit_user = ::Shackles.global_config[:username]
62
+ return @shareable if @shareable_environment_key == [environment, explicit_user]
63
+ @shareable_environment_key = [environment, explicit_user]
64
+ username = self.config[:username]
65
+ username = self.config[environment][:username] if environment && self.config[environment] && self.config[environment][:username]
66
+ username = explicit_user if explicit_user
67
+ @shareable = self.config[:adapter] != 'sqlite3' && username !~ /%?\{[a-zA-Z0-9_]+\}/
68
+ end
69
+
70
+ def shards
71
+ if self.id == Rails.env
72
+ Shard.where("database_server_id IS NULL OR database_server_id=?", self.id)
73
+ else
74
+ Shard.where(:database_server_id => self.id)
75
+ end
76
+ end
77
+
78
+ def create_new_shard(options = {})
79
+ db_name = options[:db_name]
80
+ create_schema = options[:schema]
81
+ # look for another shard associated with this db
82
+ other_shard = self.shards.where("name<>':memory:' OR name IS NULL").order(:id).first
83
+ temp_db_name = other_shard.try(:name) unless id == Rails.env
84
+ temp_db_name = Shard.default.name if id == Rails.env
85
+
86
+ case config[:adapter]
87
+ when 'postgresql'
88
+ temp_db_name ||= 'public'
89
+ create_statement = lambda { "CREATE SCHEMA #{db_name}" }
90
+ password = " PASSWORD #{::ActiveRecord::Base.connection.quote(config[:password])}" if config[:password]
91
+ when 'sqlite3'
92
+ if db_name
93
+ # Try to create a db on-disk even if the only shards for sqlite are in-memory
94
+ temp_db_name = nil if temp_db_name == ':memory:'
95
+ # Put it in the db directory if there are no other sqlite shards
96
+ temp_db_name ||= 'db/dummy'
97
+ temp_db_name = File.join(File.dirname(temp_db_name), "#{db_name}.sqlite3")
98
+ # If they really asked for :memory:, give them :memory:
99
+ temp_db_name = db_name if db_name == ':memory:'
100
+ db_name = temp_db_name
101
+ end
102
+ else
103
+ temp_db_name ||= self.config[:database] % self.config
104
+ create_statement = lambda { "CREATE DATABASE #{db_name}" }
105
+ end
106
+ sharding_config = Switchman.config
107
+ config_create_statement = sharding_config[config[:adapter]].try(:[], :create_statement)
108
+ config_create_statement ||= sharding_config[:create_statement]
109
+ if config_create_statement
110
+ create_commands = Array(config_create_statement).dup
111
+ create_statement = lambda {
112
+ create_commands.map { |statement| statement.gsub('%{db_name}', db_name).gsub('%{password}', password || '') }
113
+ }
114
+ end
115
+
116
+ shard = Shard.create!(:name => temp_db_name,
117
+ :database_server => self) do |shard|
118
+ shard.id = options[:id] if options[:id]
119
+ end
120
+ begin
121
+ if db_name.nil?
122
+ base_db_name = self.config[:database] % self.config
123
+ base_db_name = $1 if base_db_name =~ /(?:.*\/)(.+)_shard_\d+(?:\.sqlite3)?$/
124
+ base_db_name = nil if base_db_name == ':memory:'
125
+ base_db_name << '_' if base_db_name
126
+ db_name = "#{base_db_name}shard_#{shard.id}"
127
+ if config[:adapter] == 'sqlite3'
128
+ # Try to create a db on-disk even if the only shards for sqlite are in-memory
129
+ temp_db_name = nil if temp_db_name == ':memory:'
130
+ # Put it in the db directory if there are no other sqlite shards
131
+ temp_db_name ||= 'db/dummy'
132
+ db_name = File.join(File.dirname(temp_db_name), "#{db_name}.sqlite3")
133
+ shard.name = db_name
134
+ end
135
+ end
136
+ shard.activate(Shard.categories) do
137
+ ::Shackles.activate(:deploy) do
138
+ begin
139
+ if create_statement
140
+ Array(create_statement.call).each do |stmt|
141
+ ::ActiveRecord::Base.connection.execute(stmt)
142
+ end
143
+ # have to disconnect and reconnect to the correct db
144
+ if self.shareable? && other_shard
145
+ other_shard.activate { ::ActiveRecord::Base.connection }
146
+ else
147
+ ::ActiveRecord::Base.connection_pool.current_pool.disconnect!
148
+ end
149
+ end
150
+ shard.name = db_name
151
+ old_proc = ::ActiveRecord::Base.connection.raw_connection.set_notice_processor {} if config[:adapter] == 'postgresql'
152
+ old_verbose = ::ActiveRecord::Migration.verbose
153
+ ::ActiveRecord::Migration.verbose = false
154
+ ::ActiveRecord::Migrator.migrate(Rails.root + "db/migrate/") unless create_schema == false
155
+ ensure
156
+ ::ActiveRecord::Migration.verbose = old_verbose
157
+ ::ActiveRecord::Base.connection.raw_connection.set_notice_processor(&old_proc) if old_proc
158
+ end
159
+ end
160
+ end
161
+ shard.save!
162
+ shard
163
+ rescue
164
+ shard.destroy
165
+ raise
166
+ end
167
+ end
168
+
169
+ def cache_store
170
+ if !@cache_store
171
+ @cache_store = Rails.cache_without_sharding if self.id == Rails.env
172
+ # TODO: get the config from somewhere
173
+ config = nil
174
+ @cache_store ||= ActiveSupport::Cache.lookup_store(config) if config
175
+ @cache_store ||= Rails.cache_without_sharding
176
+ @cache_store.options[:namespace] = lambda { Shard.current.default? ? nil : "shard_#{Shard.current.id}" }
177
+ end
178
+ @cache_store
179
+ end
180
+
181
+ def shard_name(shard)
182
+ if config[:shard_name]
183
+ config[:shard_name]
184
+ elsif config[:adapter] == 'postgresql'
185
+ if shard == :bootstrap
186
+ # rescue nil because the database may not exist yet; if it doesn't,
187
+ # it will shortly, and this will be re-invoked
188
+ ::ActiveRecord::Base.connection.schemas.first rescue nil
189
+ else
190
+ shard.activate { ::ActiveRecord::Base.connection_pool.default_schema }
191
+ end
192
+ else
193
+ config[:database]
194
+ end
195
+ end
196
+ end
197
+ end
@@ -0,0 +1,28 @@
1
+ require_dependency 'switchman/database_server'
2
+
3
+ class Switchman::DefaultShard
4
+ def id; 'default'; end
5
+ def activate(*categories); yield; end
6
+ def activate!(*categories); end
7
+ def default?; true; end
8
+ def relative_id_for(local_id, target = nil); local_id; end
9
+ def global_id_for(local_id); local_id; end
10
+ def database_server_id; nil; end
11
+ def database_server; ::Switchman::DatabaseServer.find(nil); end
12
+ def name
13
+ unless instance_variable_defined?(:@name)
14
+ @name = nil # prevent taking this branch on recursion
15
+ @name = database_server.shard_name(:bootstrap)
16
+ end
17
+ @name
18
+ end
19
+ def description; Rails.env; end
20
+ # The default's shard is always the default shard
21
+ def shard; self; end
22
+ def _dump(depth)
23
+ ''
24
+ end
25
+ def self._load(str)
26
+ Shard.default
27
+ end
28
+ end
@@ -0,0 +1,91 @@
1
+ module Switchman
2
+ class Engine < ::Rails::Engine
3
+ isolate_namespace Switchman
4
+
5
+ initializer 'switchman.extend_ar', :before => "active_record.initialize_database" do
6
+ ActiveSupport.on_load(:active_record) do
7
+ #require 'active_record/associations/preloader/belongs_to'
8
+
9
+ require "switchman/active_record/abstract_adapter"
10
+ require "switchman/active_record/association"
11
+ require "switchman/active_record/attribute_methods"
12
+ require "switchman/active_record/base"
13
+ require "switchman/active_record/calculations"
14
+ require "switchman/active_record/connection_handler"
15
+ require "switchman/active_record/connection_pool"
16
+ require "switchman/active_record/finder_methods"
17
+ require "switchman/active_record/log_subscriber"
18
+ require "switchman/active_record/query_cache"
19
+ require "switchman/active_record/query_methods"
20
+ require "switchman/active_record/relation"
21
+ require "switchman/cache_extensions"
22
+
23
+ include ActiveRecord::Base
24
+ include ActiveRecord::AttributeMethods
25
+ ::ActiveRecord::Associations::Association.send(:include, ActiveRecord::Association)
26
+ ::ActiveRecord::Associations::CollectionProxy.send(:include, ActiveRecord::CollectionProxy)
27
+ ::ActiveRecord::Associations::Builder::Association.send(:include, ActiveRecord::Builder::Association)
28
+
29
+ ::ActiveRecord::Associations::Preloader::Association.send(:include, ActiveRecord::Preloader::Association)
30
+ ::ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, ActiveRecord::AbstractAdapter)
31
+ ::ActiveRecord::ConnectionAdapters::ConnectionHandler.send(:include, ActiveRecord::ConnectionHandler)
32
+ ::ActiveRecord::ConnectionAdapters::ConnectionPool.send(:include, ActiveRecord::ConnectionPool)
33
+ ::ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, ActiveRecord::QueryCache)
34
+ ::ActiveRecord::LogSubscriber.send(:include, ActiveRecord::LogSubscriber)
35
+ ::ActiveRecord::Relation.send(:include, ActiveRecord::Calculations)
36
+ ::ActiveRecord::Relation.send(:include, ActiveRecord::FinderMethods)
37
+ ::ActiveRecord::Relation.send(:include, ActiveRecord::QueryMethods)
38
+ ::ActiveRecord::Relation.send(:include, ActiveRecord::Relation)
39
+ Rails.send(:include, CacheExtensions)
40
+ end
41
+ end
42
+
43
+ def self.foreign_key_check(name, type, options)
44
+ if name.to_s =~ /_id\z/ && type.to_s == 'integer' && options[:limit].to_i < 8
45
+ puts "WARNING: All foreign keys need to be 8-byte integers. #{name} looks like a foreign key. If so, please add the option: `:limit => 8`"
46
+ end
47
+ end
48
+
49
+ initializer 'switchman.extend_connection_adapters', :after => "active_record.initialize_database" do
50
+ ActiveSupport.on_load(:active_record) do
51
+ ::ActiveRecord::ConnectionAdapters::AbstractAdapter.descendants.each do |klass|
52
+ klass.class_eval do
53
+ def add_column_with_foreign_key_check(table, name, type, options = {})
54
+ Switchman::Engine.foreign_key_check(name, type, options)
55
+ add_column_without_foreign_key_check(table, name, type, options)
56
+ end
57
+ alias_method_chain(:add_column, :foreign_key_check)
58
+ end
59
+ end
60
+
61
+ ::ActiveRecord::ConnectionAdapters::TableDefinition.class_eval do
62
+ def column_with_foreign_key_check(name, type, options = {})
63
+ Switchman::Engine.foreign_key_check(name, type, options)
64
+ column_without_foreign_key_check(name, type, options)
65
+ end
66
+ alias_method_chain(:column, :foreign_key_check)
67
+ end
68
+
69
+ if defined?(::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter)
70
+ require "switchman/active_record/postgresql_adapter"
71
+ ::ActiveRecord::ConnectionAdapters::PostgreSQLAdapter.send(:include, ActiveRecord::PostgreSQLAdapter)
72
+ end
73
+ end
74
+ end
75
+
76
+ initializer 'switchman.eager_load' do
77
+ ActiveSupport.on_load(:before_eager_load) do
78
+ # This needs to be loaded before Switchman::Shard, otherwise it won't autoload it correctly
79
+ require_dependency('active_record/base')
80
+ end
81
+ end
82
+
83
+ initializer 'switchman.extend_shackles', :after => "shackles.extend_ar" do
84
+ ActiveSupport.on_load(:active_record) do
85
+ require "switchman/shackles"
86
+
87
+ ::Shackles.send(:include, Shackles)
88
+ end
89
+ end
90
+ end
91
+ end
@@ -0,0 +1,124 @@
1
+ require_dependency "switchman/test_helper"
2
+
3
+ module Switchman
4
+ # including this module in your specs will give you several shards to
5
+ # work with during specs:
6
+ # * Shard.default - the test database itself
7
+ # * @shard1 - a shard possibly using the same connection as Shard.default
8
+ # * @shard2 - a shard using a dedicated connection
9
+ # * @shard3 - a shard using the same connection as @shard1 (this might
10
+ # be Shard.default if they already share a connection, or
11
+ # a separate shard)
12
+ module RSpecHelper
13
+ @@keep_the_shards = false
14
+ @@shard1 = nil
15
+
16
+ def self.included(klass)
17
+ # our before handlers have already been configured from a parent group; don't add them again
18
+ return if klass.parent_groups[1..-1].any? { |group| group.included_modules.include?(self) }
19
+
20
+ ::RSpec.configure do |config|
21
+ block = proc do
22
+ if @@shard1
23
+ # some specs are mean, and blow away our shards
24
+ begin
25
+ @@shard1.reload
26
+ @@shard2.reload
27
+ @@shard3.reload if @@shard3
28
+ rescue ::ActiveRecord::RecordNotFound
29
+ Shard.default.clone.save!
30
+ Shard.default(true)
31
+ @@shard1 = @@shard1.clone
32
+ @@shard1.save!
33
+ @@shard2 = @@shard2.clone
34
+ @@shard2.save!
35
+ if @@shard3
36
+ @@shard3 = @@shard3.clone
37
+ @@shard3.save!
38
+ end
39
+ end
40
+ end
41
+ end
42
+ # this module will be included multiple times, but we don't want to run this global hook multiple times
43
+ if config.hooks[:before][:all].all? { |hook| hook.block.source_location != block.source_location }
44
+ config.before(:all, &block)
45
+ end
46
+ end
47
+
48
+ klass.before(:all) do
49
+ unless @@shard1
50
+ puts "Setting up sharding for all specs..."
51
+ @@shard1, @@shard2 = TestHelper.recreate_persistent_test_shards
52
+ if @@shard1.is_a?(Shard)
53
+ @@keep_the_shards = true
54
+ @@shard3 = nil
55
+ else # @@shard1.is_a?(DatabaseServer)
56
+ @@shard1 = @@shard1.create_new_shard
57
+ @@shard2 = @@shard2.create_new_shard
58
+ if @@shard1.database_server == Shard.default.database_server
59
+ @@shard3 = nil
60
+ else
61
+ @@shard3 = @@shard1.database_server.create_new_shard
62
+ end
63
+ end
64
+ puts "Done!"
65
+
66
+ at_exit do
67
+ # preserve rspec's exit status
68
+ status= $!.is_a?(::SystemExit) ? $!.status : nil
69
+ puts "Tearing down sharding for all specs"
70
+ @@shard1.database_server.destroy unless @@shard1.database_server == Shard.default.database_server
71
+ unless @@keep_the_shards
72
+ @@shard1.drop_database
73
+ @@shard1.destroy
74
+ @@shard2.drop_database
75
+ @@shard2.destroy
76
+ if @@shard3
77
+ @@shard3.drop_database
78
+ @@shard3.destroy
79
+ end
80
+ end
81
+ @@shard2.database_server.destroy
82
+ exit status if status
83
+ end
84
+ end
85
+ @shard1, @shard2 = @@shard1, @@shard2
86
+ @shard3 = @@shard3 ? @@shard3 : Shard.default
87
+ end
88
+
89
+ klass.before do
90
+ Shard.clear_cache
91
+ if use_transactional_fixtures
92
+ Shard.default(true)
93
+ @shard1 = Shard.find(@shard1)
94
+ @shard2 = Shard.find(@shard2)
95
+ shards = [@shard2]
96
+ shards << @shard1 unless @shard1.database_server == Shard.default.database_server
97
+ shards.each do |shard|
98
+ shard.activate do
99
+ # this is how AR does it in fixtures.rb
100
+ ::ActiveRecord::Base.connection.increment_open_transactions
101
+ ::ActiveRecord::Base.connection.transaction_joinable = false
102
+ ::ActiveRecord::Base.connection.begin_db_transaction
103
+ end
104
+ end
105
+ end
106
+ end
107
+
108
+ klass.after do
109
+ if use_transactional_fixtures
110
+ shards = [@shard2]
111
+ shards << @shard1 unless @shard1.database_server == Shard.default.database_server
112
+ shards.each do |shard|
113
+ shard.activate do
114
+ if ::ActiveRecord::Base.connection.open_transactions != 0
115
+ ::ActiveRecord::Base.connection.rollback_db_transaction
116
+ ::ActiveRecord::Base.connection.decrement_open_transactions
117
+ end
118
+ end
119
+ end
120
+ end
121
+ end
122
+ end
123
+ end
124
+ end