Add affected_rows to sql.active_record (#1268)

Author: aidanharan

lib/active_record/connection_adapters/sqlserver/database_statements.rb

@@ -14,18 +14,19 @@ def write_query?(sql) # :nodoc:
         end
 
         def perform_query(raw_connection, sql, binds, type_casted_binds, prepare:, notification_payload:, batch:)
-          result = if id_insert_table_name = query_requires_identity_insert?(sql)
-                     # If the table name is a view, we need to get the base table name for enabling identity insert.
-                     id_insert_table_name = view_table_name(id_insert_table_name) if view_exists?(id_insert_table_name)
+          result, affected_rows = if id_insert_table_name = query_requires_identity_insert?(sql)
+                                    # If the table name is a view, we need to get the base table name for enabling identity insert.
+                                    id_insert_table_name = view_table_name(id_insert_table_name) if view_exists?(id_insert_table_name)
 
-                     with_identity_insert_enabled(id_insert_table_name, raw_connection) do
-                       internal_exec_sql_query(sql, raw_connection)
-                     end
-                   else
-                     internal_exec_sql_query(sql, raw_connection)
-                   end
+                                    with_identity_insert_enabled(id_insert_table_name, raw_connection) do
+                                      internal_exec_sql_query(sql, raw_connection)
+                                    end
+                                  else
+                                    internal_exec_sql_query(sql, raw_connection)
+                                  end
 
           verified!
+          notification_payload[:affected_rows] = affected_rows
           notification_payload[:row_count] = result.count
           result
         end
@@ -38,8 +39,18 @@ def cast_result(raw_result)
           end
         end
 
+        # Returns the affected rows from results.
         def affected_rows(raw_result)
-          raw_result.first['AffectedRows']
+          raw_result&.first&.fetch('AffectedRows', nil)
+        end
+
+        # Returns the affected rows from results or handle.
+        def affected_rows_from_results_or_handle(raw_result, handle)
+          if affected_rows_from_result = affected_rows(raw_result)
+            affected_rows_from_result
+          else
+            handle.affected_rows
+          end
         end
 
         def raw_execute(sql, name = nil, binds = [], prepare: false, async: false, allow_retry: false, materialize_transactions: true, batch: false)
@@ -53,7 +64,9 @@ def raw_execute(sql, name = nil, binds = [], prepare: false, async: false, allow
 
         def internal_exec_sql_query(sql, conn)
           handle = internal_raw_execute(sql, conn)
-          handle_to_names_and_values(handle, ar_result: true)
+          results = handle_to_names_and_values(handle, ar_result: true)
+
+          return results, affected_rows_from_results_or_handle(results, handle)
         ensure
           finish_statement_handle(handle)
         end
@@ -432,12 +445,15 @@ def handle_to_names_and_values(handle, options = {})
           end
           results = handle.each(query_options)
 
-          columns = handle.fields
-          # If query returns multiple result sets, only return the columns of the last one.
-          columns = columns.last if columns.any? && columns.all? { |e| e.is_a?(Array) }
-          columns = columns.map(&:downcase) if lowercase_schema_reflection
+          if options[:ar_result]
+            columns = handle.fields
+            columns = columns.last if columns.any? && columns.all? { |e| e.is_a?(Array) } # If query returns multiple result sets, only return the columns of the last one.
+            columns = columns.map(&:downcase) if lowercase_schema_reflection
 
-          options[:ar_result] ? ActiveRecord::Result.new(columns, results) : results
+            ActiveRecord::Result.new(columns, results)
+          else
+            results
+          end
         end
 
         def finish_statement_handle(handle)

test/cases/coerced_tests.rb

@@ -372,6 +372,13 @@ def test_payload_row_count_on_raw_sql_coerced
       Book.where(author_id: nil, name: 'row count book 3').delete_all
       Book.lease_connection.add_index(:books, [:author_id, :name], unique: true)
     end
+
+    # Fix randomly failing test. The loading of the model's schema was affecting the test.
+    coerce_tests! :test_payload_affected_rows
+    def test_payload_affected_rows_coerced
+      Book.send(:load_schema!)
+      original_test_payload_affected_rows
+    end
   end
 end