SQL Server Migrations: SET IDENTITY_INSERT ON during INSERT

dotnet · May 23, 2017 · 2601125 · 2601125
1 parent b66de51
commit 2601125
Show file tree

Hide file tree

Showing 3 changed files with 53 additions and 2 deletions.
diff --git a/src/EFCore.Relational/Migrations/MigrationsSqlGenerator.cs b/src/EFCore.Relational/Migrations/MigrationsSqlGenerator.cs
@@ -630,6 +630,13 @@ protected virtual void Generate(
             [NotNull] InsertOperation operation,
             [CanBeNull] IModel model,
             [NotNull] MigrationCommandListBuilder builder)
+            => Generate(operation, model, builder, terminate: true);
+
+        protected virtual void Generate(
+            [NotNull] InsertOperation operation,
+            [CanBeNull] IModel model,
+            [NotNull] MigrationCommandListBuilder builder,
+            bool terminate)
         {
             Check.NotNull(operation, nameof(operation));
             Check.NotNull(builder, nameof(builder));
@@ -672,8 +679,11 @@ protected virtual void Generate(
                 builder.Append(")");
             }
 
-            builder.AppendLine(Dependencies.SqlGenerationHelper.StatementTerminator);
-            EndStatement(builder);
+            if (terminate)
+            {
+                builder.AppendLine(Dependencies.SqlGenerationHelper.StatementTerminator);
+                EndStatement(builder);
+            }
         }
 
         protected virtual void Generate(

diff --git a/src/EFCore.SqlServer/Migrations/Internal/SqlServerMigrationsSqlGenerator.cs b/src/EFCore.SqlServer/Migrations/Internal/SqlServerMigrationsSqlGenerator.cs
@@ -823,6 +823,31 @@ protected override void Generate([NotNull] SqlOperation operation, [CanBeNull] I
             }
         }
 
+        protected override void Generate(
+            InsertOperation operation,
+            IModel model,
+            MigrationCommandListBuilder builder)
+        {
+            Check.NotNull(operation, nameof(operation));
+            Check.NotNull(builder, nameof(builder));
+
+            builder
+                .Append("SET IDENTITY_INSERT ")
+                .Append(Dependencies.SqlGenerationHelper.DelimitIdentifier(operation.Table, operation.Schema))
+                .Append(" ON")
+                .AppendLine(Dependencies.SqlGenerationHelper.StatementTerminator);
+
+            base.Generate(operation, model, builder, terminate: false);
+
+            builder
+                .AppendLine(Dependencies.SqlGenerationHelper.StatementTerminator)
+                .Append("SET IDENTITY_INSERT ")
+                .Append(Dependencies.SqlGenerationHelper.DelimitIdentifier(operation.Table, operation.Schema))
+                .Append(" OFF")
+                .AppendLine(Dependencies.SqlGenerationHelper.StatementTerminator)
+                .EndCommand();
+        }
+
         protected override void ColumnDefinition(
             string schema,
             string table,

diff --git a/test/EFCore.SqlServer.FunctionalTests/SqlServerMigrationSqlGeneratorTest.cs b/test/EFCore.SqlServer.FunctionalTests/SqlServerMigrationSqlGeneratorTest.cs
@@ -1113,6 +1113,22 @@ public virtual void SqlOperation_ignores_non_go()
                 Sql);
         }
 
+        public override void InsertRowsOperation()
+        {
+            base.InsertRowsOperation();
+
+            Assert.Equal(
+                "SET IDENTITY_INSERT [People] ON;" + EOL +
+                "INSERT INTO [People] ([Id], [Full Name])" + EOL +
+                "VALUES (0, NULL)," + EOL +
+                "       (1, N'Daenerys Targaryen')," + EOL +
+                "       (2, N'John Snow')," + EOL +
+                "       (3, N'Arya Stark')," + EOL +
+                "       (4, N'Harry Strickland');" + EOL +
+                "SET IDENTITY_INSERT [People] OFF;" + EOL,
+                Sql);
+        }
+
         public SqlServerMigrationSqlGeneratorTest()
             : base(SqlServerTestHelpers.Instance)
         {