lib/octopus/association.rb in ar-octopus-0.8.0 vs lib/octopus/association.rb in ar-octopus-0.8.1

@@ -1,10 +1,43 @@
 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
@@ -62,9 +95,11 @@
     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)