let EXPLAIN use a thread locals registry [John J. Wang & Xavier Noria]

Closes #10198.

5 files changed, 70 insertions, 46 deletions

diff --git a/activerecord/lib/active_record/explain.rb b/activerecord/lib/active_record/explain.rb
index 15736575a2..e65dab07ba 100644
--- a/activerecord/lib/active_record/explain.rb
+++ b/activerecord/lib/active_record/explain.rb
@@ -1,22 +1,22 @@
 require 'active_support/lazy_load_hooks'
+require 'active_record/explain_registry'
 
 module ActiveRecord
   module Explain
-    # Relation#explain needs to be able to collect the queries.
+    # Executes the block with the collect flag enabled. Queries are collected
+    # asynchronously by the subscriber and returned.
     def collecting_queries_for_explain # :nodoc:
-      current = Thread.current
-      original, current[:available_queries_for_explain] = current[:available_queries_for_explain], []
+      ExplainRegistry.collect = true
       yield
-      return current[:available_queries_for_explain]
+      ExplainRegistry.queries
     ensure
-      # Note that the return value above does not depend on this assignment.
-      current[:available_queries_for_explain] = original
+      ExplainRegistry.reset
     end
 
     # Makes the adapter execute EXPLAIN for the tuples of queries and bindings.
     # Returns a formatted string ready to be logged.
     def exec_explain(queries) # :nodoc:
-      str = queries && queries.map do |sql, bind|
+      str = queries.map do |sql, bind|
         [].tap do |msg|
           msg << "EXPLAIN for: #{sql}"
           unless bind.empty?
@@ -31,6 +31,7 @@ module ActiveRecord
       def str.inspect
         self
       end
+
       str
     end
   end
diff --git a/activerecord/lib/active_record/explain_registry.rb b/activerecord/lib/active_record/explain_registry.rb
new file mode 100644
index 0000000000..b84a692b87
--- /dev/null
+++ b/activerecord/lib/active_record/explain_registry.rb
@@ -0,0 +1,30 @@
+require 'active_support/per_thread_registry'
+
+module ActiveRecord
+  # This is a thread locals registry for EXPLAIN. For example
+  #
+  #   ActiveRecord::ExplainRegistry.queries
+  #
+  # returns the collected queries local to the current thread.
+  #
+  # See the documentation of <tt>ActiveSupport::PerThreadRegistry</tt>
+  # for further details.
+  class ExplainRegistry
+    extend ActiveSupport::PerThreadRegistry
+
+    attr_accessor :queries, :collect
+
+    def initialize
+      reset
+    end
+
+    def collect?
+      @collect
+    end
+
+    def reset
+      @collect = false
+      @queries = []
+    end
+  end
+end
\ No newline at end of file
diff --git a/activerecord/lib/active_record/explain_subscriber.rb b/activerecord/lib/active_record/explain_subscriber.rb
index 0f927496fb..a3bc56d600 100644
--- a/activerecord/lib/active_record/explain_subscriber.rb
+++ b/activerecord/lib/active_record/explain_subscriber.rb
@@ -1,4 +1,5 @@
 require 'active_support/notifications'
+require 'active_record/explain_registry'
 
 module ActiveRecord
   class ExplainSubscriber # :nodoc:
@@ -7,8 +8,8 @@ module ActiveRecord
     end
 
     def finish(name, id, payload)
-      if queries = Thread.current[:available_queries_for_explain]
-        queries << payload.values_at(:sql, :binds) unless ignore_payload?(payload)
+      if ExplainRegistry.collect? && !ignore_payload?(payload)
+        ExplainRegistry.queries << payload.values_at(:sql, :binds)
       end
     end
 
diff --git a/activerecord/test/cases/adapters/postgresql/explain_test.rb b/activerecord/test/cases/adapters/postgresql/explain_test.rb
index 619d581d5f..0b61f61572 100644
--- a/activerecord/test/cases/adapters/postgresql/explain_test.rb
+++ b/activerecord/test/cases/adapters/postgresql/explain_test.rb
@@ -22,13 +22,6 @@ module ActiveRecord
           assert_match %(EXPLAIN for: SELECT "audit_logs".* FROM "audit_logs"  WHERE "audit_logs"."developer_id" IN (1)), explain
           assert_match %(Seq Scan on audit_logs), explain
         end
-
-        def test_dont_explain_for_set_search_path
-          queries = Thread.current[:available_queries_for_explain] = []
-          ActiveRecord::Base.connection.schema_search_path = "public"
-          assert queries.empty?
-        end
-
       end
     end
   end
diff --git a/activerecord/test/cases/explain_subscriber_test.rb b/activerecord/test/cases/explain_subscriber_test.rb
index b425967678..fb53a92c89 100644
--- a/activerecord/test/cases/explain_subscriber_test.rb
+++ b/activerecord/test/cases/explain_subscriber_test.rb
@@ -1,55 +1,54 @@
 require 'cases/helper'
+require 'active_record/explain_subscriber'
+require 'active_record/explain_registry'
 
 if ActiveRecord::Base.connection.supports_explain?
   class ExplainSubscriberTest < ActiveRecord::TestCase
     SUBSCRIBER = ActiveRecord::ExplainSubscriber.new
 
-    def test_collects_nothing_if_available_queries_for_explain_is_nil
-      with_queries(nil) do
-        SUBSCRIBER.finish(nil, nil, {})
-        assert_nil Thread.current[:available_queries_for_explain]
-      end
+    def setup
+      ActiveRecord::ExplainRegistry.reset
+      ActiveRecord::ExplainRegistry.collect = true
     end
 
     def test_collects_nothing_if_the_payload_has_an_exception
-      with_queries([]) do |queries|
-        SUBSCRIBER.finish(nil, nil, :exception => Exception.new)
-        assert queries.empty?
-      end
+      SUBSCRIBER.finish(nil, nil, exception: Exception.new)
+      assert queries.empty?
     end
 
     def test_collects_nothing_for_ignored_payloads
-      with_queries([]) do |queries|
-        ActiveRecord::ExplainSubscriber::IGNORED_PAYLOADS.each do |ip|
-          SUBSCRIBER.finish(nil, nil, :name => ip)
-        end
-        assert queries.empty?
+      ActiveRecord::ExplainSubscriber::IGNORED_PAYLOADS.each do |ip|
+        SUBSCRIBER.finish(nil, nil, name: ip)
       end
+      assert queries.empty?
+    end
+
+    def test_collects_nothing_if_collect_is_false
+      ActiveRecord::ExplainRegistry.collect = false
+      SUBSCRIBER.finish(nil, nil, name: 'SQL', sql: 'select 1 from users', binds: [1, 2])
+      assert queries.empty?
     end
 
     def test_collects_pairs_of_queries_and_binds
       sql   = 'select 1 from users'
       binds = [1, 2]
-      with_queries([]) do |queries|
-        SUBSCRIBER.finish(nil, nil, :name => 'SQL', :sql => sql, :binds => binds)
-        assert_equal 1, queries.size
-        assert_equal sql, queries[0][0]
-        assert_equal binds, queries[0][1]
-      end
+      SUBSCRIBER.finish(nil, nil, name: 'SQL', sql: sql, binds: binds)
+      assert_equal 1, queries.size
+      assert_equal sql, queries[0][0]
+      assert_equal binds, queries[0][1]
     end
 
-    def test_collects_nothing_if_unexplained_sqls
-      with_queries([]) do |queries|
-        SUBSCRIBER.finish(nil, nil, :name => 'SQL', :sql => 'SHOW max_identifier_length')
-        assert queries.empty?
-      end
+    def test_collects_nothing_if_the_statement_is_not_whitelisted
+      SUBSCRIBER.finish(nil, nil, name: 'SQL', sql: 'SHOW max_identifier_length')
+      assert queries.empty?
+    end
+
+    def teardown
+      ActiveRecord::ExplainRegistry.reset
     end
 
-    def with_queries(queries)
-      Thread.current[:available_queries_for_explain] = queries
-      yield queries
-    ensure
-      Thread.current[:available_queries_for_explain] = nil
+    def queries
+      ActiveRecord::ExplainRegistry.queries
     end
   end
 end