Merge pull request rails#54454 from skipkayhil/hm-instrument-allow-retry

Add `allow_retry` to `sql.active_record`

Author: byroot

activerecord/CHANGELOG.md

@@ -1,3 +1,10 @@
+*   Add `allow_retry` to `sql.active_record` instrumentation.
+
+    This enables identifying queries which are and are not automatically
+    retryable on connection errors.
+
+    *Hartley McGuire*
+
 *   Better support UPDATE with JOIN for Postgresql and SQLite3
 
     Previously when generating update queries with one or more JOIN clauses,

activerecord/lib/active_record/connection_adapters/abstract/database_statements.rb

@@ -551,7 +551,7 @@ def internal_exec_query(...) # :nodoc:
         # Lowest level way to execute a query. Doesn't check for illegal writes, doesn't annotate queries, yields a native result object.
         def raw_execute(sql, name = nil, binds = [], prepare: false, async: false, allow_retry: false, materialize_transactions: true, batch: false)
           type_casted_binds = type_casted_binds(binds)
-          log(sql, name, binds, type_casted_binds, async: async) do |notification_payload|
+          log(sql, name, binds, type_casted_binds, async: async, allow_retry: allow_retry) do |notification_payload|
             with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
               result = ActiveSupport::Dependencies.interlock.permit_concurrent_loads do
                 perform_query(conn, sql, binds, type_casted_binds, prepare: prepare, notification_payload: notification_payload, batch: batch)

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -1131,14 +1131,15 @@ def translate_exception_class(native_error, sql, binds)
           active_record_error
         end
 
-        def log(sql, name = "SQL", binds = [], type_casted_binds = [], async: false, &block) # :doc:
+        def log(sql, name = "SQL", binds = [], type_casted_binds = [], async: false, allow_retry: false, &block) # :doc:
           instrumenter.instrument(
             "sql.active_record",
             sql:               sql,
             name:              name,
             binds:             binds,
             type_casted_binds: type_casted_binds,
             async:             async,
+            allow_retry:       allow_retry,
             connection:        self,
             transaction:       current_transaction.user_transaction.presence,
             affected_rows:     0,

activerecord/test/cases/adapter_test.rb

@@ -700,74 +700,49 @@ def teardown
         assert_predicate @connection, :active?
       end
 
-      test "idempotent SELECT queries are retried and result in a reconnect" do
-        Post.first
-
-        remote_disconnect @connection
-
-        assert Post.first
-        assert_predicate @connection, :active?
-
-        remote_disconnect @connection
-
-        assert Post.where(id: [1, 2]).first
-        assert_predicate @connection, :active?
-      end
-
-      test "#find and #find_by queries with known attributes are retried and result in a reconnect" do
-        Post.first
-
-        remote_disconnect @connection
-
-        assert Post.find(1)
-        assert_predicate @connection, :active?
-
-        remote_disconnect @connection
-
-        assert Post.find_by(title: "Welcome to the weblog")
-        assert_predicate @connection, :active?
-      end
-
-      test "#exists? queries are retried and result in a reconnect" do
-        Post.first
-
-        remote_disconnect @connection
-
-        assert_predicate Post, :exists?
-        assert_predicate @connection, :active?
+      test "idempotent SELECT queries allow retries" do
+        notifications = capture_notifications("sql.active_record") do
+          assert Post.first
+          assert Post.where(id: [1, 2]).first
+          assert Post.find(1)
+          assert Post.find_by(title: "Welcome to the weblog")
+          assert_predicate Post, :exists?
+        end.select { |n| n.payload[:name] != "SCHEMA" }
+
+        assert_equal 5, notifications.length
+
+        notifications.each do |n|
+          assert n.payload[:allow_retry]
+        end
       end
 
-      test "queries containing SQL fragments are not retried" do
-        Post.first
+      test "queries containing SQL fragments do not allow retries" do
+        notifications = capture_notifications("sql.active_record") do
+          Post.where("1 = 1").to_a
+          Post.select("title AS custom_title").first
+          Book.find_by("updated_at < ?", 2.weeks.ago)
+        end.select { |n| n.payload[:name] != "SCHEMA" }
 
-        remote_disconnect @connection
-
-        assert_raises(ActiveRecord::ConnectionFailed) { Post.where("1 = 1").to_a }
-        assert_not_predicate @connection, :active?
+        assert_equal 3, notifications.length
 
-        remote_disconnect @connection
-
-        assert_raises(ActiveRecord::ConnectionFailed) { Post.select("title AS custom_title").first }
-        assert_not_predicate @connection, :active?
-
-        remote_disconnect @connection
-
-        assert_raises(ActiveRecord::ConnectionFailed) { Post.find_by("updated_at < ?", 2.weeks.ago) }
-        assert_not_predicate @connection, :active?
+        notifications.each do |n|
+          assert_not n.payload[:allow_retry]
+        end
       end
 
-      test "queries containing SQL functions are not retried" do
-        Post.first
-
-        remote_disconnect @connection
-
+      test "queries containing SQL functions do not allow retries" do
         tags_count_attr = Post.arel_table[:tags_count]
         abs_tags_count = Arel::Nodes::NamedFunction.new("ABS", [tags_count_attr])
 
-        assert_raises(ActiveRecord::ConnectionFailed) do
+        notifications = capture_notifications("sql.active_record") do
           Post.where(abs_tags_count.eq(2)).first
+        end.select { |n| n.payload[:name] != "SCHEMA" }
+
+        assert_equal 1, notifications.length
+
+        notifications.each do |n|
+          assert_not n.payload[:allow_retry]
         end
-        assert_not_predicate @connection, :active?
       end
 
       test "transaction restores after remote disconnection" do

guides/source/active_support_instrumentation.md

@@ -362,6 +362,7 @@ The `:cache_hits` key is only included if the collection is rendered with `cache
 | `:binds`             | Bind parameters                                        |
 | `:type_casted_binds` | Typecasted bind parameters                             |
 | `:async`             | `true` if query is loaded asynchronously               |
+| `:allow_retry`       | `true` if the query can be automatically retried       |
 | `:connection`        | Connection object                                      |
 | `:transaction`       | Current transaction, if any                            |
 | `:affected_rows`     | Number of rows affected by the query                   |
@@ -378,6 +379,7 @@ Adapters may add their own data as well.
   binds: [<ActiveModel::Attribute::WithCastValue:0x00007fe19d15dc00>],
   type_casted_binds: [11],
   async: false,
+  allow_retry: true,
   connection: <ActiveRecord::ConnectionAdapters::SQLite3Adapter:0x00007f9f7a838850>,
   transaction: <ActiveRecord::ConnectionAdapters::RealTransaction:0x0000000121b5d3e0>
   affected_rows: 0