Merge pull request #6 from jasondoc3/add-one-off-tracking

Support tracking queries in a block

README.md

@@ -29,6 +29,24 @@ To start tracking, simply start your rails application server. When your server
 
 `sql_tracker` can also track sql queries when running rails tests (e.g. your controller or integration tests), it will dump the data after all the tests are finished.
 
+### Tracking Using a Block
+
+It is also possible to track queries executed within a block. This method uses a new subscriber to `sql.active_record` event notifications for each invocation. Results using this method are not saved to a file.
+
+```ruby
+query_data = SqlTracker.track do
+  # Run some active record queries
+end
+
+query_data.values
+# =>
+# [{
+#  :sql=>"SELECT * FROM users",
+#  :count=>1,
+#  :duration=>1.0,
+#  :source=>["app/models/user.rb:12"]
+# }]
+```
 
 ## Reporting
 

lib/sql_tracker.rb

@@ -8,11 +8,20 @@ def self.initialize!
 
     config = SqlTracker::Config.apply_defaults
     handler = SqlTracker::Handler.new(config)
-    ActiveSupport::Notifications.subscribe('sql.active_record', handler)
+    handler.subscribe
     @already_initialized = true
 
     at_exit { handler.save }
   end
+
+  def self.track
+    config = SqlTracker::Config.apply_defaults
+    handler = SqlTracker::Handler.new(config)
+    handler.subscribe
+    yield
+    handler.unsubscribe
+    handler.data
+  end
 end
 
 if defined?(::Rails) && ::Rails::VERSION::MAJOR.to_i >= 3

lib/sql_tracker/handler.rb

@@ -12,6 +12,21 @@ def initialize(config)
       @data = {} # {key: {sql:, count:, duration, source: []}, ...}
     end
 
+    def subscribe
+      @subscription ||= ActiveSupport::Notifications.subscribe(
+        'sql.active_record',
+        self
+      )
+    end
+
+    def unsubscribe
+      return unless @subscription
+
+      ActiveSupport::Notifications.unsubscribe(@subscription)
+
+      @subscription = nil
+    end
+
     def call(_name, started, finished, _id, payload)
       return unless @config.enabled
 

sql_tracker.gemspec

@@ -19,7 +19,7 @@ Gem::Specification.new do |spec|
   spec.executables   = ['sql_tracker']
   spec.require_paths = ['lib']
 
-  spec.add_development_dependency 'bundler', '~> 1.12'
+  spec.add_development_dependency 'bundler'
   spec.add_development_dependency 'rake', '~> 10.0'
   spec.add_development_dependency 'minitest', '~> 5.0'
   spec.add_development_dependency 'activesupport', '>= 3.0.0'

test/sql_tracker_test.rb

@@ -0,0 +1,37 @@
+require 'test_helper'
+
+module SqlTracker
+  class HandlerTest < Minitest::Test
+    def test_tracking_queries_with_a_block
+      config = SqlTracker::Config.apply_defaults
+      config.enabled = true
+      config.tracked_sql_command = %w(SELECT INSERT)
+
+      expected_queries = [
+        'SELECT * FROM users',
+        'insert into users VALUES (xxx)'
+      ]
+
+      query_data = SqlTracker.track do
+        expected_queries.each { |q| instrument_query(q) }
+        instrument_query('DELETE FROM users WHERE id = 1')
+      end
+
+      instrument_query('SELECT * FROM comments')
+
+      assert_equal(
+        expected_queries,
+        query_data.values.map { |v| v[:sql] }
+      )
+    end
+
+    private
+
+    def instrument_query(query)
+      ActiveSupport::Notifications.instrument(
+        'sql.active_record',
+        sql: query
+      )
+    end
+  end
+end