association.rb in ar-octopus-0.8.2

- old
+ new
@@ -1,105 +1,11 @@
 module Octopus::Association
-  def self.extended(base)
-    base.send(:include, InstanceMethods)
+  def self.included(base)
+    base.send(:include, Octopus::ShardTracking::Dynamic)
   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
-
+  def current_shard
+    owner.current_shard
   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
-    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
-    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)