feat: enable explicitly inserting IDENTITY values into mssql (#6199)

Allow explicit insertion of preset values for mssql IDENTITY columns. These are the equivalent to
auto-incrementing primary keys in MySQL, but require special conditional treatment.

Closes: #2199 for mssql

Co-authored-by: Daniel Klischies <Daniel.Klischies@umlaut.com>

src/query-builder/InsertQueryBuilder.ts

@@ -352,6 +352,18 @@ export class InsertQueryBuilder<Entity> extends QueryBuilder<Entity> {
             query += ` RETURNING ${returningExpression}`;
         }
 
+
+        // Inserting a specific value for an auto-increment primary key in mssql requires enabling IDENTITY_INSERT
+        // IDENTITY_INSERT can only be enabled for tables where there is an IDENTITY column and only if there is a value to be inserted (i.e. supplying DEFAULT is prohibited if IDENTITY_INSERT is enabled)
+        if (this.connection.driver instanceof SqlServerDriver
+            && this.expressionMap.mainAlias!.hasMetadata
+            && this.expressionMap.mainAlias!.metadata.columns
+                .filter((column) => this.expressionMap.insertColumns.length > 0 ? this.expressionMap.insertColumns.indexOf(column.propertyPath) !== -1 : column.isInsert)
+                .some((column) => this.isOverridingAutoIncrementBehavior(column))
+        ) {
+            query = `SET IDENTITY_INSERT ${tableName} ON; ${query}; SET IDENTITY_INSERT ${tableName} OFF`;
+        }
+
         return query;
     }
 
@@ -377,7 +389,8 @@ export class InsertQueryBuilder<Entity> extends QueryBuilder<Entity> {
                 && !(this.connection.driver instanceof OracleDriver)
                 && !(this.connection.driver instanceof AbstractSqliteDriver)
                 && !(this.connection.driver instanceof MysqlDriver)
-                && !(this.connection.driver instanceof AuroraDataApiDriver))
+                && !(this.connection.driver instanceof AuroraDataApiDriver)
+                && !(this.connection.driver instanceof SqlServerDriver && this.isOverridingAutoIncrementBehavior(column)))
                 return false;
 
             return true;
@@ -617,4 +630,19 @@ export class InsertQueryBuilder<Entity> extends QueryBuilder<Entity> {
         throw new InsertValuesMissingError();
     }
 
+    /**
+     * Checks if column is an auto-generated primary key, but the current insertion specifies a value for it.
+     * 
+     * @param column
+     */
+    protected isOverridingAutoIncrementBehavior(column: ColumnMetadata): boolean {
+        return column.isPrimary 
+                && column.isGenerated 
+                && column.generationStrategy === "increment"
+                && this.getValueSets().some((valueSet) => 
+                    column.getEntityValue(valueSet) !== undefined 
+                    && column.getEntityValue(valueSet) !== null
+                );
+    }
+
 }

test/github-issues/2199/issue-2199.ts

@@ -1,15 +1,16 @@
 import { expect } from "chai";
 import "reflect-metadata";
 import { Connection } from "../../../src/connection/Connection";
+import { SqlServerDriver } from '../../../src/driver/sqlserver/SqlServerDriver';
 import { closeTestingConnections, createTestingConnections, reloadTestingDatabases } from "../../utils/test-utils";
 import { Bar } from "./entity/Bar";
 
-describe("github issues > #2199 - Inserting value for @PrimaryGeneratedColumn() for mysql and sqlite", () => {
+describe("github issues > #2199 - Inserting value for @PrimaryGeneratedColumn() for mysql, sqlite and mssql", () => {
 
     let connections: Connection[];
     before(async () => connections = await createTestingConnections({
         entities: [__dirname + "/entity/*{.js,.ts}"],
-        enabledDrivers: ["mysql", "mariadb", "sqlite", "better-sqlite3"],
+        enabledDrivers: ["mysql", "mariadb", "sqlite", "better-sqlite3", "mssql"],
         schemaCreate: true,
         dropSchema: true
     }));
@@ -44,4 +45,25 @@ describe("github issues > #2199 - Inserting value for @PrimaryGeneratedColumn()
         const thirdBar = await connection.manager.save(thirdBarQuery);
         expect(thirdBar.id).to.eql(5);
     })));
+
+    it("should reset mssql's INSERT_IDENTITY flag correctly after failed queries", () => Promise.all(connections
+        .filter(connection => connection.driver instanceof SqlServerDriver)
+        .map(async connection => {
+            // Run a query that failes at the database level
+            await expect(connection.createQueryBuilder()
+                .insert()
+                .into(Bar)
+                .values({
+                    id: 20,
+                    description: () => "NONEXISTINGFUNCTION()",
+                })
+                .execute()
+            ).to.be.rejectedWith("Error: 'NONEXISTINGFUNCTION' is not a recognized built-in function name.");
+            // And now check that IDENTITY_INSERT is disabled by inserting something without an ID value and see if that works
+            const successfulBarQuery = connection.manager.create(Bar, {
+                description: "default id value"
+            });
+            const bar = await connection.manager.save(successfulBarQuery);
+            expect(bar.id).to.be.a('number');
+    })));
 });