Merge pull request rails#53277 from skipkayhil/hm-instrument-affected-rows

Add affected_rows to sql.active_record

Author: byroot

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
 }
 ```