bullet 5.7.0 → 5.8.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: 7c9650d2dbc870d2ac04be4ca0e19723ef6a85c2
-  data.tar.gz: 35839d605fb1e11687836dbcbe12572b1eb07775
+  metadata.gz: 0b65855ae5365567a18d1fb6ffdee5046b8cec35
+  data.tar.gz: 8df6b7afd18400dd2fe9bf39ec4a537fa04228af
 SHA512:
-  metadata.gz: 5064948ab9db44cc9b4082397990a3d91e93dc6986db2767a528a4a05734b1d1ba9231af337624d79d7e312d13e075afb53217191d017a3c511a0889c436d0a0
-  data.tar.gz: 2c12bd045e5e41b4db0761b07562a763d3be089231d5e420fc3d7d8d27a40f184216c9090667883a606e35fcdfd2eda3f276838e60696c642bff0939ab35e85f
+  metadata.gz: 212cf02f1520dfaf752f794e42e9b2611fa0eca8b786bf5ef281c36c416ddf24016a524029614447c6241ebb734a2491ab554fa505544d7a7bc6f2e5b767527d
+  data.tar.gz: d2d71ceb54033f58cb6d675ea88268f130a0ea8bcfd7b4987486e374a2d2a3fc9169691ff20a9ca0fd9cf2787656dacc19649a8473879e9d1d68484597ccc18b

data/CHANGELOG.md

@@ -1,4 +1,19 @@
-# Next Release
+## Next Release
+
+## 5.8.0 (10/29/2018)
+
+* Fix through reflection for rails 5.x
+* Fix false positive in after_save/after_create callbacks
+* Don't triger a preload error on "manual" preloads
+* Avoid Bullet from making extra queries in mongoid6
+* Support option for #first and #last on mongoid6.x
+* Fix duplicate logs in mongoid 4.x and 5.x version
+* Use caller for ruby 1.9 while caller_locations for 2.0+
+* Extend stacktrace matching for sub-file precision
+* Exclude configured bundler path in addition to '/vendor'
+* Fix `caller_path` in `excluded_stacktrace_path`
+* Update `uniform_notifier` dependency to add Sentry support
+* Integrate awesomecode.io and refactor code
 
 ## 5.7.0 (12/03/2017)
 

data/Gemfile.rails-5.2

@@ -2,7 +2,7 @@ source "https://rubygems.org"
 
 gemspec
 
-gem 'rails', '~> 5.2.0.beta1'
+gem 'rails', '~> 5.2.0'
 gem 'sqlite3'
 gem 'activerecord-jdbcsqlite3-adapter', platforms: [:jruby]
 gem 'activerecord-import'

data/Guardfile

@@ -1,7 +1,7 @@
 # A sample Guardfile
 # More info at https://github.com/guard/guard#readme
 
-guard 'rspec', :version => 2, :all_after_pass => false, :all_on_start => false, :cli => '--color --format nested --fail-fast' do
+guard 'rspec', version: 2, all_after_pass: false, all_on_start: false, cli: '--color --format nested --fail-fast' do
   watch(%r{^spec/.+_spec\.rb$})
   watch(%r{^lib/(.+)\.rb$})     { |m| "spec/#{m[1]}_spec.rb" }
   watch('spec/spec_helper.rb')  { 'spec' }

data/README.md

@@ -2,7 +2,7 @@
 
 [![Gem Version](https://badge.fury.io/rb/bullet.png)](http://badge.fury.io/rb/bullet)
 [![Build Status](https://secure.travis-ci.org/flyerhzm/bullet.png)](http://travis-ci.org/flyerhzm/bullet)
-<a href="https://codeclimate.com/github/flyerhzm/bullet"><img src="https://codeclimate.com/github/flyerhzm/bullet.png" /></a>
+[![AwesomeCode Status for flyerhzm/bullet](https://awesomecode.io/projects/6755235b-e2c1-459e-bf92-b8b13d0c0472/status)](https://awesomecode.io/repos/flyerhzm/bullet)
 [![Coderwall Endorse](http://api.coderwall.com/flyerhzm/endorsecount.png)](http://coderwall.com/flyerhzm)
 
 The Bullet gem is designed to help you increase your application's performance by reducing the number of queries it makes. It will watch your queries while you develop your application and notify you when you should add eager loading (N+1 queries), when you're using eager loading that isn't necessary and when you should use counter cache.
@@ -48,6 +48,7 @@ Bullet won't do ANYTHING unless you tell it to explicitly. Append to
 ```ruby
 config.after_initialize do
   Bullet.enable = true
+  Bullet.sentry = true
   Bullet.alert = true
   Bullet.bullet_logger = true
   Bullet.console = true
@@ -63,7 +64,7 @@ config.after_initialize do
   Bullet.rollbar = true
   Bullet.add_footer = true
   Bullet.stacktrace_includes = [ 'your_gem', 'your_middleware' ]
-  Bullet.stacktrace_excludes = [ 'their_gem', 'their_middleware' ]
+  Bullet.stacktrace_excludes = [ 'their_gem', 'their_middleware', ['my_file.rb', 'my_method'], ['my_file.rb', 16..20] ]
   Bullet.slack = { webhook_url: 'http://some.slack.url', channel: '#default', username: 'notifier' }
 end
 ```
@@ -82,9 +83,12 @@ The code above will enable all of the Bullet notification systems:
 * `Bullet.bugsnag`: add notifications to bugsnag
 * `Bullet.airbrake`: add notifications to airbrake
 * `Bullet.rollbar`: add notifications to rollbar
+* `Bullet.sentry`: add notifications to sentry
 * `Bullet.add_footer`: adds the details in the bottom left corner of the page. Double click the footer or use close button to hide footer.
 * `Bullet.stacktrace_includes`: include paths with any of these substrings in the stack trace, even if they are not in your main app
 * `Bullet.stacktrace_excludes`: ignore paths with any of these substrings in the stack trace, even if they are not in your main app.
+   Each item can be a string (match substring), a regex, or an array where the first item is a path to match, and the second
+   item is a line number, a Range of line numbers, or a (bare) method name, to exclude only particular lines in a file.
 * `Bullet.slack`: add notifications to slack
 * `Bullet.raise`: raise errors, useful for making your specs fail unless they have optimized queries
 
@@ -124,10 +128,11 @@ class ApplicationController < ActionController::Base
   around_action :skip_bullet
 
   def skip_bullet
+    previous_value = Bullet.enable?
     Bullet.enable = false
     yield
   ensure
-    Bullet.enable = true
+    Bullet.enable = previous_value
   end
 end
 ```
@@ -178,7 +183,7 @@ If you find Bullet does not work for you, *please disable your browser's cache*.
 
 ### Profile a job
 
-The Bullet gem uses rack middleware to profile requests. If you want to use Bullet without an http server, like to profile a job, you can use use profile method and fetch warnings
+The Bullet gem uses rack middleware to profile requests. If you want to use Bullet without an http server, like to profile a job, you can use the profile method and fetch warnings
 
 ```ruby
 Bullet.profile do

data/Rakefile

@@ -1,4 +1,4 @@
-$LOAD_PATH.unshift File.expand_path('../lib', __FILE__)
+$LOAD_PATH.unshift File.expand_path('lib', __dir__)
 require 'bundler'
 Bundler.setup
 
@@ -12,11 +12,11 @@ task :build do
   system 'gem build bullet.gemspec'
 end
 
-task :install => :build do
+task install: :build do
   system "sudo gem install bullet-#{Bullet::VERSION}.gem"
 end
 
-task :release => :build do
+task release: :build do
   puts "Tagging #{Bullet::VERSION}..."
   system "git tag -a #{Bullet::VERSION} -m 'Tagging #{Bullet::VERSION}'"
   puts 'Pushing to Github...'
@@ -30,7 +30,7 @@ RSpec::Core::RakeTask.new(:spec) do |spec|
 end
 
 RSpec::Core::RakeTask.new('spec:progress') do |spec|
-  spec.rspec_opts = %w(--format progress)
+  spec.rspec_opts = %w[--format progress]
   spec.pattern = 'spec/**/*_spec.rb'
 end
 
@@ -48,4 +48,4 @@ rescue LoadError
   puts 'RDocTask is not supported for this platform'
 end
 
-task :default => :spec
+task default: :spec

data/bullet.gemspec

@@ -1,5 +1,7 @@
-lib = File.expand_path('../lib/', __FILE__)
-$:.unshift lib unless $:.include?(lib)
+# frozen_string_literal: true
+
+lib = File.expand_path('lib', __dir__)
+$LOAD_PATH.unshift lib unless $LOAD_PATH.include?(lib)
 
 require 'bullet/version'
 
@@ -9,19 +11,22 @@ Gem::Specification.new do |s|
   s.platform    = Gem::Platform::RUBY
   s.authors     = ['Richard Huang']
   s.email       = ['flyerhzm@gmail.com']
-  s.homepage    = 'http://github.com/flyerhzm/bullet'
+  s.homepage    = 'https://github.com/flyerhzm/bullet'
   s.summary     = 'help to kill N+1 queries and unused eager loading.'
   s.description = 'help to kill N+1 queries and unused eager loading.'
+  s.metadata    = {
+    'changelog_uri' => 'https://github.com/flyerhzm/bullet/blob/master/CHANGELOG.md',
+    'source_code_uri' => 'https://github.com/flyerhzm/bullet'
+  }
 
-  s.license     = 'MIT'
+  s.license = 'MIT'
 
   s.required_rubygems_version = '>= 1.3.6'
 
   s.add_runtime_dependency 'activesupport', '>= 3.0.0'
-  s.add_runtime_dependency 'uniform_notifier', '~> 1.10.0'
+  s.add_runtime_dependency 'uniform_notifier', '~> 1.11'
 
   s.files         = `git ls-files`.split("\n")
   s.test_files    = `git ls-files -- {test,spec,features}/*`.split("\n")
   s.require_paths = ['lib']
 end
-

data/lib/bullet/active_record4.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Bullet
   module ActiveRecord
     def self.enable
@@ -46,21 +48,12 @@ module Bullet
       end
 
       ::ActiveRecord::Persistence.class_eval do
-        def save_with_bullet(*args, &proc)
-          was_new_record = new_record?
-          save_without_bullet(*args, &proc).tap do |result|
-            Bullet::Detector::NPlusOneQuery.add_impossible_object(self) if result && was_new_record
-          end
-        end
-        alias_method_chain :save, :bullet
-
-        def save_with_bullet!(*args, &proc)
-          was_new_record = new_record?
-          save_without_bullet!(*args, &proc).tap do |result|
-            Bullet::Detector::NPlusOneQuery.add_impossible_object(self) if result && was_new_record
+        def _create_record_with_bullet(*args)
+          _create_record_without_bullet(*args).tap do
+            Bullet::Detector::NPlusOneQuery.add_impossible_object(self)
           end
         end
-        alias_method_chain :save!, :bullet
+        alias_method_chain :_create_record, :bullet
       end
 
       ::ActiveRecord::Associations::Preloader.class_eval do
@@ -72,6 +65,7 @@ module Bullet
           if Bullet.start?
             records = [records].flatten.compact.uniq
             return if records.empty?
+
             records.each do |record|
               Bullet::Detector::Association.add_object_associations(record, associations)
             end
@@ -105,7 +99,7 @@ module Bullet
           records = origin_instantiate(rows)
 
           if Bullet.start?
-            @bullet_eager_loadings.each do |klazz, eager_loadings_hash|
+            @bullet_eager_loadings.each do |_klazz, eager_loadings_hash|
               objects = eager_loadings_hash.keys
               Bullet::Detector::UnusedEagerLoading.add_eager_loadings(objects, eager_loadings_hash[objects.first].to_a)
             end
@@ -187,6 +181,7 @@ module Bullet
       ::ActiveRecord::Associations::HasManyAssociation.class_eval do
         alias_method :origin_has_cached_counter?, :has_cached_counter?
 
+        # rubocop:disable Style/MethodCallWithoutArgsParentheses
         def has_cached_counter?(reflection = reflection())
           result = origin_has_cached_counter?(reflection)
           if Bullet.start? && !result
@@ -194,6 +189,15 @@ module Bullet
           end
           result
         end
+        # rubocop:enable Style/MethodCallWithoutArgsParentheses
+      end
+
+      ::ActiveRecord::Associations::BelongsToAssociation.class_eval do
+        def writer_with_bullet(record)
+          Bullet::Detector::Association.add_object_associations(owner, reflection.name) if Bullet.start?
+          writer_without_bullet(record)
+        end
+        alias_method_chain :writer, :bullet
       end
     end
   end

data/lib/bullet/active_record41.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Bullet
   module ActiveRecord
     def self.enable
@@ -48,21 +50,12 @@ module Bullet
       end
 
       ::ActiveRecord::Persistence.class_eval do
-        def save_with_bullet(*args, &proc)
-          was_new_record = new_record?
-          save_without_bullet(*args, &proc).tap do |result|
-            Bullet::Detector::NPlusOneQuery.add_impossible_object(self) if result && was_new_record
-          end
-        end
-        alias_method_chain :save, :bullet
-
-        def save_with_bullet!(*args, &proc)
-          was_new_record = new_record?
-          save_without_bullet!(*args, &proc).tap do |result|
-            Bullet::Detector::NPlusOneQuery.add_impossible_object(self) if result && was_new_record
+        def _create_record_with_bullet(*args)
+          _create_record_without_bullet(*args).tap do
+            Bullet::Detector::NPlusOneQuery.add_impossible_object(self)
           end
         end
-        alias_method_chain :save!, :bullet
+        alias_method_chain :_create_record, :bullet
       end
 
       ::ActiveRecord::Associations::Preloader.class_eval do
@@ -87,6 +80,7 @@ module Bullet
         alias_method :origin_find_with_associations, :find_with_associations
         def find_with_associations
           return origin_find_with_associations { |r| yield r } if block_given?
+
           records = origin_find_with_associations
           if Bullet.start?
             associations = (eager_load_values + includes_values).uniq
@@ -108,7 +102,7 @@ module Bullet
           records = origin_instantiate(result_set, aliases)
 
           if Bullet.start?
-            @bullet_eager_loadings.each do |klazz, eager_loadings_hash|
+            @bullet_eager_loadings.each do |_klazz, eager_loadings_hash|
               objects = eager_loadings_hash.keys
               Bullet::Detector::UnusedEagerLoading.add_eager_loadings(objects, eager_loadings_hash[objects.first].to_a)
             end
@@ -185,6 +179,14 @@ module Bullet
           origin_count_records
         end
       end
+
+      ::ActiveRecord::Associations::BelongsToAssociation.class_eval do
+        def writer_with_bullet(record)
+          Bullet::Detector::Association.add_object_associations(owner, reflection.name) if Bullet.start?
+          writer_without_bullet(record)
+        end
+        alias_method_chain :writer, :bullet
+      end
     end
   end
 end

data/lib/bullet/active_record42.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Bullet
   module ActiveRecord
     def self.enable
@@ -42,21 +44,12 @@ module Bullet
       end
 
       ::ActiveRecord::Persistence.class_eval do
-        def save_with_bullet(*args, &proc)
-          was_new_record = new_record?
-          save_without_bullet(*args, &proc).tap do |result|
-            Bullet::Detector::NPlusOneQuery.add_impossible_object(self) if result && was_new_record
+        def _create_record_with_bullet(*args)
+          _create_record_without_bullet(*args).tap do
+            Bullet::Detector::NPlusOneQuery.add_impossible_object(self)
           end
         end
-        alias_method_chain :save, :bullet
-
-        def save_with_bullet!(*args, &proc)
-          was_new_record = new_record?
-          save_without_bullet!(*args, &proc).tap do |result|
-            Bullet::Detector::NPlusOneQuery.add_impossible_object(self) if result && was_new_record
-          end
-        end
-        alias_method_chain :save!, :bullet
+        alias_method_chain :_create_record, :bullet
       end
 
       ::ActiveRecord::Relation.class_eval do
@@ -102,6 +95,7 @@ module Bullet
         alias_method :origin_find_with_associations, :find_with_associations
         def find_with_associations
           return origin_find_with_associations { |r| yield r } if block_given?
+
           records = origin_find_with_associations
           if Bullet.start?
             associations = (eager_load_values + includes_values).uniq
@@ -124,7 +118,7 @@ module Bullet
           records = origin_instantiate(result_set, aliases)
 
           if Bullet.start?
-            @bullet_eager_loadings.each do |klazz, eager_loadings_hash|
+            @bullet_eager_loadings.each do |_klazz, eager_loadings_hash|
               objects = eager_loadings_hash.keys
               Bullet::Detector::UnusedEagerLoading.add_eager_loadings(objects, eager_loadings_hash[objects.first].to_a)
             end
@@ -138,14 +132,14 @@ module Bullet
               parent.children.each do |node|
                 key = aliases.column_alias(node, node.primary_key)
                 id = row[key]
-                if id.nil?
-                  associations = node.reflection.name
-                  Bullet::Detector::Association.add_object_associations(ar_parent, associations)
-                  Bullet::Detector::NPlusOneQuery.call_association(ar_parent, associations)
-                  @bullet_eager_loadings[ar_parent.class] ||= {}
-                  @bullet_eager_loadings[ar_parent.class][ar_parent] ||= Set.new
-                  @bullet_eager_loadings[ar_parent.class][ar_parent] << associations
-                end
+                next unless id.nil?
+
+                associations = node.reflection.name
+                Bullet::Detector::Association.add_object_associations(ar_parent, associations)
+                Bullet::Detector::NPlusOneQuery.call_association(ar_parent, associations)
+                @bullet_eager_loadings[ar_parent.class] ||= {}
+                @bullet_eager_loadings[ar_parent.class][ar_parent] ||= Set.new
+                @bullet_eager_loadings[ar_parent.class][ar_parent] << associations
               end
             end
           end
@@ -246,6 +240,14 @@ module Bullet
           origin_count_records
         end
       end
+
+      ::ActiveRecord::Associations::BelongsToAssociation.class_eval do
+        def writer_with_bullet(record)
+          Bullet::Detector::Association.add_object_associations(owner, reflection.name) if Bullet.start?
+          writer_without_bullet(record)
+        end
+        alias_method_chain :writer, :bullet
+      end
     end
   end
 end

data/lib/bullet/active_record5.rb

@@ -1,16 +1,11 @@
+# frozen_string_literal: true
+
 module Bullet
   module SaveWithBulletSupport
-    def save(*args)
-      was_new_record = new_record?
-      super(*args).tap do |result|
-        Bullet::Detector::NPlusOneQuery.add_impossible_object(self) if result && was_new_record
-      end
-    end
-
-    def save!(*args)
-      was_new_record = new_record?
-      super(*args).tap do |result|
-        Bullet::Detector::NPlusOneQuery.add_impossible_object(self) if result && was_new_record
+    def _create_record(*)
+      super do
+        Bullet::Detector::NPlusOneQuery.add_impossible_object(self)
+        yield(self) if block_given?
       end
     end
   end
@@ -18,7 +13,7 @@ module Bullet
   module ActiveRecord
     def self.enable
       require 'active_record'
-      ::ActiveRecord::Base.extend(Module.new {
+      ::ActiveRecord::Base.extend(Module.new do
         def find_by_sql(sql, binds = [], preparable: nil, &block)
           result = super
           if Bullet.start?
@@ -37,11 +32,11 @@ module Bullet
           end
           result
         end
-      })
+      end)
 
       ::ActiveRecord::Base.prepend(SaveWithBulletSupport)
 
-      ::ActiveRecord::Relation.prepend(Module.new {
+      ::ActiveRecord::Relation.prepend(Module.new do
         # if select a collection of objects, then these objects have possible to cause N+1 query.
         # if select only one object, then the only one object has impossible to cause N+1 query.
         def records
@@ -59,9 +54,9 @@ module Bullet
           end
           result
         end
-      })
+      end)
 
-      ::ActiveRecord::Associations::Preloader.prepend(Module.new {
+      ::ActiveRecord::Associations::Preloader.prepend(Module.new do
         def preloaders_for_one(association, records, scope)
           if Bullet.start?
             records.compact!
@@ -74,12 +69,13 @@ module Bullet
           end
           super
         end
-      })
+      end)
 
-      ::ActiveRecord::FinderMethods.prepend(Module.new {
+      ::ActiveRecord::FinderMethods.prepend(Module.new do
         # add includes in scope
         def find_with_associations
           return super { |r| yield r } if block_given?
+
           records = super
           if Bullet.start?
             associations = (eager_load_values + includes_values).uniq
@@ -90,20 +86,37 @@ module Bullet
           end
           records
         end
-      })
-
-      ::ActiveRecord::Associations::JoinDependency.prepend(Module.new {
-        def instantiate(result_set, aliases)
-          @bullet_eager_loadings = {}
-          records = super
-
-          if Bullet.start?
-            @bullet_eager_loadings.each do |klazz, eager_loadings_hash|
-              objects = eager_loadings_hash.keys
-              Bullet::Detector::UnusedEagerLoading.add_eager_loadings(objects, eager_loadings_hash[objects.first].to_a)
+      end)
+
+      ::ActiveRecord::Associations::JoinDependency.prepend(Module.new do
+        if ::ActiveRecord::Associations::JoinDependency.instance_method(:instantiate).parameters.last[0] == :block
+          # ActiveRecord >= 5.1.5
+          def instantiate(result_set, &block)
+            @bullet_eager_loadings = {}
+            records = super
+
+            if Bullet.start?
+              @bullet_eager_loadings.each do |_klazz, eager_loadings_hash|
+                objects = eager_loadings_hash.keys
+                Bullet::Detector::UnusedEagerLoading.add_eager_loadings(objects, eager_loadings_hash[objects.first].to_a)
+              end
             end
+            records
+          end
+        else
+          # ActiveRecord <= 5.1.4
+          def instantiate(result_set, aliases)
+            @bullet_eager_loadings = {}
+            records = super
+
+            if Bullet.start?
+              @bullet_eager_loadings.each do |_klazz, eager_loadings_hash|
+                objects = eager_loadings_hash.keys
+                Bullet::Detector::UnusedEagerLoading.add_eager_loadings(objects, eager_loadings_hash[objects.first].to_a)
+              end
+            end
+            records
           end
-          records
         end
 
         def construct(ar_parent, parent, row, rs, seen, model_cache, aliases)
@@ -112,14 +125,14 @@ module Bullet
               parent.children.each do |node|
                 key = aliases.column_alias(node, node.primary_key)
                 id = row[key]
-                if id.nil?
-                  associations = node.reflection.name
-                  Bullet::Detector::Association.add_object_associations(ar_parent, associations)
-                  Bullet::Detector::NPlusOneQuery.call_association(ar_parent, associations)
-                  @bullet_eager_loadings[ar_parent.class] ||= {}
-                  @bullet_eager_loadings[ar_parent.class][ar_parent] ||= Set.new
-                  @bullet_eager_loadings[ar_parent.class][ar_parent] << associations
-                end
+                next unless id.nil?
+
+                associations = node.reflection.name
+                Bullet::Detector::Association.add_object_associations(ar_parent, associations)
+                Bullet::Detector::NPlusOneQuery.call_association(ar_parent, associations)
+                @bullet_eager_loadings[ar_parent.class] ||= {}
+                @bullet_eager_loadings[ar_parent.class][ar_parent] ||= Set.new
+                @bullet_eager_loadings[ar_parent.class][ar_parent] << associations
               end
             end
           end
@@ -142,19 +155,28 @@ module Bullet
 
           result
         end
-      })
+      end)
 
-      ::ActiveRecord::Associations::CollectionAssociation.prepend(Module.new {
+      ::ActiveRecord::Associations::CollectionAssociation.prepend(Module.new do
         def load_target
           records = super
 
           if Bullet.start?
-            if self.is_a? ::ActiveRecord::Associations::ThroughAssociation
-              Bullet::Detector::NPlusOneQuery.call_association(owner, through_reflection.name)
-              association = self.owner.association self.through_reflection.name
+            if is_a? ::ActiveRecord::Associations::ThroughAssociation
+              refl = reflection.through_reflection
+              Bullet::Detector::NPlusOneQuery.call_association(owner, refl.name)
+              association = owner.association refl.name
               Array(association.target).each do |through_record|
                 Bullet::Detector::NPlusOneQuery.call_association(through_record, source_reflection.name)
               end
+
+              if refl.through_reflection?
+                while refl.through_reflection?
+                  refl = refl.through_reflection
+                end
+
+                Bullet::Detector::NPlusOneQuery.call_association(owner, refl.name)
+              end
             end
             Bullet::Detector::NPlusOneQuery.call_association(owner, reflection.name) unless @inversed
             if records.first.class.name !~ /^HABTM_/
@@ -183,9 +205,9 @@ module Bullet
           end
           super
         end
-      })
+      end)
 
-      ::ActiveRecord::Associations::SingularAssociation.prepend(Module.new {
+      ::ActiveRecord::Associations::SingularAssociation.prepend(Module.new do
         # call has_one and belongs_to associations
         def target
           result = super()
@@ -201,9 +223,9 @@ module Bullet
           end
           result
         end
-      })
+      end)
 
-      ::ActiveRecord::Associations::HasManyAssociation.prepend(Module.new {
+      ::ActiveRecord::Associations::HasManyAssociation.prepend(Module.new do
         def empty?
           result = super
           if Bullet.start? && !reflection.has_cached_counter?
@@ -214,12 +236,19 @@ module Bullet
 
         def count_records
           result = reflection.has_cached_counter?
-          if Bullet.start? && !result && !self.is_a?(::ActiveRecord::Associations::ThroughAssociation)
+          if Bullet.start? && !result && !is_a?(::ActiveRecord::Associations::ThroughAssociation)
             Bullet::Detector::CounterCache.add_counter_cache(owner, reflection.name)
           end
           super
         end
-      })
+      end)
+
+      ::ActiveRecord::Associations::BelongsToAssociation.prepend(Module.new do
+        def writer(record)
+          Bullet::Detector::Association.add_object_associations(owner, reflection.name) if Bullet.start?
+          super
+        end
+      end)
     end
   end
 end