Add row_count field to sql.active_record notification

marvin-bitterlich · marvin-bitterlich · commit e9a2288c1380 · 2024-01-27T00:12:19.000Z
This field returns the amount of rows returned by the query that emitted the notification.    
This metric is useful in cases where one wants to detect queries with big result sets.
diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
@@ -1,3 +1,11 @@
+*   Add row_count field to sql.active_record notification
+
+    This field returns the amount of rows returned by the query that emitted the notification.
+
+    This metric is useful in cases where one wants to detect queries with big result sets.
+
+    *Marvin Bitterlich*
+
 *   Consistently raise an `ArgumentError` when passing an invalid argument to a nested attributes association writer.
 
     Previously, this would only raise on collection associations and produce a generic error on singular associations.
diff --git a/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb b/activerecord/lib/active_record/connection_adapters/abstract_adapter.rb
@@ -1156,6 +1156,7 @@ def log(sql, name = "SQL", binds = [], type_casted_binds = [], statement_name =
             statement_name:    statement_name,
             async:             async,
             connection:        self,
+            row_count:         0,
             &block
           )
         rescue ActiveRecord::StatementInvalid => ex
diff --git a/activerecord/lib/active_record/connection_adapters/mysql2/database_statements.rb b/activerecord/lib/active_record/connection_adapters/mysql2/database_statements.rb
@@ -98,12 +98,13 @@ def with_multi_statements
           end
 
           def raw_execute(sql, name, async: false, allow_retry: false, materialize_transactions: true)
-            log(sql, name, async: async) do
+            log(sql, name, async: async) do |notification_payload|
               with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
                 sync_timezone_changes(conn)
                 result = conn.query(sql)
                 verified!
                 handle_warnings(sql)
+                notification_payload[:row_count] = result&.size || 0
                 result
               end
             end
@@ -117,7 +118,7 @@ def exec_stmt_and_free(sql, name, binds, cache_stmt: false, async: false)
 
             type_casted_binds = type_casted_binds(binds)
 
-            log(sql, name, binds, type_casted_binds, async: async) do
+            log(sql, name, binds, type_casted_binds, async: async) do |notification_payload|
               with_raw_connection do |conn|
                 sync_timezone_changes(conn)
 
@@ -143,6 +144,7 @@ def exec_stmt_and_free(sql, name, binds, cache_stmt: false, async: false)
                 end
 
                 ret = yield stmt, result
+                notification_payload[:row_count] = result&.size || 0
                 result.free if result
                 stmt.close unless cache_stmt
                 ret
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb b/activerecord/lib/active_record/connection_adapters/postgresql/database_statements.rb
@@ -14,10 +14,11 @@ def explain(arel, binds = [], options = [])
         def query(sql, name = nil) # :nodoc:
           mark_transaction_written_if_write(sql)
 
-          log(sql, name) do
+          log(sql, name) do |notification_payload|
             with_raw_connection do |conn|
               result = conn.async_exec(sql).map_types!(@type_map_for_results).values
               verified!
+              notification_payload[:row_count] = result.count
               result
             end
           end
@@ -50,11 +51,12 @@ def execute(...) # :nodoc:
         end
 
         def raw_execute(sql, name, async: false, allow_retry: false, materialize_transactions: true)
-          log(sql, name, async: async) do
+          log(sql, name, async: async) do |notification_payload|
             with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
               result = conn.async_exec(sql)
               verified!
               handle_warnings(result)
+              notification_payload[:row_count] = result.count
               result
             end
           end
diff --git a/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb b/activerecord/lib/active_record/connection_adapters/postgresql_adapter.rb
@@ -880,10 +880,11 @@ def exec_no_cache(sql, name, binds, async:, allow_retry:, materialize_transactio
           update_typemap_for_default_timezone
 
           type_casted_binds = type_casted_binds(binds)
-          log(sql, name, binds, type_casted_binds, async: async) do
+          log(sql, name, binds, type_casted_binds, async: async) do |notification_payload|
             with_raw_connection(allow_retry: false, materialize_transactions: materialize_transactions) do |conn|
               result = conn.exec_params(sql, type_casted_binds)
               verified!
+              notification_payload[:row_count] = result.count
               result
             end
           end
@@ -898,9 +899,10 @@ def exec_cache(sql, name, binds, async:, allow_retry:, materialize_transactions:
             stmt_key = prepare_statement(sql, binds, conn)
             type_casted_binds = type_casted_binds(binds)
 
-            log(sql, name, binds, type_casted_binds, stmt_key, async: async) do
+            log(sql, name, binds, type_casted_binds, stmt_key, async: async) do |notification_payload|
               result = conn.exec_prepared(stmt_key, type_casted_binds)
               verified!
+              notification_payload[:row_count] = result.count
               result
             end
           end
diff --git a/activerecord/lib/active_record/connection_adapters/sqlite3/database_statements.rb b/activerecord/lib/active_record/connection_adapters/sqlite3/database_statements.rb
@@ -29,7 +29,7 @@ def internal_exec_query(sql, name = nil, binds = [], prepare: false, async: fals
 
           type_casted_binds = type_casted_binds(binds)
 
-          log(sql, name, binds, type_casted_binds, async: async) do
+          log(sql, name, binds, type_casted_binds, async: async) do |notification_payload|
             with_raw_connection do |conn|
               # Don't cache statements if they are not prepared
               unless prepare
@@ -52,7 +52,9 @@ def internal_exec_query(sql, name = nil, binds = [], prepare: false, async: fals
               end
               verified!
 
-              build_result(columns: cols, rows: records)
+              result = build_result(columns: cols, rows: records)
+              notification_payload[:row_count] = result.length
+              result
             end
           end
         end
@@ -113,10 +115,11 @@ def high_precision_current_timestamp
 
         private
           def raw_execute(sql, name, async: false, allow_retry: false, materialize_transactions: false)
-            log(sql, name, async: async) do
+            log(sql, name, async: async) do |notification_payload|
               with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
                 result = conn.execute(sql)
                 verified!
+                notification_payload[:row_count] = result.length
                 result
               end
             end
@@ -136,10 +139,11 @@ def execute_batch(statements, name = nil)
             check_if_write_query(sql)
             mark_transaction_written_if_write(sql)
 
-            log(sql, name) do
+            log(sql, name) do |notification_payload|
               with_raw_connection do |conn|
                 result = conn.execute_batch2(sql)
                 verified!
+                notification_payload[:row_count] = result.length
                 result
               end
             end
diff --git a/activerecord/lib/active_record/connection_adapters/trilogy/database_statements.rb b/activerecord/lib/active_record/connection_adapters/trilogy/database_statements.rb
@@ -43,12 +43,13 @@ def exec_delete(sql, name = nil, binds = []) # :nodoc:
 
         private
           def raw_execute(sql, name, async: false, allow_retry: false, materialize_transactions: true)
-            log(sql, name, async: async) do
+            log(sql, name, async: async) do |notification_payload|
               with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
                 sync_timezone_changes(conn)
                 result = conn.query(sql)
                 verified!
                 handle_warnings(sql)
+                notification_payload[:row_count] = result.count
                 result
               end
             end
diff --git a/activerecord/test/cases/instrumentation_test.rb b/activerecord/test/cases/instrumentation_test.rb
@@ -111,6 +111,42 @@ def test_payload_name_on_grouped_count
       ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
     end
 
+    def test_payload_row_count_on_select_all
+      10.times { Book.create(name: "row count book 1") }
+      subscriber = ActiveSupport::Notifications.subscribe("sql.active_record") do |_, _, _, _, payload|
+        if payload[:sql].match?("SELECT")
+          assert_equal 10, payload[:row_count]
+        end
+      end
+      Book.where(name: "row count book 1").to_a
+    ensure
+      ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+    end
+
+    def test_payload_row_count_on_pluck
+      10.times { Book.create(name: "row count book 2") }
+      subscriber = ActiveSupport::Notifications.subscribe("sql.active_record") do |_, _, _, _, payload|
+        if payload[:sql].match?("SELECT")
+          assert_equal 10, payload[:row_count]
+        end
+      end
+      Book.where(name: "row count book 2").pluck(:name)
+    ensure
+      ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+    end
+
+    def test_payload_row_count_on_raw_sql
+      10.times { Book.create(name: "row count book 3") }
+      subscriber = ActiveSupport::Notifications.subscribe("sql.active_record") do |_, _, _, _, payload|
+        if payload[:sql].match?("SELECT")
+          assert_equal 10, payload[:row_count]
+        end
+      end
+      ActiveRecord::Base.connection.execute("SELECT * FROM books WHERE name='row count book 3';")
+    ensure
+      ActiveSupport::Notifications.unsubscribe(subscriber) if subscriber
+    end
+
     def test_payload_connection_with_query_cache_disabled
       connection = Book.connection
       subscriber = ActiveSupport::Notifications.subscribe("sql.active_record") do |_, _, _, _, payload|
diff --git a/guides/source/active_support_instrumentation.md b/guides/source/active_support_instrumentation.md
@@ -365,6 +365,7 @@ The `:cache_hits` key is only included if the collection is rendered with `cache
 | `:statement_name`    | SQL Statement name                       |
 | `:async`             | `true` if query is loaded asynchronously |
 | `:cached`            | `true` is added when cached queries used |
+| `:row_count`         | Number of rows returned by the query     |
 
 Adapters may add their own data as well.
 
@@ -375,7 +376,8 @@ Adapters may add their own data as well.
   connection: <ActiveRecord::ConnectionAdapters::SQLite3Adapter:0x00007f9f7a838850>,
   binds: [<ActiveModel::Attribute::WithCastValue:0x00007fe19d15dc00>],
   type_casted_binds: [11],
-  statement_name: nil
+  statement_name: nil,
+  row_count: 5
 }
 ```
 

Original file line number	Diff line number	Diff line change
`@@ -1156,6 +1156,7 @@ def log(sql, name = "SQL", binds = [], type_casted_binds = [], statement_name =`
`1156`	`1156`	`statement_name: statement_name,`
`1157`	`1157`	`async: async,`
`1158`	`1158`	`connection: self,`
	`1159`	`+ row_count: 0,`
`1159`	`1160`	`&block`
`1160`	`1161`	`)`
`1161`	`1162`	`rescue ActiveRecord::StatementInvalid => ex`