From 9da481fc3757dbcb7051b3e7e1493125155b21bc Mon Sep 17 00:00:00 2001
From: Jason Dougherty <jasondoc3@gmail.com>
Date: Thu, 5 Mar 2020 15:08:41 -0800
Subject: [PATCH 1/3] Support tracking queries in a block

---
 lib/sql_tracker.rb         | 11 ++++++++++-
 lib/sql_tracker/handler.rb | 15 +++++++++++++++
 test/sql_tracker_test.rb   | 37 +++++++++++++++++++++++++++++++++++++
 3 files changed, 62 insertions(+), 1 deletion(-)
 create mode 100644 test/sql_tracker_test.rb

diff --git a/lib/sql_tracker.rb b/lib/sql_tracker.rb
index 3e2770c..3a32187 100644
--- a/lib/sql_tracker.rb
+++ b/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
diff --git a/lib/sql_tracker/handler.rb b/lib/sql_tracker/handler.rb
index 6bc84a8..9797c36 100644
--- a/lib/sql_tracker/handler.rb
+++ b/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
 
diff --git a/test/sql_tracker_test.rb b/test/sql_tracker_test.rb
new file mode 100644
index 0000000..bdd5109
--- /dev/null
+++ b/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(
+        query_data.values.map { |v| v[:sql] },
+        expected_queries
+      )
+    end
+
+    private
+
+    def instrument_query(query)
+      ActiveSupport::Notifications.instrument(
+        'sql.active_record',
+        sql: query
+      )
+    end
+  end
+end

From c32c18601a0ce27c1c04c7509845ffd09c1f033f Mon Sep 17 00:00:00 2001
From: Jason Dougherty <jasondoc3@gmail.com>
Date: Fri, 20 Mar 2020 21:31:03 -0700
Subject: [PATCH 2/3] Relax bundler dependency

---
 sql_tracker.gemspec | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/sql_tracker.gemspec b/sql_tracker.gemspec
index e08429b..000e5b4 100644
--- a/sql_tracker.gemspec
+++ b/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'

From dde6f21649718028ffc8300f62de94655e64100c Mon Sep 17 00:00:00 2001
From: Jason Dougherty <jasondoc3@gmail.com>
Date: Sat, 21 Mar 2020 12:15:26 -0700
Subject: [PATCH 3/3] Update README.md

---
 README.md                | 18 ++++++++++++++++++
 test/sql_tracker_test.rb |  4 ++--
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index d6dc869..72ecaf3 100644
--- a/README.md
+++ b/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
 
diff --git a/test/sql_tracker_test.rb b/test/sql_tracker_test.rb
index bdd5109..c3f26bc 100644
--- a/test/sql_tracker_test.rb
+++ b/test/sql_tracker_test.rb
@@ -20,8 +20,8 @@ def test_tracking_queries_with_a_block
       instrument_query('SELECT * FROM comments')
 
       assert_equal(
-        query_data.values.map { |v| v[:sql] },
-        expected_queries
+        expected_queries,
+        query_data.values.map { |v| v[:sql] }
       )
     end