ar-octopus 0.8.1 → 0.10.2

data/gemfiles/{rails32.gemfile → rails5.gemfile}

@@ -2,6 +2,6 @@
 
 source "https://rubygems.org"
 
-gem "activerecord", "~> 3.2.0"
+gem "activerecord", "~> 5.0.0"
 
-gemspec :path=>"../"
+gemspec path: "../"

data/gemfiles/{rails4.gemfile → rails51.gemfile}

@@ -2,6 +2,6 @@
 
 source "https://rubygems.org"
 
-gem "activerecord", "~> 4.0.0"
+gem "activerecord", "~> 5.1.0"
 
-gemspec :path=>"../"
+gemspec path: "../"

data/gemfiles/rails52.gemfile

@@ -0,0 +1,7 @@
+# This file was generated by Appraisal
+
+source "https://rubygems.org"
+
+gem "activerecord", "~> 5.2.0"
+
+gemspec path: "../"

data/lib/ar-octopus.rb

@@ -1 +1 @@
-require "octopus"
+require 'octopus'

data/lib/octopus/{rails3/abstract_adapter.rb → abstract_adapter.rb}

@@ -2,13 +2,7 @@
 module Octopus
   module AbstractAdapter
     module OctopusShard
-
-      parent = ActiveSupport::BasicObject
-      if Octopus.rails4?
-        parent = ActiveSupport::ProxyObject
-      end
-
-      class InstrumenterDecorator < parent
+      class InstrumenterDecorator < ActiveSupport::ProxyObject
         def initialize(adapter, instrumenter)
           @adapter = adapter
           @instrumenter = instrumenter
@@ -24,21 +18,16 @@ module Octopus
         end
       end
 
-      def self.included(base)
-        base.alias_method_chain :initialize, :octopus_shard
-      end
-
       def octopus_shard
         @config[:octopus_shard]
       end
 
-      def initialize_with_octopus_shard(*args)
-        initialize_without_octopus_shard(*args)
+      def initialize(*args)
+        super
         @instrumenter = InstrumenterDecorator.new(self, @instrumenter)
       end
-
     end
   end
 end
 
-ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:include, Octopus::AbstractAdapter::OctopusShard)
+ActiveRecord::ConnectionAdapters::AbstractAdapter.send(:prepend, Octopus::AbstractAdapter::OctopusShard)

data/lib/octopus/association.rb

@@ -1,105 +1,14 @@
-module Octopus::Association
-  def self.extended(base)
-    base.send(:include, InstanceMethods)
-  end
-
-  module QueryOnCurrentShard
-
-    METHODS= %w[
-      all
-      average
-      count
-      empty?
-      exists?
-      find
-      find_by_sql
-      first
-      last
-      maximum
-      minimum
-      pluck
-      scoping
-      size
-      sum
-      to_a
-    ]
-
-    METHODS.each do |m|
-      define_method m.to_sym do |*args,&block|
-        if self.respond_to?(:proxy_association) and self.proxy_association
-          self.proxy_association.owner.run_on_shard { super(*args, &block) }
-        else
-          super(*args, &block)
-        end
-      end
-    end
-
-  end
-
-  module InstanceMethods
-    def set_connection_on_association(record)
-      return unless ::Octopus.enabled?
-      return if !self.class.connection.respond_to?(:current_shard) || !self.respond_to?(:current_shard)
-      if !record.current_shard.nil? && !self.current_shard.nil? && record.current_shard != self.current_shard
-        raise "Association Error: Records are from different shards"
-      end
-
-      record.current_shard = self.class.connection.current_shard = self.current_shard if should_set_current_shard?
-    end
-  end
-
-  if Octopus.rails4?
-    def has_many(association_id, scope=nil, options={}, &extension)
-      if options == {} && scope.is_a?(Hash)
-        default_octopus_opts(scope)
-      else
-        default_octopus_opts(options)
-      end
-      super
-    end
-  else
-    def has_many(association_id, options={}, &extension)
-      default_octopus_opts(options)
-      super
+module Octopus
+  module Association
+    def self.included(base)
+      base.send(:include, Octopus::ShardTracking::Dynamic)
+      base.sharded_methods :target_scope
     end
-  end
 
-
-  if Octopus.rails4?
-    def has_and_belongs_to_many(association_id, scope=nil, options={}, &extension)
-      if options == {} && scope.is_a?(Hash)
-        default_octopus_opts(scope)
-      else
-        default_octopus_opts(options)
-      end
-      super
-    end
-  else
-    def has_and_belongs_to_many(association_id, options={}, &extension)
-      default_octopus_opts(options)
-      super
+    def current_shard
+      owner.current_shard
     end
   end
-
-  def default_octopus_opts(options)
-    if options[:before_add].is_a?(Array)
-      options[:before_add] << :set_connection_on_association
-    elsif options[:before_add].is_a?(Symbol)
-      options[:before_add] = [:set_connection_on_association, options[:before_add]]
-    else
-      options[:before_add] = :set_connection_on_association
-    end
-
-    if options[:before_remove].is_a?(Array)
-      options[:before_remove] << :set_connection_on_association
-    elsif options[:before_remove].is_a?(Symbol)
-      options[:before_remove] = [:set_connection_on_association, options[:before_remove]]
-    else
-      options[:before_remove] = :set_connection_on_association
-    end
-
-    options[:extend] = [ Octopus::Association::QueryOnCurrentShard, options[:extend] ].flatten.compact
-  end
 end
 
-ActiveRecord::Base.extend(Octopus::Association)
+ActiveRecord::Associations::Association.send(:include, Octopus::Association)

data/lib/octopus/association_shard_tracking.rb

@@ -0,0 +1,74 @@
+module Octopus
+  module AssociationShardTracking
+    class MismatchedShards < StandardError
+      attr_reader :record, :current_shard
+
+      def initialize(record, current_shard)
+        @record = record
+        @current_shard = current_shard
+      end
+
+      def message
+        [
+          "Association Error: Records are from different shards",
+          "Record: #{record.inspect}",
+          "Current Shard: #{current_shard.inspect}",
+          "Current Record Shard: #{record.current_shard.inspect}",
+        ].join(" ")
+      end
+    end
+
+    def self.extended(base)
+      base.send(:include, InstanceMethods)
+    end
+
+    module InstanceMethods
+      def connection_on_association=(record)
+        return unless ::Octopus.enabled?
+        return if !self.class.connection.respond_to?(:current_shard) || !self.respond_to?(:current_shard)
+
+        if !record.current_shard.nil? && !current_shard.nil? && record.current_shard.to_s != current_shard.to_s
+          raise MismatchedShards.new(record, current_shard)
+        end
+
+        record.current_shard = self.class.connection.current_shard = current_shard if should_set_current_shard?
+      end
+    end
+
+    def has_many(name, scope = nil, **options, &extension)
+      if options == {} && scope.is_a?(Hash)
+        default_octopus_opts(scope)
+      else
+        default_octopus_opts(options)
+      end
+      super
+    end
+
+    def has_and_belongs_to_many(association_id, scope = nil, options = {}, &extension)
+      assign_octopus_opts(scope, options)
+      super
+    end
+
+    def default_octopus_opts(options)
+      options[:before_add] = [ :connection_on_association=, options[:before_add] ].compact.flatten
+      options[:before_remove] = [ :connection_on_association=, options[:before_remove] ].compact.flatten
+    end
+
+    def assign_octopus_opts(scope, options)
+      if options == {} && scope.is_a?(Hash)
+        default_octopus_opts(scope)
+      else
+        default_octopus_opts(options)
+      end
+    end
+
+    if Octopus.atleast_rails51?
+      def has_and_belongs_to_many(association_id, scope = nil, **options, &extension)
+        assign_octopus_opts(scope, options)
+        super
+      end
+    end
+  end
+end
+
+ActiveRecord::Base.extend(Octopus::AssociationShardTracking)

data/lib/octopus/collection_association.rb

@@ -0,0 +1,17 @@
+module Octopus
+  module CollectionAssociation
+    def self.included(base)
+      if Octopus.rails51? || Octopus.rails52?
+        base.sharded_methods :reader, :writer, :ids_reader, :ids_writer, :create, :create!,
+                             :build, :include?,
+                             :load_target, :reload, :size, :select
+      else
+        base.sharded_methods :reader, :writer, :ids_reader, :ids_writer, :create, :create!,
+                             :build, :any?, :count, :empty?, :first, :include?, :last, :length,
+                             :load_target, :many?, :reload, :size, :select, :uniq
+      end
+    end
+  end
+end
+
+ActiveRecord::Associations::CollectionAssociation.send(:include, Octopus::CollectionAssociation)

data/lib/octopus/collection_proxy.rb

@@ -0,0 +1,16 @@
+module Octopus
+  module CollectionProxy
+    def self.included(base)
+      base.send(:include, Octopus::ShardTracking::Dynamic)
+      base.sharded_methods :any?, :build, :count, :create, :create!, :concat, :delete, :delete_all,
+                           :destroy, :destroy_all, :empty?, :find, :first, :include?, :last, :length,
+                           :many?, :pluck, :replace, :select, :size, :sum, :to_a, :uniq
+    end
+
+    def current_shard
+      @association.owner.current_shard
+    end
+  end
+end
+
+ActiveRecord::Associations::CollectionProxy.send(:include, Octopus::CollectionProxy)

data/lib/octopus/exception.rb

@@ -0,0 +1,4 @@
+module Octopus
+  class Exception < ::Exception
+  end
+end

data/lib/octopus/finder_methods.rb

@@ -0,0 +1,8 @@
+# find_nth / find_nth! must be public here to allow Octopus to call
+# them on the scope proxy object
+if Octopus.rails42? || Octopus.rails50?
+  module ActiveRecord::FinderMethods
+    public :find_nth
+    public :find_nth!
+  end
+end

data/lib/octopus/load_balancing/round_robin.rb

@@ -0,0 +1,20 @@
+require 'octopus/load_balancing'
+
+# The round-robin load balancing of slaves belonging to the same shard.
+# It is a pool that contains slaves which queries are distributed to.
+module Octopus
+  module LoadBalancing
+    class RoundRobin
+      def initialize(slaves_list)
+        raise Octopus::Exception.new("No slaves available") if slaves_list.empty?
+        @slaves_list = slaves_list
+        @slave_index = 0
+      end
+
+      # Returns the next available slave in the pool
+      def next(options)
+        @slaves_list[@slave_index = (@slave_index + 1) % @slaves_list.length]
+      end
+    end
+  end
+end

data/lib/octopus/load_balancing.rb

@@ -0,0 +1,4 @@
+module Octopus
+  module LoadBalancing
+  end
+end

data/lib/octopus/{rails3/log_subscriber.rb → log_subscriber.rb}

@@ -3,8 +3,12 @@ module Octopus
   module LogSubscriber
     def self.included(base)
       base.send(:attr_accessor, :octopus_shard)
-      base.alias_method_chain :sql, :octopus_shard
-      base.alias_method_chain :debug, :octopus_shard
+
+      base.send :alias_method, :sql_without_octopus_shard, :sql
+      base.send :alias_method, :sql, :sql_with_octopus_shard
+
+      base.send :alias_method, :debug_without_octopus_shard, :debug
+      base.send :alias_method, :debug, :debug_with_octopus_shard
     end
 
     def sql_with_octopus_shard(event)