From 5cb2520a44ebefa0e3e6b7b84b2846900ccbcc64 Mon Sep 17 00:00:00 2001
From: Aidan Haran <aidanharan@yahoo.com>
Date: Fri, 10 May 2024 16:07:10 +0100
Subject: [PATCH 1/2] Add row_count field to sql.active_record notification

---
 .../sqlserver/database_statements.rb          | 34 +++++++++----------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/lib/active_record/connection_adapters/sqlserver/database_statements.rb b/lib/active_record/connection_adapters/sqlserver/database_statements.rb
index f219ce20f..c2463eb92 100644
--- a/lib/active_record/connection_adapters/sqlserver/database_statements.rb
+++ b/lib/active_record/connection_adapters/sqlserver/database_statements.rb
@@ -14,8 +14,6 @@ def write_query?(sql) # :nodoc:
         end
 
         def raw_execute(sql, name, async: false, allow_retry: false, materialize_transactions: true)
-          result = nil
-
           log(sql, name, async: async) do
             with_raw_connection(allow_retry: allow_retry, materialize_transactions: materialize_transactions) do |conn|
               result = if id_insert_table_name = query_requires_identity_insert?(sql)
@@ -24,14 +22,12 @@ def raw_execute(sql, name, async: false, allow_retry: false, materialize_transac
                          internal_raw_execute(sql, conn, perform_do: true)
                        end
               verified!
+              result
             end
           end
-
-          result
         end
 
         def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: false, allow_retry: false)
-          result = nil
           sql = transform_query(sql)
 
           check_if_write_query(sql)
@@ -42,20 +38,21 @@ def internal_exec_query(sql, name = "SQL", binds = [], prepare: false, async: fa
             sql = sp_executesql_sql(sql, types, params, name)
           end
 
-          log(sql, name, binds, async: async) do
+          log(sql, name, binds, async: async) do |notification_payload|
             with_raw_connection do |conn|
-              if id_insert_table_name = query_requires_identity_insert?(sql)
-                with_identity_insert_enabled(id_insert_table_name, conn) do
-                  result = internal_exec_sql_query(sql, conn)
-                end
-              else
-                result = internal_exec_sql_query(sql, conn)
-              end
+              result = if id_insert_table_name = query_requires_identity_insert?(sql)
+                         with_identity_insert_enabled(id_insert_table_name, conn) do
+                           internal_exec_sql_query(sql, conn)
+                         end
+                       else
+                         internal_exec_sql_query(sql, conn)
+                       end
+
               verified!
+              notification_payload[:row_count] = result.count
+              result
             end
           end
-
-          result
         end
 
         def internal_exec_sql_query(sql, conn)
@@ -174,7 +171,7 @@ def execute_procedure(proc_name, *variables)
                  end.join(", ")
           sql = "EXEC #{proc_name} #{vars}".strip
 
-          log(sql, "Execute Procedure") do
+          log(sql, "Execute Procedure") do |notification_payload|
             with_raw_connection do |conn|
               result = internal_raw_execute(sql, conn)
               verified!
@@ -185,10 +182,11 @@ def execute_procedure(proc_name, *variables)
                 yield(r) if block_given?
               end
 
-              result.each.map { |row| row.is_a?(Hash) ? row.with_indifferent_access : row }
+              result = result.each.map { |row| row.is_a?(Hash) ? row.with_indifferent_access : row }
+              notification_payload[:row_count] = result.count
+              result
             end
           end
-
         end
 
         def with_identity_insert_enabled(table_name, conn)

From 92e5d603b3e071e9a33e3bce7c46e431839547de Mon Sep 17 00:00:00 2001
From: Aidan Haran <aidanharan@yahoo.com>
Date: Fri, 10 May 2024 16:29:40 +0100
Subject: [PATCH 2/2] Coerced tests

---
 test/cases/coerced_tests.rb | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/test/cases/coerced_tests.rb b/test/cases/coerced_tests.rb
index aabf4ffbf..04c258cd1 100644
--- a/test/cases/coerced_tests.rb
+++ b/test/cases/coerced_tests.rb
@@ -351,6 +351,28 @@ def test_payload_name_on_load_coerced
       Book.send(:load_schema!)
       original_test_payload_name_on_load
     end
+
+    # Need to remove index as SQL Server considers NULLs on a unique-index to be equal unlike PostgreSQL/MySQL/SQLite.
+    coerce_tests! :test_payload_row_count_on_select_all
+    def test_payload_row_count_on_select_all_coerced
+      connection.remove_index(:books, column: [:author_id, :name])
+
+      original_test_payload_row_count_on_select_all
+    ensure
+      Book.where(author_id: nil, name: 'row count book 1').delete_all
+      Book.lease_connection.add_index(:books, [:author_id, :name], unique: true)
+    end
+
+    # Need to remove index as SQL Server considers NULLs on a unique-index to be equal unlike PostgreSQL/MySQL/SQLite.
+    coerce_tests! :test_payload_row_count_on_raw_sql
+    def test_payload_row_count_on_raw_sql_coerced
+      connection.remove_index(:books, column: [:author_id, :name])
+
+      original_test_payload_row_count_on_raw_sql
+    ensure
+      Book.where(author_id: nil, name: 'row count book 3').delete_all
+      Book.lease_connection.add_index(:books, [:author_id, :name], unique: true)
+    end
   end
 end