ar-octopus 0.8.1 → 0.8.2

data/lib/octopus/relation_proxy.rb

@@ -0,0 +1,39 @@
+module Octopus
+  class RelationProxy
+    include Octopus::ShardTracking::Attribute
+
+    attr_accessor :ar_relation
+
+    def initialize(shard, ar_relation)
+      @current_shard = shard
+      @ar_relation = ar_relation
+    end
+
+    def method_missing(method, *args, &block)
+      run_on_shard { @ar_relation.send(method, *args, &block) }
+    end
+
+    def respond_to?(*args)
+      super || @ar_relation.respond_to?(*args)
+    end
+
+    # these methods are not normally sent to method_missing
+    def inspect
+      method_missing(:inspect)
+    end
+
+    def as_json(options = nil)
+      method_missing(:as_json, options)
+    end
+
+    def ==(other)
+      case other
+      when Octopus::RelationProxy
+        method_missing(:==, other.ar_relation)
+      else
+        method_missing(:==, other)
+      end
+    end
+    alias :eql? :==
+  end
+end

data/lib/octopus/scope_proxy.rb

@@ -1,33 +1,30 @@
 class Octopus::ScopeProxy
-  attr_accessor :shard, :klass
+  include Octopus::ShardTracking::Attribute
+  attr_accessor :klass
 
   def initialize(shard, klass)
-    @shard = shard
+    @current_shard = shard
     @klass = klass
   end
 
   def using(shard)
     raise "Nonexistent Shard Name: #{shard}" if @klass.connection.instance_variable_get(:@shards)[shard].nil?
-    @shard = shard
+    @current_shard = shard
     return self
   end
 
   # Transaction Method send all queries to a specified shard.
   def transaction(options = {}, &block)
-    @klass.connection.run_queries_on_shard(@shard) do
-      @klass = @klass.connection().transaction(options, &block)
-    end
+    run_on_shard { @klass = klass.transaction(options, &block) }
   end
 
   def connection
-    @klass.connection().current_shard = @shard
+    @klass.connection().current_shard = @current_shard
     @klass.connection()
   end
 
   def method_missing(method, *args, &block)
-    result = @klass.connection.run_queries_on_shard(@shard) do
-      @klass.send(method, *args, &block)
-    end
+    result = run_on_shard { @klass.send(method, *args, &block) }
 
     if result.respond_to?(:scoped)
       @klass = result

data/lib/octopus/shard_tracking.rb

@@ -0,0 +1,45 @@
+module Octopus::ShardTracking
+  def self.included(base)
+    base.extend(ClassMethods)
+  end
+
+  module ClassMethods
+    # If the class which includes this module responds to the class
+    # method sharded_methods, then automagically alias_method_chain
+    # a sharding-friendly version of each of those methods into existence
+    def sharded_methods(*methods)
+      methods.each { |m| create_sharded_method(m) }
+    end
+
+    def create_sharded_method(name)
+      name.to_s =~ /([^!?]+)([!?])?/
+      method, punctuation = [ $1, $2 ]
+      with = :"#{method}_with_octopus#{punctuation}"
+      without = :"#{method}_without_octopus#{punctuation}"
+      define_method with do |*args, &block|
+        run_on_shard { send(without, *args, &block) }
+      end
+      alias_method_chain name.to_sym, :octopus
+    end
+  end
+
+  # Adds run_on_shard method, but does not implement current_shard method
+  def run_on_shard(&block)
+    cs = current_shard
+    if !!cs
+      r = ActiveRecord::Base.connection_proxy.run_queries_on_shard(current_shard, &block)
+      # Use a case statement to avoid any path through ActiveRecord::Delegation's
+      # respond_to? code. We want to avoid the respond_to? code because it can have
+      # the side effect of causing a call to load_target
+      # return r
+      case r
+      when ActiveRecord::Relation
+        Octopus::RelationProxy.new(cs, r)
+      else
+        r
+      end
+    else
+      yield
+    end
+  end
+end

data/lib/octopus/shard_tracking/attribute.rb

@@ -0,0 +1,24 @@
+# Adds current_shard as an attribute; provide a default
+# implementation of set_current_shard which considers
+# only the current ActiveRecord::Base.connection_proxy
+module Octopus::ShardTracking::Attribute
+  def self.included(base)
+    base.send(:include, Octopus::ShardTracking)
+    base.extend(ClassMethods)
+    base.track_current_shard_as_attribute
+  end
+
+  module ClassMethods
+    def track_current_shard_as_attribute
+      attr_accessor :current_shard
+    end
+  end
+
+  def set_current_shard
+    return unless Octopus.enabled?
+
+    if ActiveRecord::Base.connection_proxy.block
+      self.current_shard = ActiveRecord::Base.connection_proxy.current_shard
+    end
+  end
+end

data/lib/octopus/shard_tracking/dynamic.rb

@@ -0,0 +1,7 @@
+require 'octopus/shard_tracking'
+
+module Octopus::ShardTracking::Dynamic
+  def self.included(base)
+    base.send(:include, Octopus::ShardTracking)
+  end
+end

data/lib/octopus/singular_association.rb

@@ -0,0 +1,7 @@
+module Octopus::SingularAssociation
+  def self.included(base)
+    base.sharded_methods :reader, :writer, :create, :create!, :build
+  end
+end
+
+ActiveRecord::Associations::SingularAssociation.send(:include, Octopus::SingularAssociation)

data/lib/octopus/slave_group.rb

@@ -0,0 +1,11 @@
+class Octopus::SlaveGroup
+  def initialize(slaves)
+    slaves = HashWithIndifferentAccess.new(slaves)
+    slaves_list = slaves.values
+    @load_balancer = Octopus::LoadBalancing::RoundRobin.new(slaves_list)
+  end
+
+  def next
+    @load_balancer.next
+  end
+end

data/lib/octopus/version.rb

@@ -1,3 +1,3 @@
 module Octopus
-  VERSION = '0.8.1'
+  VERSION = '0.8.2'
 end

data/spec/config/shards.yml

@@ -142,6 +142,59 @@ not_entire_sharded:
       database: octopus_shard_5
       <<: *mysql
 
+sharded_replicated_slave_grouped:
+  replicated: true
+  fully_replicated: true
+  shards:
+    russia:
+      database: octopus_shard_1
+      <<: *mysql
+      slaves1:
+        russia_slave11:
+          database: octopus_shard_1
+          <<: *mysql
+        russia_slave21:
+          database: octopus_shard_2
+          <<: *mysql
+      slaves2:
+        russia_slave21:
+          database: octopus_shard_3
+          <<: *mysql
+        russia_slave22:
+            database: octopus_shard_1
+            <<: *mysql
+    europe:
+      database: octopus_shard_2
+      <<: *mysql
+      slaves1:
+        europe_slave11:
+          database: octopus_shard_3
+          <<: *mysql
+      slaves2:
+        europe_slave21:
+          database: octopus_shard_1
+          <<: *mysql
+
+replicated_slave_grouped:
+  replicated: true
+  fully_replicated: true
+  shards:
+    slaves1:
+      slave11:
+        database: octopus_shard_2
+        <<: *mysql
+    slaves2:
+      slave21:
+        database: octopus_shard_1
+        <<: *mysql
+    slaves3:
+      slave31:
+        database: octopus_shard_1
+        <<: *mysql
+      slave32:
+        database: octopus_shard_2
+        <<: *mysql
+
 modify_config:
   replicated: true
   shards:

data/spec/octopus/{association_spec.rb → association_shard_tracking_spec.rb}

@@ -1,6 +1,6 @@
 require "spec_helper"
 
-describe Octopus::Association, :shards => [:brazil, :master, :canada] do
+describe Octopus::AssociationShardTracking, :shards => [:brazil, :master, :canada] do
   describe "when you have a 1 x 1 relationship" do
     before(:each) do
       @computer_brazil = Computer.using(:brazil).create!(:name => "Computer Brazil")

data/spec/octopus/collection_proxy_spec.rb

@@ -0,0 +1,15 @@
+require "spec_helper"
+
+describe Octopus::CollectionProxy do
+  describe "method dispatch" do
+    before :each do
+      @client = Client.using(:canada).create!
+      @client.items << Item.using(:canada).create!
+    end
+
+    it "computes the size of the collection without loading it" do
+      @client.items.size.should eq(1)
+      @client.items.should_not be_loaded
+    end
+  end
+end

data/spec/octopus/model_spec.rb

@@ -361,14 +361,14 @@ describe Octopus::Model do
         @user = User.using(:brazil).create!(:name => "User1")
         User.using(:brazil).update_all({:updated_at => Time.now - 3.months}, {:id => @user.id})
         @user.touch
-        @user.reload.updated_at.to_date.should eq(Date.today)
+        @user.reload.updated_at.in_time_zone('GMT').to_date.should eq(Time.now.in_time_zone('GMT').to_date)
       end
 
       it "updates passed in attribute name" do
         @user = User.using(:brazil).create!(:name => "User1")
         User.using(:brazil).update_all({:created_at => Time.now - 3.months}, {:id => @user.id})
         @user.touch(:created_at)
-        @user.reload.created_at.to_date.should eq(Date.today)
+        @user.reload.created_at.in_time_zone('GMT').to_date.should eq(Time.now.in_time_zone('GMT').to_date)
       end
     end
 

data/spec/octopus/octopus_spec.rb

@@ -87,4 +87,38 @@ describe Octopus, :shards => [] do
       Octopus.should_not be_enabled
     end
   end
+
+  describe "#fully_replicated" do
+    before do
+      OctopusHelper.using_environment :production_replicated do
+        OctopusHelper.clean_all_shards([:slave1, :slave2, :slave3, :slave4])
+        4.times { |i| User.using(:"slave#{i+1}").create!(:name => "Slave User") }
+      end
+    end
+
+    it "sends queries to slaves" do
+      OctopusHelper.using_environment :production_replicated do
+        User.count.should eq(0)
+        4.times do |i|
+          Octopus.fully_replicated do
+            User.count.should eq(1)
+          end
+        end
+      end
+    end
+
+    it "allows nesting" do
+      OctopusHelper.using_environment :production_replicated do
+        Octopus.fully_replicated do
+          User.count.should eq(1)
+
+          Octopus.fully_replicated do
+            User.count.should eq(1)
+          end
+
+          User.count.should eq(1)
+        end
+      end
+    end
+  end
 end

data/spec/octopus/relation_proxy_spec.rb

@@ -0,0 +1,77 @@
+require "spec_helper"
+
+describe Octopus::RelationProxy do
+  describe "shard tracking" do
+    before :each do
+      @client = Client.using(:canada).create!
+      @client.items << Item.using(:canada).create!
+      @relation = @client.items
+    end
+
+    it "remembers the shard on which a relation was created" do
+      @relation.current_shard.should eq(:canada)
+    end
+
+    context "when comparing to other Relation objects" do
+      before :each do
+        @relation.reset
+      end
+
+      it "is equal to its clone" do
+        @relation.should eq(@relation.clone)
+      end
+    end
+
+    if Octopus.rails4?
+      context "under Rails 4" do
+        it "is an Octopus::RelationProxy" do
+          @relation.class.should eq(Octopus::RelationProxy)
+        end
+
+        it "should be able to return its ActiveRecord::Relation" do
+          @relation.ar_relation.is_a?(ActiveRecord::Relation).should be_true
+        end
+
+        it "is equal to an identically-defined, but different, RelationProxy" do
+          i = @client.items
+          @relation.should eq(i)
+          @relation.object_id.should_not eq(i.object_id)
+        end
+
+        it "is equal to its own underlying ActiveRecord::Relation" do
+          @relation.should eq(@relation.ar_relation)
+          @relation.ar_relation.should eq(@relation)
+        end
+      end
+    end
+
+    context "when no explicit shard context is provided" do
+      it "uses the correct shard" do
+        @relation.count.should eq(1)
+      end
+
+      it "lazily evaluates on the correct shard" do
+        # Do something to force Client.connection_proxy.current_shard to change
+        other_count = Client.using(:brazil).count
+        @relation.select(:client_id).count.should == 1
+      end
+    end
+
+    context "when an explicit, but different, shard context is provided" do
+      it "uses the correct shard" do
+        Item.using(:brazil).count.should eq(0)
+        clients_on_brazil = Client.using(:brazil).all
+        Client.using(:brazil) do
+          @relation.count.should eq(1)
+        end
+      end
+
+      it "lazily evaluates on the correct shard" do
+        Item.using(:brazil).count.should eq(0)
+        Client.using(:brazil) do
+          @relation.select(:client_id).count.should == 1
+        end
+      end
+    end
+  end
+end

data/spec/octopus/replicated_slave_grouped_spec.rb

@@ -0,0 +1,64 @@
+require "spec_helper"
+
+describe "when the database is replicated and has slave groups" do
+
+  it "should pick the slave group based on current_slave_grup when you have a replicated model" do
+
+    OctopusHelper.using_environment :replicated_slave_grouped do
+      # The following two calls of `create!` both creates cats in :master(The database `octopus_shard_1`)
+      # which is configured through RAILS_ENV and database.yml
+      Cat.create!(:name => "Thiago1")
+      Cat.create!(:name => "Thiago2")
+
+      # See "replicated_slave_grouped" defined in shards.yml
+      # We have:
+      #   The database `octopus_shard_1` as :slave21 which is a member of the slave group :slaves2, and as :master
+      #   The databse `octopus_shard_2` as :slave11 which is a member of the slave group :slaves1
+      # When a select-count query is sent to `octopus_shard_1`, it should return 2 because we have create two cats in :master .
+      # When a select-count query is sent to `octopus_shard_2`, it should return 0.
+
+      # The query goes to `octopus_shard_1`
+      Cat.using(:master).count.should == 2
+      # The query goes to `octopus_shard_1`
+      Cat.count.should == 2
+      # The query goes to `octopus_shard_2`
+      Cat.using(slave_group: :slaves1).count.should == 0
+      # The query goes to `octopus_shard_1`
+      Cat.using(slave_group: :slaves2).count.should == 2
+    end
+  end
+
+  it "should distribute queries between slaves in a slave group in round-robin" do
+    OctopusHelper.using_environment :replicated_slave_grouped do
+      # The query goes to :master(`octopus_shard_1`)
+      Cat.create!(:name => "Thiago1")
+      # The query goes to :master(`octopus_shard_1`)
+      Cat.create!(:name => "Thiago2")
+
+      # The query goes to :slave32(`octopus_shard_2`)
+      Cat.using(slave_group: :slaves3).count.should == 0
+      # The query goes to :slave31(`octopus_shard_1`)
+      Cat.using(slave_group: :slaves3).count.should == 2
+      # The query goes to :slave32(`octopus_shard_2`)
+      Cat.using(slave_group: :slaves3).count.should == 0
+    end
+  end
+
+  it "should make queries to master when slave groups are configured but not selected" do
+    OctopusHelper.using_environment :replicated_slave_grouped do
+      # All the queries go to :master(`octopus_shard_1`)
+
+      Cat.create!(:name => "Thiago1")
+      Cat.create!(:name => "Thiago2")
+
+      # In `database.yml` and `shards.yml`, we have configured 1 master and 4 slaves.
+      # So we can ensure Octopus is not distributing queries between them
+      # by asserting 1 + 4 = 5 queries go to :master(`octopus_shard_1`)
+      Cat.count.should == 2
+      Cat.count.should == 2
+      Cat.count.should == 2
+      Cat.count.should == 2
+      Cat.count.should == 2
+    end
+  end
+end