lib/bullet/mongoid3x.rb in bullet-4.9.0 vs lib/bullet/mongoid3x.rb in bullet-4.10.0

@@ -39,16 +39,17 @@
           origin_eager_load(docs)
         end
       end
 
       ::Mongoid::Relations::Accessors.class_eval do
-        alias_method :origin_set_relation, :set_relation
+        alias_method :origin_get_relation, :get_relation
 
-        def set_relation(name, relation)
-          if relation && relation.metadata.macro !~ /embed/
+        def get_relation(name, metadata, reload = false)
+          result = origin_get_relation(name, metadata, reload)
+          if metadata.macro !~ /embed/
             Bullet::Detector::NPlusOneQuery.call_association(self, name)
           end
-          origin_set_relation(name, relation)
+          result
         end
       end
     end
   end
 end