Fix composite primary key with different data type with triggers (#1164)

Author: vividmuimui

README.md

@@ -49,6 +49,10 @@ adapter.exclude_output_inserted_table_names['my_table_name'] = true
 
 # Explicitly set the data type for the temporary key table.
 adapter.exclude_output_inserted_table_names['my_uuid_table_name'] = 'uniqueidentifier'
+
+
+# Explicitly set data types when data type is different for composite primary keys.
+adapter.exclude_output_inserted_table_names['my_composite_pk_table_name'] = { pk_col_one: "uniqueidentifier", pk_col_two: "int" }
 ```
 
 

lib/active_record/connection_adapters/sqlserver/database_statements.rb

@@ -278,13 +278,17 @@ def sql_for_insert(sql, pk, binds, returning)
                   exclude_output_inserted = exclude_output_inserted_table_name?(table_name, sql)
 
                   if exclude_output_inserted
-                    quoted_pk = Array(pk).map { |subkey| SQLServer::Utils.extract_identifiers(subkey).quoted }
+                    pk_and_types = Array(pk).map do |subkey|
+                      {
+                        quoted: SQLServer::Utils.extract_identifiers(subkey).quoted,
+                        id_sql_type: exclude_output_inserted_id_sql_type(subkey, exclude_output_inserted)
+                      }
+                    end
 
-                    id_sql_type = exclude_output_inserted.is_a?(TrueClass) ? "bigint" : exclude_output_inserted
                     <<~SQL.squish
-                      DECLARE @ssaIdInsertTable table (#{quoted_pk.map { |subkey| "#{subkey} #{id_sql_type}"}.join(", ") });
-                      #{sql.dup.insert sql.index(/ (DEFAULT )?VALUES/i), " OUTPUT #{ quoted_pk.map { |subkey| "INSERTED.#{subkey}" }.join(", ") } INTO @ssaIdInsertTable"}
-                      SELECT #{quoted_pk.map {|subkey| "CAST(#{subkey} AS #{id_sql_type}) #{subkey}"}.join(", ")} FROM @ssaIdInsertTable
+                      DECLARE @ssaIdInsertTable table (#{pk_and_types.map { |pk_and_type| "#{pk_and_type[:quoted]} #{pk_and_type[:id_sql_type]}"}.join(", ") });
+                      #{sql.dup.insert sql.index(/ (DEFAULT )?VALUES/i), " OUTPUT #{ pk_and_types.map { |pk_and_type| "INSERTED.#{pk_and_type[:quoted]}" }.join(", ") } INTO @ssaIdInsertTable"}
+                      SELECT #{pk_and_types.map {|pk_and_type| "CAST(#{pk_and_type[:quoted]} AS #{pk_and_type[:id_sql_type]}) #{pk_and_type[:quoted]}"}.join(", ")} FROM @ssaIdInsertTable
                     SQL
                   else
                     returning_columns = returning || Array(pk)
@@ -382,6 +386,12 @@ def exclude_output_inserted_table_name?(table_name, sql)
           self.class.exclude_output_inserted_table_names[table_name]
         end
 
+        def exclude_output_inserted_id_sql_type(pk, exclude_output_inserted)
+          return "bigint" if exclude_output_inserted.is_a?(TrueClass)
+          return exclude_output_inserted[pk.to_sym] if exclude_output_inserted.is_a?(Hash)
+          exclude_output_inserted
+        end
+
         def query_requires_identity_insert?(sql)
           return false unless insert_sql?(sql)
 

test/cases/trigger_test_sqlserver.rb

@@ -38,4 +38,14 @@ class SQLServerTriggerTest < ActiveRecord::TestCase
     _(obj.pk_col_two).must_equal 42
     _(obj.pk_col_one.to_s).must_equal SSTestTriggerHistory.first.id_source
   end
+
+  it "can insert into a table with composite pk with different data type with output inserted - with a hash setting for table name" do
+    exclude_output_inserted_table_names["sst_table_with_composite_pk_trigger_with_different_data_type"] = { pk_col_one: "uniqueidentifier", pk_col_two: "int" }
+    assert SSTestTriggerHistory.all.empty?
+    obj = SSTestTriggerCompositePkWithDefferentDataType.create! pk_col_two: 123, event_name: "test trigger"
+    _(obj.event_name).must_equal "test trigger"
+    _(obj.pk_col_one).must_be :present?
+    _(obj.pk_col_two).must_equal 123
+    _(obj.pk_col_one.to_s).must_equal SSTestTriggerHistory.first.id_source
+  end
 end

test/models/sqlserver/trigger.rb

@@ -11,3 +11,7 @@ class SSTestTriggerUuid < ActiveRecord::Base
 class SSTestTriggerCompositePk < ActiveRecord::Base
   self.table_name = "sst_table_with_composite_pk_trigger"
 end
+
+class SSTestTriggerCompositePkWithDefferentDataType < ActiveRecord::Base
+  self.table_name = "sst_table_with_composite_pk_trigger_with_different_data_type"
+end

test/schema/sqlserver_specific_schema.rb

@@ -249,6 +249,23 @@
     SELECT pk_col_one AS id_source, event_name FROM INSERTED
   SQL
 
+  execute "IF EXISTS(SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = 'sst_table_with_composite_pk_trigger_with_different_data_type') DROP TABLE sst_table_with_composite_pk_trigger_with_different_data_type"
+  execute <<-SQL
+    CREATE TABLE sst_table_with_composite_pk_trigger_with_different_data_type(
+      pk_col_one uniqueidentifier DEFAULT NEWID(),
+      pk_col_two int NOT NULL,
+      event_name nvarchar(255),
+      CONSTRAINT PK_sst_table_with_composite_pk_trigger_with_different_data_type PRIMARY KEY (pk_col_one, pk_col_two)
+    )
+  SQL
+  execute <<-SQL
+    CREATE TRIGGER sst_table_with_composite_pk_trigger_with_different_data_type_t ON sst_table_with_composite_pk_trigger_with_different_data_type
+    FOR INSERT
+    AS
+    INSERT INTO sst_table_with_trigger_history (id_source, event_name)
+    SELECT pk_col_one AS id_source, event_name FROM INSERTED
+  SQL
+
   # Another schema.
 
   create_table :sst_schema_columns, force: true do |t|