Add affected_rows to sql.active_record

The [recently added][1] `row_count` value is very useful for identifying
cases where a query would return a large result set as large results can
end up using a lot of memory or even be blocked by databases like
Vitess.

However, somewhere that `row_count` falls short is for queries that do
not necessarily return their results back to the client. These queries
that affect too many rows can lead to their own set of problems, such as
overwhelming replication and causing replication lag.

This commit adds the number of `affected_rows` to the
`sql.active_record` Notification so that these kinds of queries can be
identified as well.

[1]: e9a2288

Author: skipkayhil

activerecord/CHANGELOG.md

@@ -1,3 +1,7 @@
+*   Add `affected_rows` to `sql.active_record` Notification.
+
+    *Hartley McGuire*
+
 *   Fix `sum` when performing a grouped calculation.
 
     `User.group(:friendly).sum` no longer worked. This is fixed.

activerecord/lib/active_record/connection_adapters/abstract_adapter.rb

@@ -1135,6 +1135,7 @@ def log(sql, name = "SQL", binds = [], type_casted_binds = [], async: false, &bl
             async:             async,
             connection:        self,
             transaction:       current_transaction.user_transaction.presence,
+            affected_rows:     0,
             row_count:         0,
             &block
           )

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

@@ -65,9 +65,11 @@ def perform_query(raw_connection, sql, binds, type_casted_binds, prepare:, notif
               raw_connection.query(sql)
             end
 
+            @affected_rows_before_warnings = raw_connection.affected_rows
+
+            notification_payload[:affected_rows] = @affected_rows_before_warnings
             notification_payload[:row_count] = result&.size || 0
 
-            @affected_rows_before_warnings = raw_connection.affected_rows
             raw_connection.abandon_results!
 
             verified!

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

@@ -164,6 +164,8 @@ def perform_query(raw_connection, sql, binds, type_casted_binds, prepare:, notif
 
             verified!
             handle_warnings(result)
+
+            notification_payload[:affected_rows] = result.cmd_tuples
             notification_payload[:row_count] = result.count
             result
           end

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

@@ -109,6 +109,7 @@ def perform_query(raw_connection, sql, binds, type_casted_binds, prepare:, notif
             @last_affected_rows = raw_connection.changes
             verified!
 
+            notification_payload[:affected_rows] = @last_affected_rows
             notification_payload[:row_count] = result&.length || 0
             result
           end

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

@@ -30,6 +30,8 @@ def perform_query(raw_connection, sql, binds, type_casted_binds, prepare:, notif
             end
             verified!
             handle_warnings(sql)
+
+            notification_payload[:affected_rows] = result.affected_rows
             notification_payload[:row_count] = result.count
             result
           ensure

activerecord/test/cases/instrumentation_test.rb

@@ -172,6 +172,27 @@ def test_payload_row_count_on_cache
       assert_equal 1, events[1][:row_count]
     end
 
+    def test_payload_affected_rows
+      affected_row_values = []
+
+      ActiveSupport::Notifications.subscribed(
+        -> (event) { affected_row_values << event.payload[:affected_rows] },
+        "sql.active_record",
+      ) do
+        # The combination of MariaDB + Trilogy returns 0 for affected_rows with
+        # INSERT ... RETURNING
+        Book.insert_all!([{ name: "One" }, { name: "Two" }, { name: "Three" }, { name: "Four" }], returning: false)
+
+        Book.where(name: ["One", "Two", "Three"]).update_all(status: :published)
+
+        Book.where(name: ["Three", "Four"]).delete_all
+
+        Book.where(name: ["Three", "Four"]).delete_all
+      end
+
+      assert_equal [4, 3, 2, 0], affected_row_values
+    end
+
     def test_payload_connection_with_query_cache_disabled
       connection = ClothingItem.lease_connection
       subscriber = ActiveSupport::Notifications.subscribe("sql.active_record") do |_, _, _, _, payload|

guides/source/active_support_instrumentation.md

@@ -367,6 +367,7 @@ The `:cache_hits` key is only included if the collection is rendered with `cache
 | `:async`             | `true` if query is loaded asynchronously |
 | `:cached`            | `true` is added when cached queries used |
 | `:row_count`         | Number of rows returned by the query     |
+| `:affected_rows`     | Number of rows affected by the query     |
 
 Adapters may add their own data as well.
 
@@ -379,7 +380,8 @@ Adapters may add their own data as well.
   binds: [<ActiveModel::Attribute::WithCastValue:0x00007fe19d15dc00>],
   type_casted_binds: [11],
   statement_name: nil,
-  row_count: 5
+  row_count: 5,
+  affected_rows: 0
 }
 ```