gistfile1.diff

diff --git a/lib/optimizations/associations/macro.rb b/lib/optimizations/associations/macro.rb
index 6d83873..f8a7d8c 100644
--- a/lib/optimizations/associations/macro.rb
+++ b/lib/optimizations/associations/macro.rb
@@ -25,6 +25,56 @@ module Scrooge
       end
       
       module SingletonMethods
+      
+        def self.extended( base )
+          eigen = class << base; self; end
+          eigen.instance_eval do
+            # Let :scrooge_callsite be a valid find option
+            #
+            remove_const(:VALID_FIND_OPTIONS)
+            const_set( :VALID_FIND_OPTIONS, [ :conditions, :include, :joins, :limit, :offset, :order, :select, :readonly, :group, :having, :from, :lock, :scrooge_callsite ] )
+          end
+          eigen.alias_method_chain :find, :scrooge
+          eigen.alias_method_chain :find_every, :scrooge
+        end
+      
+        # Let .find inject / remove callsite tracked associations
+        #
+        def find_with_scrooge(*args)
+          options = args.extract_options!
+          validate_find_options(options)
+          set_readonly_option!(options)
+
+          options[:scrooge_callsite] = callsite_signature( caller, options.except(:conditions, :limit, :offset) )  
+          options[:include] = scrooge_callsite(options[:scrooge_callsite]).preload( options[:include] )
+
+          case args.first
+            when :first then find_initial(options)
+            when :last  then find_last(options)
+            when :all   then find_every(options)
+            else             find_from_ids(args, options)
+          end
+        end
+      
+        # Override find_ever to pass along the callsite signature
+        #
+        def find_every_with_scrooge(options)
+          include_associations = merge_includes(scope(:find, :include), options[:include])
+
+          if include_associations.any? && references_eager_loaded_tables?(options)
+            records = find_with_associations(options)
+          else
+            records = find_by_sql(construct_finder_sql(options), options[:scrooge_callsite])
+            if include_associations.any?
+              preload_associations(records, include_associations)
+            end
+          end
+
+          records.each { |record| record.readonly! } if options[:readonly]
+
+          records
+        end      
+      
       end
       
       module InstanceMethods
diff --git a/lib/optimizations/columns/macro.rb b/lib/optimizations/columns/macro.rb
index d5cdae1..da94b15 100644
--- a/lib/optimizations/columns/macro.rb
+++ b/lib/optimizations/columns/macro.rb
@@ -80,9 +80,10 @@ module Scrooge
             end
           end        
         
-          # Find and instantiate as usual
+          # Find and instantiate as usual.Unset a registered callsite if given.
           #
-          def find_by_sql_without_scrooge( sql )
+          def find_by_sql_without_scrooge( sql, callsite = nil )
+            scrooge_unlink_callsite!( callsite ) if callsite
             connection.select_all(sanitize_sql(sql), "#{name} Load").collect! do |record|
               instantiate( UnscroogedAttributes.setup(record) )
             end
diff --git a/lib/scrooge.rb b/lib/scrooge.rb
index b1df0a6..40f9224 100644
--- a/lib/scrooge.rb
+++ b/lib/scrooge.rb
@@ -13,50 +13,6 @@ module ActiveRecord
     ScroogeCallsiteSample = 0..10
 
     class << self
-      
-      # Let :scrooge_callsite be a valid find option
-      #
-      remove_const(:VALID_FIND_OPTIONS)
-      VALID_FIND_OPTIONS = [ :conditions, :include, :joins, :limit, :offset,
-                             :order, :select, :readonly, :group, :having, :from, :lock, :scrooge_callsite ]
-
-      
-      # Let .find
-      #
-      def find(*args)
-        options = args.extract_options!
-        validate_find_options(options)
-        set_readonly_option!(options)
-
-        options[:scrooge_callsite] = callsite_signature( caller, options.except(:conditions, :limit, :offset) )  
-        options[:include] = scrooge_callsite(options[:scrooge_callsite]).preload( options[:include] )
-
-        case args.first
-          when :first then find_initial(options)
-          when :last  then find_last(options)
-          when :all   then find_every(options)
-          else             find_from_ids(args, options)
-        end
-      end      
-
-      # Override find_ever to pass along the callsite signature
-      #
-      def find_every(options)
-        include_associations = merge_includes(scope(:find, :include), options[:include])
-
-        if include_associations.any? && references_eager_loaded_tables?(options)
-          records = find_with_associations(options)
-        else
-          records = find_by_sql(construct_finder_sql(options), options[:scrooge_callsite])
-          if include_associations.any?
-            preload_associations(records, include_associations)
-          end
-        end
-
-        records.each { |record| record.readonly! } if options[:readonly]
-
-        records
-      end
 
       # Determine if a given SQL string is a candidate for callsite <=> columns
       # optimization.
@@ -65,7 +21,7 @@ module ActiveRecord
         if scope_with_scrooge?(sql)
           find_by_sql_with_scrooge(sql, callsite_signature)
         else
-          find_by_sql_without_scrooge(sql)
+          find_by_sql_without_scrooge(sql, callsite_signature)
         end
       end
       
@@ -90,6 +46,12 @@ module ActiveRecord
 
       private
 
+        # Removes a single callsite
+        #
+        def scrooge_unlink_callsite!( callsite_signature )
+          @@scrooge_callsites.delete(callsite_signature)  
+        end 
+
         # Initialize a callsite
         #
         def callsite( signature )

(END)