active_record_shards 2.0.0.beta1

data/README.rdoc

@@ -0,0 +1,33 @@
+= active_record_shards
+
+== Example
+  
+  Ticket.on_slave do
+    Ticket.first(50)
+  end
+
+== Install
+
+gem install active_record_shards
+
+Add the slave database to config/database.yml:
+  development_slave:
+    adapter: mysql
+    ...
+
+== Note on Patches/Pull Requests
+ 
+* Fork the project.
+* Make your feature addition or bug fix.
+* Add tests for it. This is important so I don't break it in a
+  future version unintentionally.
+* Commit, do not mess with rakefile, version, or history.
+  (if you want to have your own version, that is fine but bump version in a commit by itself I can ignore when I pull)
+* Send me a pull request. Bonus points for topic branches.
+
+== Copyright
+
+Copyright (c) 2009 Zendesk. See LICENSE for details.
+
+== Authors
+Mick Staugaard

data/lib/active_record_shards.rb

@@ -0,0 +1,13 @@
+require 'active_record'
+require 'active_record/base'
+require 'active_record_shards/configuration_parser'
+require 'active_record_shards/model'
+require 'active_record_shards/shard_selection'
+require 'active_record_shards/connection_switcher'
+require 'active_record_shards/association_collection_connection_selection'
+require 'active_record_shards/connection_pool'
+
+ActiveRecord::Base.extend(ActiveRecordShards::ConfigurationParser)
+ActiveRecord::Base.extend(ActiveRecordShards::Model)
+ActiveRecord::Base.extend(ActiveRecordShards::ConnectionSwitcher)
+ActiveRecord::Associations::AssociationCollection.send(:include, ActiveRecordShards::AssociationCollectionConnectionSelection)

data/lib/active_record_shards/association_collection_connection_selection.rb

@@ -0,0 +1,25 @@
+module ActiveRecordShards
+  module AssociationCollectionConnectionSelection
+    def on_slave_if(condition)
+      condition ? on_slave : self
+    end
+
+    def on_slave_unless(condition)
+      on_slave_if(!condition)
+    end
+
+    def on_slave
+      SlaveProxy.new(self)
+    end
+
+    class SlaveProxy
+      def initialize(association_collection)
+        @association_collection = association_collection
+      end
+
+      def method_missing(method, *args, &block)
+        @association_collection.proxy_reflection.klass.on_slave_block { @association_collection.send(method, *args, &block) }
+      end
+    end
+  end
+end

data/lib/active_record_shards/configuration_parser.rb

@@ -0,0 +1,46 @@
+module ActiveRecordShards
+  module ConfigurationParser
+    module_function
+    
+    def explode(conf)
+      conf.keys.each do |env_name|
+        env_config = conf[env_name]
+        if shards = env_config.delete('shards')
+          shards.each do |shard_name, shard_conf|
+            expand_child!(env_config, shard_conf)
+            conf["#{env_name}_shard_#{shard_name}"] = shard_conf
+          end
+        end
+      end
+
+      conf.keys.each do |env_name|
+        env_config = conf[env_name]
+        if slave_conf = env_config.delete('slave')
+          expand_child!(env_config, slave_conf)
+          conf["#{env_name}_slave"] = slave_conf
+        end
+      end
+      conf
+    end
+
+    def expand_child!(parent, child)
+      parent.each do |key, value|
+        unless ['slave', 'shards'].include?(key) || value.is_a?(Hash)
+          child[key] ||= value
+        end
+      end
+    end
+
+    def configurations_with_shard_explosion=(conf)
+      self.configurations_without_shard_explosion = explode(conf)
+    end
+
+    def ConfigurationParser.extended(klass)
+      klass.singleton_class.alias_method_chain :configurations=, :shard_explosion
+
+      if !klass.configurations.nil? && !klass.configurations.empty?
+        klass.configurations = klass.configurations
+      end
+    end
+  end
+end

data/lib/active_record_shards/connection_pool.rb

@@ -0,0 +1,29 @@
+ActiveRecord::ConnectionAdapters::ConnectionHandler.class_eval do
+  # The only difference here is that we use klass.connection_pool_name
+  # instead of klass.name as the pool key
+  def retrieve_connection_pool(klass)
+    pool = @connection_pools[klass.connection_pool_name]
+    return pool if pool
+    return nil if ActiveRecord::Base == klass
+    retrieve_connection_pool klass.superclass
+  end
+
+  def remove_connection(klass)
+    pool = @connection_pools[klass.connection_pool_name]
+    @connection_pools.delete_if { |key, value| value == pool }
+    pool.disconnect! if pool
+    pool.spec.config if pool
+  end
+end
+
+ActiveRecord::Base.singleton_class.class_eval do
+  def establish_connection_with_connection_pool_name(spec = nil)
+    case spec
+    when ActiveRecord::Base::ConnectionSpecification
+      connection_handler.establish_connection(connection_pool_name, spec)
+    else
+      establish_connection_without_connection_pool_name(spec)
+    end
+  end
+  alias_method_chain :establish_connection, :connection_pool_name
+end

data/lib/active_record_shards/connection_switcher.rb

@@ -0,0 +1,104 @@
+module ActiveRecordShards
+  module ConnectionSwitcher
+
+    def on_shard(shard, &block)
+      old_shard = current_shard_selection.shard
+      switch_connection(:shard => shard)
+      yield
+    ensure
+      switch_connection(:shard => old_shard)
+    end
+
+    def on_slave_if(condition, &block)
+      condition ? on_slave(&block) : yield
+    end
+
+    def on_slave_unless(condition, &block)
+      on_slave_if(!condition, &block)
+    end
+
+    # Executes queries using the slave database. Fails over to master if no slave is found.
+    # if you want to execute a block of code on the slave you can go:
+    #   Account.on_slave do
+    #     Account.first
+    #   end
+    # the first account will be found on the slave DB
+    #
+    # For one-liners you can simply do
+    #   Account.on_slave.first
+    def on_slave(&block)
+      if block_given?
+        on_slave_block(&block)
+      else
+        SlaveProxy.new(self)
+      end
+    end
+
+    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
+      with_scope({:find => {:readonly => current_shard_selection.on_slave?}}, :merge, &block)
+    ensure
+      switch_connection(:slave => old_slave)
+    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
+    end
+
+    private
+
+    def current_shard_selection
+      Thread.current[:shard_selection] ||= ShardSelection.new
+    end
+
+    def switch_connection(options)
+      if options.any?
+        if options.has_key?(:slave)
+          current_shard_selection.on_slave = options[:slave]
+        end
+
+        if options.has_key?(:shard)
+          current_shard_selection.shard = options[:shard]
+        end
+
+        establish_shard_connection unless connected_to_shard?
+      end
+    end
+
+    def establish_shard_connection
+      name = current_shard_selection.shard_name(self)
+      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"))
+    end
+
+    def connected_to_shard?
+      connection_handler.connection_pools.has_key?(connection_pool_name)
+    end
+
+    class SlaveProxy
+      def initialize(target)
+        @target = target
+      end
+
+      def method_missing(method, *args, &block)
+        @target.on_slave_block { @target.send(method, *args, &block) }
+      end
+    end
+  end
+end

data/lib/active_record_shards/model.rb

@@ -0,0 +1,20 @@
+module ActiveRecordShards
+  module Model
+    def not_sharded
+      if self == ActiveRecord::Base || self != base_class
+        raise "You should only call not_sharded on direct descendants of ActiveRecord::Base"
+      end
+      @sharded = false
+    end
+
+    def is_sharded?
+      if self == ActiveRecord::Base
+        true
+      elsif self == base_class
+        @sharded != false
+      else
+        base_class.is_sharded?
+      end
+    end
+  end
+end

data/lib/active_record_shards/shard_selection.rb

@@ -0,0 +1,41 @@
+module ActiveRecordShards
+  class ShardSelection
+    def initialize
+      @on_slave = false
+    end
+
+    def shard
+      @shard
+    end
+
+    def shard=(new_shard)
+      @shard = new_shard
+    end
+
+    def on_slave?
+      @on_slave
+    end
+
+    def on_slave=(new_slave)
+      @on_slave = (new_slave == true)
+    end
+
+    def connection_configuration_name
+      shard_name#(RAILS_ENV)
+    end
+
+    def shard_name(klass = nil)
+      s = "#{RAILS_ENV}"
+      if @shard && (klass.nil? || klass.is_sharded?)
+        s << '_shard_'
+        s << @shard.to_s
+      end
+      s << "_slave" if @on_slave
+      s
+    end
+
+    def any?
+      @on_slave || @shard.present?
+    end
+  end
+end

data/test/configuration_parser_test.rb

@@ -0,0 +1,88 @@
+require File.expand_path('helper', File.dirname(__FILE__))
+
+class ConfigurationParserTest < ActiveSupport::TestCase
+  context "exploding the database.yml" do
+    setup do
+      @exploded_conf = ActiveRecordShards::ConfigurationParser.explode(YAML::load(IO.read(File.dirname(__FILE__) + '/database_parse_test.yml')))
+    end
+
+    context "main slave" do
+      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"
+        }, @conf)
+      end
+    end
+
+    context "shard a" do
+      context "master" do
+        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"
+          }, @conf)
+        end
+      end
+
+      context "slave" do
+        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"
+          }, @conf)
+        end
+      end
+    end
+
+    context "shard b" do
+      context "master" do
+        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"
+          }, @conf)
+        end
+      end
+
+      context "slave" do
+        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"
+          }, @conf)
+        end
+      end
+    end
+  end
+end

data/test/connection_switching_test.rb

@@ -0,0 +1,167 @@
+require File.expand_path('helper', File.dirname(__FILE__))
+
+class ConnectionSwitchenTest < ActiveSupport::TestCase
+  context "shard switching" do
+    should "only switch connection on sharded models" do
+      assert_using_database('ars_test', Ticket)
+      assert_using_database('ars_test', Account)
+
+      ActiveRecord::Base.on_shard(0) do
+        assert_using_database('ars_test_shard0', Ticket)
+        assert_using_database('ars_test', Account)
+      end
+    end
+
+    should "switch to shard and back" do
+      assert_using_database('ars_test')
+      ActiveRecord::Base.on_slave { assert_using_database('ars_test_slave') }
+
+      ActiveRecord::Base.on_shard(0) do
+        assert_using_database('ars_test_shard0')
+        ActiveRecord::Base.on_slave { assert_using_database('ars_test_shard0_slave') }
+
+        ActiveRecord::Base.on_shard(nil) do
+          assert_using_database('ars_test')
+          ActiveRecord::Base.on_slave { assert_using_database('ars_test_slave') }
+        end
+
+        assert_using_database('ars_test_shard0')
+        ActiveRecord::Base.on_slave { assert_using_database('ars_test_shard0_slave') }
+      end
+
+      assert_using_database('ars_test')
+      ActiveRecord::Base.on_slave { assert_using_database('ars_test_slave') }
+    end
+  end
+  
+  context "slave driving" do
+    context "without slave configuration" do
+
+      setup do
+        ActiveRecord::Base.configurations.delete('test_slave')
+        ActiveRecord::Base.connection_handler.connection_pools.clear
+        ActiveRecord::Base.establish_connection('test')
+      end
+
+      should "default to the master database" do
+        Account.create!
+
+        ActiveRecord::Base.on_slave { assert_using_master_db }
+        Account.on_slave { assert_using_master_db }
+        Ticket.on_slave  { assert_using_master_db }
+      end
+
+      should "successfully execute queries" do
+        Account.create!
+        assert_using_master_db
+
+        assert_equal Account.count, ActiveRecord::Base.on_slave { Account.count }
+        assert_equal Account.count, Account.on_slave { Account.count }
+      end
+
+    end
+
+    context "with slave configuration" do
+
+      should "successfully execute queries" do
+        assert_using_master_db
+        Account.create!
+
+        assert_equal(1, Account.count)
+        assert_equal(0, ActiveRecord::Base.on_slave { Account.count })
+      end
+
+      should "support global on_slave blocks" do
+        assert_using_master_db
+        assert_using_master_db
+
+        ActiveRecord::Base.on_slave do
+          assert_using_slave_db
+          assert_using_slave_db
+        end
+
+        assert_using_master_db
+        assert_using_master_db
+      end
+
+      should "support conditional methods" do
+        assert_using_master_db
+
+        Account.on_slave_if(true) do
+          assert_using_slave_db
+        end
+
+        assert_using_master_db
+
+        Account.on_slave_if(false) do
+          assert_using_master_db
+        end
+
+        Account.on_slave_unless(true) do
+          assert_using_master_db
+        end
+
+        Account.on_slave_unless(false) do
+          assert_using_slave_db
+        end
+
+      end
+
+      context "a model loaded with the slave" do
+        setup do
+          Account.connection.execute("INSERT INTO accounts (id, name, created_at, updated_at) VALUES(1000, 'master_name', '2009-12-04 20:18:48', '2009-12-04 20:18:48')")
+          assert(Account.find(1000))
+          assert_equal('master_name', Account.find(1000).name)
+
+          Account.on_slave.connection.execute("INSERT INTO accounts (id, name, created_at, updated_at) VALUES(1000, 'slave_name', '2009-12-04 20:18:48', '2009-12-04 20:18:48')")
+
+          @model = Account.on_slave.find(1000)
+          assert(@model)
+          assert_equal('slave_name', @model.name)
+        end
+
+        should "read from master on reload" do
+          @model.reload
+          assert_equal('master_name', @model.name)
+        end
+
+        should "be marked as read only" do
+          assert(@model.readonly?)
+        end
+      end
+
+      context "a model loaded with the master" do
+        setup do
+          Account.connection.execute("INSERT INTO accounts (id, name, created_at, updated_at) VALUES(1000, 'master_name', '2009-12-04 20:18:48', '2009-12-04 20:18:48')")
+          @model = Account.first
+          assert(@model)
+          assert_equal('master_name', @model.name)
+        end
+
+        should "not be marked as read only" do
+          assert(!@model.readonly?)
+        end
+      end
+    end
+
+    context "slave proxy" do
+      should "successfully execute queries" do
+        assert_using_master_db
+        Account.create!
+
+        assert_not_equal Account.count, Account.on_slave.count
+      end
+
+      should "work on association collections" do
+        assert_using_master_db
+        account = Account.create!
+
+        Ticket.connection.expects(:select_all).with("SELECT * FROM `tickets` WHERE (`tickets`.account_id = #{account.id})  LIMIT 1", anything).returns([])
+        Ticket.on_slave.connection.expects(:select_all).with("SELECT * FROM `tickets` WHERE (`tickets`.account_id = #{account.id})  LIMIT 1", anything).returns([])
+
+        account.tickets.first
+        account.tickets.on_slave.first
+      end
+    end
+  end
+end