Merge pull request rails#50887 from marvin-bitterlich/marvin.bitterlich/result-set

rafaelfranca · web-flow · commit e7073e4c95ce · 2024-01-29T13:02:16.000-05:00
Add row_count field to sql.active_record notification
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`