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,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