active_record_shards 2.0.0.beta2 → 2.0.0.beta5

data/README.md

@@ -0,0 +1,98 @@
+# ActiveRecord Shards
+
+ActiveRecord Shards is an extension for ActiveRecord that provides support for sharded database and slaves. Basically it is just a nice way to
+switch between database connection. We've made the implementation very small, and have tried not to reinvent any wheels already present in ActiveRecord.
+
+ActiveRecord Shards has only been used and tested on Rails 2.3.x, and has in some form or another been used on in production on a large rails app for
+more than a year.
+
+## Installation
+
+    $ gem install active_record_shards
+
+and make sure to require 'active\_record\_shards' in some way.
+
+## Configuration
+
+Add the slave and shard configuration to config/database.yml:
+
+    production:
+      adapter: mysql
+      encoding: utf8
+      database: my_app_main
+      pool: 5
+      host: db1
+      username: root
+      password:
+      slave:
+        host: db1_slave
+      shards:
+        1:
+          host: db_shard1
+          database: my_app_shard
+          slave:
+            host: db_shard1_slave
+        2:
+          host: db_shard2
+          database: my_app_shard
+          slave:
+            host: db_shard2_slave
+
+basically connections inherit configuration from the parent the configuration file.
+
+## Usage
+
+Normally you have some models that live on a shared database, and you might need to query this data in order to know what shard to switch to.
+All the model that live on the shared database must be marked as not\_sharded:
+
+    class Account < ActiveRecord::Base
+      not_sharded
+      
+      has_many :projects
+    end
+    
+    class Project < ActiveRecord::Base
+      belongs_to :account
+    end
+
+So in this setup the accounts live on the shared database, but the projects are sharded. If accounts had a shard\_id column, you could lookup the account
+in a rack middleware and switch to the right shard:
+
+    class AccountMiddleware
+      def initialize(app)
+        @app = app
+      end
+      
+      def call(env)
+        account = lookup_account(env)
+        
+        if account
+          ActiveRecord::Base.on_shard(account.shard_id) do
+            @app.call(env)
+          end
+        else
+          @app.call(env)
+        end
+      end
+      
+      def lookup_account(env)
+        ...
+      end
+    end
+
+Any where in your app, you can switch to the slave databases, by wrapping you code an on\_slave block:
+
+    ActiveRecord::Base.on_slave do
+      Account.find_by_big_expensive_query
+    end
+
+This will perform the query on the slave, and mark the returned instances as read only. There is also a shortcut for this:
+
+    Account.on_slave.find_by_big_expensive_query
+
+## Copyright
+
+Copyright (c) 2011 Zendesk. See LICENSE for details.
+
+## Authors
+Mick Staugaard, Eric Chapweske

data/lib/active_record_shards/configuration_parser.rb

@@ -6,6 +6,7 @@ module ActiveRecordShards
       conf.keys.each do |env_name|
         env_config = conf[env_name]
         if shards = env_config.delete('shards')
+          env_config['shard_names'] = shards.keys
           shards.each do |shard_name, shard_conf|
             expand_child!(env_config, shard_conf)
             conf["#{env_name}_shard_#{shard_name}"] = shard_conf

data/lib/active_record_shards/connection_switcher.rb

@@ -1,12 +1,30 @@
 module ActiveRecordShards
   module ConnectionSwitcher
+    def default_shard=(new_default_shard)
+      ActiveRecordShards::ShardSelection.default_shard = new_default_shard
+      switch_connection(:shard => new_default_shard)
+    end
 
     def on_shard(shard, &block)
-      old_shard = current_shard_selection.shard
-      switch_connection(:shard => shard)
+      old_options = current_shard_selection.options
+      switch_connection(:shard => shard) if supports_sharding?
       yield
     ensure
-      switch_connection(:shard => old_shard)
+      switch_connection(old_options)
+    end
+
+    def on_all_shards(&block)
+      old_options = current_shard_selection.options
+      if supports_sharding?
+        shard_names.each do |shard|
+          switch_connection(:shard => shard)
+          yield
+        end
+      else
+        yield
+      end
+    ensure
+      switch_connection(old_options)
     end
 
     def on_slave_if(condition, &block)
@@ -40,28 +58,26 @@ module ActiveRecordShards
     alias_method :with_slave_unless, :on_slave_unless
 
     def on_slave_block(&block)
-      old_slave = current_shard_selection.on_slave?
-      begin
-        switch_connection(:slave => true)
-      rescue ActiveRecord::AdapterNotSpecified => e
-        switch_connection(:slave => old_slave)
-        logger.warn("Failed to establish shard connection: #{e.message} - defaulting to master")
-      end
+      old_options = current_shard_selection.options
+      switch_connection(:slave => true)
       with_scope({:find => {:readonly => current_shard_selection.on_slave?}}, :merge, &block)
     ensure
-      switch_connection(:slave => old_slave)
+      switch_connection(old_options)
     end
 
     # Name of the connection pool. Used by ConnectionHandler to retrieve the current connection pool.
     def connection_pool_name # :nodoc:
-      current_shard_selection.shard_name(self)
-#      if current_shard_selection.any?
-#        current_shard_selection.shard_name(name, self)
-#      elsif self == ActiveRecord::Base
-#        name
-#      else
-#        superclass.connection_pool_name
-#      end
+      name = current_shard_selection.shard_name(self)
+
+      if configurations[name].nil? && current_shard_selection.on_slave?
+        current_shard_selection.shard_name(self, false)
+      else
+        name
+      end
+    end
+
+    def supports_sharding?
+      shard_names.any?
     end
 
     private
@@ -84,9 +100,18 @@ module ActiveRecordShards
       end
     end
 
+    def shard_names
+      configurations[shard_env]['shard_names'] || []
+    end
+
+    def shard_env
+      @shard_env = defined?(Rails.env) ? Rails.env : RAILS_ENV
+    end
+
     def establish_shard_connection
-      name = current_shard_selection.shard_name(self)
+      name = connection_pool_name
       spec = configurations[name]
+
       raise ActiveRecord::AdapterNotSpecified.new("No database defined by #{name} in database.yml") if spec.nil?
 
       connection_handler.establish_connection(connection_pool_name, ActiveRecord::Base::ConnectionSpecification.new(spec, "#{spec['adapter']}_connection"))

data/lib/active_record_shards/shard_selection.rb

@@ -1,15 +1,24 @@
 module ActiveRecordShards
   class ShardSelection
+    NO_SHARD = :_no_shard
+    cattr_accessor :default_shard
+
     def initialize
       @on_slave = false
     end
 
-    def shard
-      @shard
+    def shard(klass = nil)
+      if (@shard || self.class.default_shard) && (klass.nil? || klass.is_sharded?)
+        if @shard == NO_SHARD
+          nil
+        else
+          (@shard || self.class.default_shard).to_s
+        end
+      end
     end
 
     def shard=(new_shard)
-      @shard = new_shard
+      @shard = (new_shard || NO_SHARD)
     end
 
     def on_slave?
@@ -20,22 +29,20 @@ module ActiveRecordShards
       @on_slave = (new_slave == true)
     end
 
-    def connection_configuration_name
-      shard_name#(RAILS_ENV)
-    end
-
-    def shard_name(klass = nil)
+    def shard_name(klass = nil, try_slave = true)
       s = "#{RAILS_ENV}"
-      if @shard && (klass.nil? || klass.is_sharded?)
+      if the_shard = shard(klass)
         s << '_shard_'
-        s << @shard.to_s
+        s << the_shard
+      end
+      if @on_slave && try_slave
+        s << "_slave" if @on_slave
       end
-      s << "_slave" if @on_slave
       s
     end
 
-    def any?
-      @on_slave || @shard.present?
+    def options
+      {:shard => @shard, :slave => @on_slave}
     end
   end
 end

data/lib/active_record_shards/tasks.rb

@@ -0,0 +1,52 @@
+require 'active_record_shards'
+
+Rake::TaskManager.class_eval do
+  def remove_task(task_name)
+    @tasks.delete(task_name.to_s)
+  end
+end
+
+def remove_task(task_name)
+  Rake.application.remove_task(task_name)
+end
+
+remove_task 'db:drop'
+remove_task 'db:create'
+remove_task 'db:abort_if_pending_migrations'
+
+namespace :db do
+  desc 'Drops the database for the current RAILS_ENV including shards and slaves'
+  task :drop => :load_config do
+    env_name = defined?(Rails.env) ? Rails.env : RAILS_ENV || 'development'
+    ActiveRecord::Base.configurations.each do |key, conf|
+      if key.starts_with?(env_name) && !key.ends_with?("_slave")
+        drop_database(conf)
+      end
+    end
+  end
+
+  desc "Create the database defined in config/database.yml for the current RAILS_ENV including shards and slaves"
+  task :create => :load_config do
+    env_name = defined?(Rails.env) ? Rails.env : RAILS_ENV || 'development'
+    ActiveRecord::Base.configurations.each do |key, conf|
+      if key.starts_with?(env_name) && !key.ends_with?("_slave")
+        create_database(conf)
+      end
+    end
+  end
+
+  desc "Raises an error if there are pending migrations"
+  task :abort_if_pending_migrations => :environment do
+    if defined? ActiveRecord
+      pending_migrations = ActiveRecord::Base.on_shard(nil) { ActiveRecord::Migrator.new(:up, 'db/migrate').pending_migrations }
+
+      if pending_migrations.any?
+        puts "You have #{pending_migrations.size} pending migrations:"
+        pending_migrations.each do |pending_migration|
+          puts '  %4d %s' % [pending_migration.version, pending_migration.name]
+        end
+        abort %{Run "rake db:migrate" to update your database then try again.}
+      end
+    end
+  end
+end

data/test/configuration_parser_test.rb

@@ -10,13 +10,14 @@ class ConfigurationParserTest < ActiveSupport::TestCase
       setup { @conf = @exploded_conf['test_slave'] }
       should "be exploded" do
         assert_equal({
-          "adapter"  => "mysql",
-          "encoding" => "utf8",
-          "database" => "ars_test",
-          "port"     => 123,
-          "username" => "root",
-          "password" => nil,
-          "host"     => "main_slave_host"
+          "adapter"     => "mysql",
+          "encoding"    => "utf8",
+          "database"    => "ars_test",
+          "port"        => 123,
+          "username"    => "root",
+          "password"    => nil,
+          "host"        => "main_slave_host",
+          "shard_names" => ["a", "b"]
         }, @conf)
       end
     end
@@ -26,13 +27,14 @@ class ConfigurationParserTest < ActiveSupport::TestCase
         setup { @conf = @exploded_conf['test_shard_a'] }
         should "be exploded" do
           assert_equal({
-            "adapter"  => "mysql",
-            "encoding" => "utf8",
-            "database" => "ars_test_shard_a",
-            "port"     => 123,
-            "username" => "root",
-            "password" => nil,
-            "host"     => "shard_a_host"
+            "adapter"     => "mysql",
+            "encoding"    => "utf8",
+            "database"    => "ars_test_shard_a",
+            "port"        => 123,
+            "username"    => "root",
+            "password"    => nil,
+            "host"        => "shard_a_host",
+            "shard_names" => ["a", "b"]
           }, @conf)
         end
       end
@@ -41,13 +43,14 @@ class ConfigurationParserTest < ActiveSupport::TestCase
         setup { @conf = @exploded_conf['test_shard_a_slave'] }
         should "be exploded" do
           assert_equal({
-            "adapter"  => "mysql",
-            "encoding" => "utf8",
-            "database" => "ars_test_shard_a",
-            "port"     => 123,
-            "username" => "root",
-            "password" => nil,
-            "host"     => "shard_a_slave_host"
+            "adapter"     => "mysql",
+            "encoding"    => "utf8",
+            "database"    => "ars_test_shard_a",
+            "port"        => 123,
+            "username"    => "root",
+            "password"    => nil,
+            "host"        => "shard_a_slave_host",
+            "shard_names" => ["a", "b"]
           }, @conf)
         end
       end
@@ -58,13 +61,14 @@ class ConfigurationParserTest < ActiveSupport::TestCase
         setup { @conf = @exploded_conf['test_shard_b'] }
         should "be exploded" do
           assert_equal({
-            "adapter"  => "mysql",
-            "encoding" => "utf8",
-            "database" => "ars_test_shard_b",
-            "port"     => 123,
-            "username" => "root",
-            "password" => nil,
-            "host"     => "shard_b_host"
+            "adapter"     => "mysql",
+            "encoding"    => "utf8",
+            "database"    => "ars_test_shard_b",
+            "port"        => 123,
+            "username"    => "root",
+            "password"    => nil,
+            "host"        => "shard_b_host",
+            "shard_names" => ["a", "b"]
           }, @conf)
         end
       end
@@ -73,13 +77,14 @@ class ConfigurationParserTest < ActiveSupport::TestCase
         setup { @conf = @exploded_conf['test_shard_b_slave'] }
         should "be exploded" do
           assert_equal({
-            "adapter"  => "mysql",
-            "encoding" => "utf8",
-            "database" => "ars_test_shard_b_slave",
-            "port"     => 123,
-            "username" => "root",
-            "password" => nil,
-            "host"     => "shard_b_host"
+            "adapter"     => "mysql",
+            "encoding"    => "utf8",
+            "database"    => "ars_test_shard_b_slave",
+            "port"        => 123,
+            "username"    => "root",
+            "password"    => nil,
+            "host"        => "shard_b_host",
+            "shard_names" => ["a", "b"]
           }, @conf)
         end
       end

data/test/connection_switching_test.rb

@@ -32,8 +32,102 @@ class ConnectionSwitchenTest < ActiveSupport::TestCase
       assert_using_database('ars_test')
       ActiveRecord::Base.on_slave { assert_using_database('ars_test_slave') }
     end
+
+    context "on_all_shards" do
+      setup do
+        @shard_0_master = ActiveRecord::Base.on_shard(0) {ActiveRecord::Base.connection}
+        @shard_1_master = ActiveRecord::Base.on_shard(1) {ActiveRecord::Base.connection}
+        assert_not_equal(@shard_0_master.select_value("SELECT DATABASE()"), @shard_1_master.select_value("SELECT DATABASE()"))
+
+        @database_names = []
+        ActiveRecord::Base.on_all_shards do
+          @database_names << ActiveRecord::Base.connection.select_value("SELECT DATABASE()")
+        end
+      end
+
+      should "execute the block on all shard masters" do
+        assert_equal(2, @database_names.size)
+        assert_contains(@database_names, @shard_0_master.select_value("SELECT DATABASE()"))
+        assert_contains(@database_names, @shard_1_master.select_value("SELECT DATABASE()"))
+      end
+    end
   end
-  
+
+  context "default shard selection" do
+    context "of nil" do
+      setup do
+        ActiveRecord::Base.default_shard = nil
+      end
+
+      should "use unsharded db for sharded models" do
+        assert_using_database('ars_test', Ticket)
+        assert_using_database('ars_test', Account)
+      end
+    end
+
+    context "value" do
+      setup do
+        ActiveRecord::Base.default_shard = 0
+      end
+
+      teardown do
+        ActiveRecord::Base.default_shard = nil
+      end
+
+      should "use default shard db for sharded models" do
+        assert_using_database('ars_test_shard0', Ticket)
+        assert_using_database('ars_test', Account)
+      end
+
+      should "still be able to switch to shard nil" do
+        ActiveRecord::Base.on_shard(nil) do
+          assert_using_database('ars_test', Ticket)
+          assert_using_database('ars_test', Account)
+        end
+      end
+    end
+  end
+
+  context "in an unsharded environment" do
+    setup do
+      silence_warnings { ::RAILS_ENV = 'test2' }
+      ActiveRecord::Base.establish_connection(::RAILS_ENV)
+      assert_using_database('ars_test2', Ticket)
+    end
+
+    teardown do
+      silence_warnings { ::RAILS_ENV = 'test' }
+      ActiveRecord::Base.establish_connection(::RAILS_ENV)
+      assert_using_database('ars_test', Ticket)
+    end
+
+    context "shard switching" do
+      should "just stay on the main db" do
+        assert_using_database('ars_test2', Ticket)
+        assert_using_database('ars_test2', Account)
+
+        ActiveRecord::Base.on_shard(0) do
+          assert_using_database('ars_test2', Ticket)
+          assert_using_database('ars_test2', Account)
+        end
+      end
+    end
+
+    context "on_all_shards" do
+      setup do
+        @database_names = []
+        ActiveRecord::Base.on_all_shards do
+          @database_names << ActiveRecord::Base.connection.select_value("SELECT DATABASE()")
+        end
+      end
+
+      should "execute the block on all shard masters" do
+        @database_names
+        assert_equal([ActiveRecord::Base.connection.select_value("SELECT DATABASE()")], @database_names)
+      end
+    end
+  end
+
   context "slave driving" do
     context "without slave configuration" do