extend "query_requires_identity_insert?" for merge queries with inserts

Author: andyundso

lib/active_record/connection_adapters/sqlserver/database_statements.rb

@@ -161,15 +161,6 @@ def build_insert_sql(insert) # :nodoc:
             # if we do not have any columns that might have conflicting values, just execute a regular insert
             return build_sql_for_regular_insert(insert) if conflict_columns.flatten.empty?
 
-            primary_keys_for_insert = insert_all.primary_keys.to_set
-
-            # if we receive a composite primary key, MSSQL will not be happy when we want to call "IDENTITY_INSERT"
-            # as there is likely no IDENTITY column
-            # so we need to check if there is exacty one
-            # TODO: Refactor to use existing "SET IDENTITY_INSERT" settings
-            enable_identity_insert = primary_keys_for_insert.length == 1 &&
-              (insert_all.primary_keys.to_set & insert.keys).present?
-
             # why is the "PARTITION BY" clause needed?
             # in every DBMS system, insert_all / upsert_all is usually implemented with INSERT, that allows to define what happens
             # when duplicates are found (SKIP OR UPDATE)
@@ -179,7 +170,7 @@ def build_insert_sql(insert) # :nodoc:
             # this works easiest by using PARTITION and make sure that any record
             # we are trying to insert is "the first one seen across all the potential conflicted columns"
             sql = <<~SQL
-              #{"SET IDENTITY_INSERT #{insert.model.quoted_table_name} ON;" if enable_identity_insert}
+
               MERGE INTO #{insert.model.quoted_table_name} WITH (UPDLOCK, HOLDLOCK) AS target
               USING (
                 SELECT *
@@ -230,7 +221,6 @@ def build_insert_sql(insert) # :nodoc:
             sql << build_sql_for_returning(insert:, insert_all: insert.send(:insert_all))
 
             sql << ";"
-            sql << "SET IDENTITY_INSERT #{insert.model.quoted_table_name} OFF;" if enable_identity_insert
 
             return sql
           end
@@ -527,11 +517,18 @@ def query_requires_identity_insert?(sql)
           raw_table_name = get_raw_table_name(sql)
           id_column = identity_columns(raw_table_name).first
 
-          (id_column && sql =~ /^\s*(INSERT|EXEC sp_executesql N'INSERT)[^(]+\([^)]*\b(#{id_column.name})\b,?[^)]*\)/i) ? SQLServer::Utils.extract_identifiers(raw_table_name).quoted : false
+          if id_column && (
+            sql =~ /^\s*(INSERT|EXEC sp_executesql N'INSERT)[^(]+\([^)]*\b(#{id_column.name})\b,?[^)]*\)/i ||
+              sql =~ /^\s*MERGE INTO.+THEN INSERT \([^)]*\b(#{id_column.name})\b,?[^)]*\)/im
+          )
+            SQLServer::Utils.extract_identifiers(raw_table_name).quoted
+          else
+            false
+          end
         end
 
         def insert_sql?(sql)
-          !(sql =~ /\A\s*(INSERT|EXEC sp_executesql N'INSERT)/i).nil?
+          !(sql =~ /\A\s*(INSERT|EXEC sp_executesql N'INSERT|MERGE INTO.+THEN INSERT)/im).nil?
         end
 
         def identity_columns(table_name)

lib/active_record/connection_adapters/sqlserver/schema_statements.rb

@@ -759,6 +759,8 @@ def get_raw_table_name(sql)
               .match(/\s*([^(]*)/i)[0]
           elsif s.match?(/^\s*UPDATE\s+.*/i)
             s.match(/UPDATE\s+([^\(\s]+)\s*/i)[1]
+          elsif s.match?(/^\s*MERGE INTO.*/i)
+            s.match(/^\s*MERGE\s+INTO\s+(\[?[a-z_ -]+\]?\.?\[?[a-z_ -]+\]?)\s+(AS|WITH|USING)/i)[1]
           else
             s.match(/FROM[\s|\(]+((\[[^\(\]]+\])|[^\(\s]+)\s*/i)[1]
           end.strip

test/cases/adapter_test_sqlserver.rb

@@ -13,6 +13,7 @@ class AdapterTestSQLServer < ActiveRecord::TestCase
   fixtures :tasks
 
   let(:basic_insert_sql) { "INSERT INTO [funny_jokes] ([name]) VALUES('Knock knock')" }
+  let(:basic_merge_sql) { "MERGE INTO [ships] WITH (UPDLOCK, HOLDLOCK) AS target USING ( SELECT * FROM ( SELECT [id], [name], ROW_NUMBER() OVER ( PARTITION BY [id] ORDER BY [id] DESC ) AS rn_0 FROM ( VALUES (101, N'RSS Sir David Attenborough') ) AS t1 ([id], [name]) ) AS ranked_source WHERE rn_0 = 1 ) AS source ON (target.[id] = source.[id]) WHEN MATCHED THEN UPDATE SET target.[name] = source.[name]" }
   let(:basic_update_sql) { "UPDATE [customers] SET [address_street] = NULL WHERE [id] = 2" }
   let(:basic_select_sql) { "SELECT * FROM [customers] WHERE ([customers].[id] = 1)" }
 
@@ -91,6 +92,7 @@ class AdapterTestSQLServer < ActiveRecord::TestCase
 
   it "return unquoted table name object from basic INSERT UPDATE and SELECT statements" do
     assert_equal "funny_jokes", connection.send(:get_table_name, basic_insert_sql)
+    assert_equal "ships", connection.send(:get_table_name, basic_merge_sql)
     assert_equal "customers", connection.send(:get_table_name, basic_update_sql)
     assert_equal "customers", connection.send(:get_table_name, basic_select_sql)
   end
@@ -219,6 +221,10 @@ class AdapterTestSQLServer < ActiveRecord::TestCase
       @identity_insert_sql_unquoted_sp = "EXEC sp_executesql N'INSERT INTO funny_jokes (id, name) VALUES (@0, @1)', N'@0 int, @1 nvarchar(255)', @0 = 420, @1 = N'Knock knock'"
       @identity_insert_sql_unordered_sp = "EXEC sp_executesql N'INSERT INTO [funny_jokes] ([name],[id]) VALUES (@0, @1)', N'@0 nvarchar(255), @1  int', @0 = N'Knock knock', @1 = 420"
 
+      @identity_merge_sql = "MERGE INTO [ships] WITH (UPDLOCK, HOLDLOCK) AS target USING ( SELECT * FROM ( SELECT [id], [name], ROW_NUMBER() OVER ( PARTITION BY [id] ORDER BY [id] DESC ) AS rn_0 FROM ( VALUES (101, N'RSS Sir David Attenborough') ) AS t1 ([id], [name]) ) AS ranked_source WHERE rn_0 = 1 ) AS source ON (target.[id] = source.[id]) WHEN MATCHED THEN UPDATE SET target.[name] = source.[name] WHEN NOT MATCHED BY TARGET THEN INSERT ([id], [name]) VALUES (source.[id], source.[name]) OUTPUT INSERTED.[id]"
+      @identity_merge_sql_unquoted = "MERGE INTO ships WITH (UPDLOCK, HOLDLOCK) AS target USING ( SELECT * FROM ( SELECT id, name, ROW_NUMBER() OVER ( PARTITION BY id ORDER BY id DESC ) AS rn_0 FROM ( VALUES (101, N'RSS Sir David Attenborough') ) AS t1 (id, name) ) AS ranked_source WHERE rn_0 = 1 ) AS source ON (target.id = source.id) WHEN MATCHED THEN UPDATE SET target.name = source.name WHEN NOT MATCHED BY TARGET THEN INSERT (id, name) VALUES (source.id, source.name) OUTPUT INSERTED.id"
+      @identity_merge_sql_unordered = "MERGE INTO [ships] WITH (UPDLOCK, HOLDLOCK) AS target USING ( SELECT * FROM ( SELECT [name], [id], ROW_NUMBER() OVER ( PARTITION BY [id] ORDER BY [id] DESC ) AS rn_0 FROM ( VALUES (101, N'RSS Sir David Attenborough') ) AS t1 ([name], [id]) ) AS ranked_source WHERE rn_0 = 1 ) AS source ON (target.[id] = source.[id]) WHEN MATCHED THEN UPDATE SET target.[name] = source.[name] WHEN NOT MATCHED BY TARGET THEN INSERT ([name], [id]) VALUES (source.[name], source.[id]) OUTPUT INSERTED.[id]"
+
       @identity_insert_sql_non_dbo = "INSERT INTO [test].[aliens] ([id],[name]) VALUES(420,'Mork')"
       @identity_insert_sql_non_dbo_unquoted = "INSERT INTO test.aliens ([id],[name]) VALUES(420,'Mork')"
       @identity_insert_sql_non_dbo_unordered = "INSERT INTO [test].[aliens] ([name],[id]) VALUES('Mork',420)"
@@ -235,6 +241,10 @@ class AdapterTestSQLServer < ActiveRecord::TestCase
       assert_equal "[funny_jokes]", connection.send(:query_requires_identity_insert?, @identity_insert_sql_unquoted_sp)
       assert_equal "[funny_jokes]", connection.send(:query_requires_identity_insert?, @identity_insert_sql_unordered_sp)
 
+      assert_equal "[ships]", connection.send(:query_requires_identity_insert?, @identity_merge_sql)
+      assert_equal "[ships]", connection.send(:query_requires_identity_insert?, @identity_merge_sql_unquoted)
+      assert_equal "[ships]", connection.send(:query_requires_identity_insert?, @identity_merge_sql_unordered)
+
       assert_equal "[test].[aliens]", connection.send(:query_requires_identity_insert?, @identity_insert_sql_non_dbo)
       assert_equal "[test].[aliens]", connection.send(:query_requires_identity_insert?, @identity_insert_sql_non_dbo_unquoted)
       assert_equal "[test].[aliens]", connection.send(:query_requires_identity_insert?, @identity_insert_sql_non_dbo_unordered)
@@ -244,7 +254,7 @@ class AdapterTestSQLServer < ActiveRecord::TestCase
     end
 
     it "return false to #query_requires_identity_insert? for normal SQL" do
-      [basic_insert_sql, basic_update_sql, basic_select_sql].each do |sql|
+      [basic_insert_sql, basic_merge_sql, basic_update_sql, basic_select_sql].each do |sql|
         assert !connection.send(:query_requires_identity_insert?, sql), "SQL was #{sql}"
       end
     end

test/cases/schema_test_sqlserver.rb

@@ -102,6 +102,24 @@ class SchemaTestSQLServer < ActiveRecord::TestCase
       end
     end
 
+    describe "MERGE statements" do
+      it do
+        assert_equal "[dashboards]", connection.send(:get_raw_table_name, "MERGE INTO [dashboards] AS target")
+      end
+
+      it do
+        assert_equal "lock_without_defaults", connection.send(:get_raw_table_name, "MERGE INTO lock_without_defaults AS target")
+      end
+
+      it do
+        assert_equal "[WITH - SPACES]", connection.send(:get_raw_table_name, "MERGE INTO [WITH - SPACES] AS target")
+      end
+
+      it do
+        assert_equal "[with].[select notation]", connection.send(:get_raw_table_name, "MERGE INTO [with].[select notation] AS target")
+      end
+    end
+
     describe "CREATE VIEW statements" do
       it do
         assert_equal "test_table_as", connection.send(:get_raw_table_name, "CREATE VIEW test_views ( test_table_a_id, test_table_b_id ) AS SELECT test_table_as.id as test_table_a_id, test_table_bs.id as test_table_b_id FROM (test_table_as with(nolock) LEFT JOIN test_table_bs with(nolock) ON (test_table_as.id = test_table_bs.test_table_a_id))")