lib/octopus/association.rb in ar-octopus-0.0.5 vs lib/octopus/association.rb in ar-octopus-0.0.6

@@ -1,6 +1,30 @@
 module Octopus::Association
+  def has_many(association_id, options = {}, &extension)
+    default_octopus_opts(options)
+    super(association_id, options, &extension)
+  end
+
+  def has_and_belongs_to_many(association_id, options = {}, &extension)
+    default_octopus_opts(options)
+    super(association_id, options, &extension)    
+  end
+
+  def default_octopus_opts(options)
+    if options[:before_add].is_a?(Array)
+      options[:before_add] << :set_connection
+    else
+      options[:before_add] = :set_connection
+    end
+
+    if options[:before_remove].is_a?(Array)
+      options[:before_remove] << :set_connection
+    else
+      options[:before_remove] = :set_connection
+    end
+  end
+  
   def collection_reader_method(reflection, association_proxy_class)
     define_method(reflection.name) do |*params|
       force_reload = params.first unless params.empty?
       reload_connection() 
 
@@ -58,16 +82,16 @@
     end
   end
 
   def association_accessor_methods(reflection, association_proxy_class)
     define_method(reflection.name) do |*params|
-      force_reload = true
       reload_connection() 
+      force_reload = params.first unless params.empty?
       association = association_instance_get(reflection.name)
 
       if association.nil? || force_reload
         association = association_proxy_class.new(self, reflection)
-        retval = have_a_valid_shard? ? reflection.klass.uncached { self.class.connection_proxy.run_query_on_shard(self.current_shard) { association.reload } } : association.reload
+        retval = have_a_valid_shard? ? reflection.klass.uncached { self.class.connection_proxy.run_queries_on_shard(self.current_shard) { association.reload } } : association.reload
         if retval.nil? and association_proxy_class == ActiveRecord::Associations::BelongsToAssociation
           association_instance_set(reflection.name, nil)
           return nil
         end
         association_instance_set(reflection.name, association)