db explain plan refactor

CHANGELOG.md

@@ -2,7 +2,7 @@
 
 ## dev
 
-Version <dev> introduces instrumentation for the aws-sdk-lambda gem and allows users to opt-in to adding labels to logs.
+Version <dev> introduces instrumentation for the aws-sdk-lambda gem, allows users to opt-in to adding labels to logs, and fixes a bug with explain plans on Rails 7.2+.
 
 - **Feature: Instrumentation for aws-sdk-lambda**
 
@@ -12,6 +12,10 @@ Version <dev> introduces instrumentation for the aws-sdk-lambda gem and allows u
 
   The Ruby agent now allows you to opt-in to adding your custom tags (labels) to agent-forwarded logs. With custom tags on logs, platform engineers can easily filter, search, and correlate log data for faster and more efficient troubleshooting, improved performance, and optimized resource utilization. [PR#2925](https://github.com/newrelic/newrelic-ruby-agent/pull/2925)
 
+- **Bugfix: Record explain plan traces on Rails 7.2+**
+
+  Rails 7.2 removed adapter-specific connection methods (ex. `ActiveRecord::Base.postgresql_connection`) and replaced them with `ActiveRecord::Base.with_connection`. Our explain plan feature relies on making a connection to the database to create an explain plan trace. Due to a bug in our tests, we missed this regression. Now, the agent uses the new method to fetch explain plans on Rails 7.2+. Thank you, [@gsar](https://github.com/gsar) and [@gstark](https://github.com/gstark) for bringing this to our attention! [Issue#2922](https://github.com/newrelic/newrelic-ruby-agent/issues/2922) [PR#2940](https://github.com/newrelic/newrelic-ruby-agent/pull/2940)
+
 ## v9.15.0
 
 Version 9.15.0 updates View Component instrumentation to use a default metric name when one is unavailable, adds a configuration option to associate the AWS account ID with the DynamoDB calls from the AWS SDK, resolves a bug in rdkafka instrumentation when using the karafka-rdkafka gem, resolves a bug in the ruby-kafka instrumentation, fixes a bug with Grape instrumentation, and addresses a bug preventing the agent from running in serverless mode in an AWS Lambda layer.

lib/new_relic/agent/database.rb

@@ -90,6 +90,42 @@ def get_connection(config, &connector)
         ConnectionManager.instance.get_connection(config, &connector)
       end
 
+      def explain_this(statement, use_execute = false)
+        if supports_with_connection?
+          explain_this_using_with_connection(statement)
+        else
+          explain_this_using_adapter_connection(statement, use_execute)
+        end
+      rescue => e
+        NewRelic::Agent.logger.error("Couldn't fetch the explain plan for statement: #{e}")
+      end
+
+      def explain_this_using_with_connection(statement)
+        ::ActiveRecord::Base.with_connection do |conn|
+          conn.exec_query("EXPLAIN #{statement.sql}", "Explain #{statement.name}", statement.binds)
+        end
+      end
+
+      def explain_this_using_adapter_connection(statement, use_execute)
+        connection = get_connection(statement.config) do
+          ::ActiveRecord::Base.send(:"#{statement.config[:adapter]}_connection", statement.config)
+        end
+
+        if use_execute
+          connection.execute("EXPLAIN #{statement.sql}")
+        else
+          connection.exec_query("EXPLAIN #{statement.sql}", "Explain #{statement.name}", statement.binds)
+        end
+      end
+
+      # ActiveRecord v7.2.0 introduced with_connection
+      def supports_with_connection?
+        return @supports_with_connection if defined?(@supports_with_connection)
+
+        @supports_with_connection = defined?(::ActiveRecord::VERSION::STRING) &&
+          Gem::Version.new(ActiveRecord::VERSION::STRING) >= Gem::Version.new('7.2.0')
+      end
+
       def close_connections
         ConnectionManager.instance.close_connections
       end

lib/new_relic/agent/instrumentation/active_record.rb

@@ -9,14 +9,7 @@ module Agent
     module Instrumentation
       module ActiveRecord
         EXPLAINER = lambda do |statement|
-          connection = NewRelic::Agent::Database.get_connection(statement.config) do
-            ::ActiveRecord::Base.send("#{statement.config[:adapter]}_connection",
-              statement.config)
-          end
-          # the following line needs else branch coverage
-          if connection && connection.respond_to?(:execute) # rubocop:disable Style/SafeNavigation
-            return connection.execute("EXPLAIN #{statement.sql}")
-          end
+          NewRelic::Agent::Database.explain_this(statement, true)
         end
 
         def self.insert_instrumentation

lib/new_relic/agent/instrumentation/active_record_subscriber.rb

@@ -70,18 +70,7 @@ def finish(name, id, payload) # THREAD_LOCAL_ACCESS
         end
 
         def get_explain_plan(statement)
-          connection = NewRelic::Agent::Database.get_connection(statement.config) do
-            ::ActiveRecord::Base.send("#{statement.config[:adapter]}_connection",
-              statement.config)
-          end
-          # the following line needs else branch coverage
-          if connection && connection.respond_to?(:exec_query) # rubocop:disable Style/SafeNavigation
-            return connection.exec_query("EXPLAIN #{statement.sql}",
-              "Explain #{statement.name}",
-              statement.binds)
-          end
-        rescue => e
-          NewRelic::Agent.logger.debug("Couldn't fetch the explain plan for #{statement} due to #{e}")
+          NewRelic::Agent::Database.explain_this(statement)
         end
 
         def active_record_config(payload)

test/multiverse/suites/active_record_pg/active_record_test.rb

@@ -375,14 +375,15 @@ def test_metrics_for_direct_sql_other
   end
 
   def test_metrics_for_direct_sql_show
-    if supports_show_tables?
-      in_web_transaction do
-        conn = Order.connection
+    skip "Adapter does not support 'show tables'" unless supports_show_tables?
+
+    in_web_transaction do
+      Order.with_connection do |conn|
         conn.execute('show tables')
       end
-
-      assert_generic_rollup_metrics('show')
     end
+
+    assert_generic_rollup_metrics('show')
   end
 
   def test_still_records_metrics_in_error_cases
@@ -454,11 +455,10 @@ def test_gathers_explain_plans
 
       sample.prepare_to_send!
       explanations = sql_node.params[:explain_plan]
-      if supports_explain_plans?
-        refute_nil explanations, "No explains in node: #{sql_node}"
-        assert_equal(2, explanations.size,
-          "No explains in node: #{sql_node}")
-      end
+
+      refute_nil explanations, "No explains in node: #{sql_node}"
+      assert_equal(2, explanations.size,
+        "No explains in node: #{sql_node}")
     end
   end
 
@@ -593,11 +593,7 @@ def adapter
   end
 
   def supports_show_tables?
-    [:mysql, :postgres].include?(adapter)
-  end
-
-  def supports_explain_plans?
-    [:mysql, :postgres].include?(adapter)
+    [:mysql, :mysql2, :trilogy].include?(adapter)
   end
 
   def current_product

test/new_relic/agent/sql_sampler_test.rb

@@ -167,15 +167,17 @@ def test_harvest_should_aggregate_similar_queries
   end
 
   def test_harvest_should_collect_explain_plans
-    @connection.expects(:execute).with('EXPLAIN select * from test') \
-      .returns(dummy_mysql_explain_result({'header0' => 'foo0', 'header1' => 'foo1', 'header2' => 'foo2'}))
-    @connection.expects(:execute).with('EXPLAIN select * from test2') \
-      .returns(dummy_mysql_explain_result({'header0' => 'bar0', 'header1' => 'bar1', 'header2' => 'bar2'}))
-
     data = NewRelic::Agent::TransactionSqlData.new
     data.set_transaction_info('/c/a', 'guid')
     data.set_transaction_name('WebTransaction/Controller/c/a')
-    explainer = NewRelic::Agent::Instrumentation::ActiveRecord::EXPLAINER
+    explainer = proc do |statement|
+      hash = if statement.sql.match?('test2')
+        {'header0' => 'bar0', 'header1' => 'bar1', 'header2' => 'bar2'}
+      else
+        {'header0' => 'foo0', 'header1' => 'foo1', 'header2' => 'foo2'}
+      end
+      dummy_mysql_explain_result(hash)
+    end
     config = {:adapter => 'mysql'}
     queries = [
       NewRelic::Agent::SlowSql.new(NewRelic::Agent::Database::Statement.new('select * from test', config, explainer),

test/new_relic/agent/transaction/trace_node_test.rb

@@ -290,12 +290,11 @@ def test_each_node_with_nest_tracking
 
   def test_explain_sql_raising_an_error
     s = NewRelic::Agent::Transaction::TraceNode.new('Custom/test/metric', Process.clock_gettime(Process::CLOCK_REALTIME))
-    config = {:adapter => 'mysql'}
     statement = NewRelic::Agent::Database::Statement.new('SELECT')
-    statement.config = config
+    statement.config = {:adapter => 'mysql'}
     statement.explainer = NewRelic::Agent::Instrumentation::ActiveRecord::EXPLAINER
     s.params = {:sql => statement}
-    NewRelic::Agent::Database.expects(:get_connection).with(config).raises(RuntimeError.new('whee'))
+    NewRelic::Agent::Database.expects(:explain_this).raises(RuntimeError.new('whee'))
     s.explain_sql
   end