fix: use OUTPUT INTO on SqlServer for returning columns (#5361)

* test: Add test for GitHub issue #5160

* fix: use OUTPUT INTO on SqlServer for returning columns

Closes #5160

* chore: fix typo in comment

* fix: use full types with length e.g. nvarchar(255)

* refactor: move output table declaration to driver

* fix: don't try to drop temporary tables

src/driver/sqlserver/SqlServerDriver.ts

@@ -700,6 +700,20 @@ export class SqlServerDriver implements Driver {
         }, {} as ObjectLiteral);
     }
 
+    buildTableVariableDeclaration(identifier: string, columns: ColumnMetadata[]): string {
+        const outputColumns = columns.map(column => {
+            return `${this.escape(column.databaseName)} ${this.createFullType(new TableColumn({
+                name: column.databaseName,
+                type: this.normalizeType(column),
+                length: column.length,
+                isNullable: column.isNullable,
+                isArray: column.isArray,
+            }))}`;
+        });
+
+        return `DECLARE ${identifier} TABLE (${outputColumns.join(", ")})`;
+    }
+
     // -------------------------------------------------------------------------
     // Protected Methods
     // -------------------------------------------------------------------------

src/driver/sqlserver/SqlServerQueryRunner.ts

@@ -1370,6 +1370,11 @@ export class SqlServerQueryRunner extends BaseQueryRunner implements QueryRunner
                 return Promise.all(dropFkQueries.map(result => result["query"]).map(dropQuery => this.query(dropQuery)));
             }));
             await Promise.all(allTablesResults.map(tablesResult => {
+                if (tablesResult["TABLE_NAME"].startsWith("#")) {
+                    // don't try to drop temporary tables
+                    return;
+                }
+
                 const dropTableSql = `DROP TABLE "${tablesResult["TABLE_CATALOG"]}"."${tablesResult["TABLE_SCHEMA"]}"."${tablesResult["TABLE_NAME"]}"`;
                 return this.query(dropTableSql);
             }));

src/query-builder/InsertQueryBuilder.ts

@@ -16,7 +16,7 @@ import {ReturningResultsEntityUpdator} from "./ReturningResultsEntityUpdator";
 import {AbstractSqliteDriver} from "../driver/sqlite-abstract/AbstractSqliteDriver";
 import {SqljsDriver} from "../driver/sqljs/SqljsDriver";
 import {BroadcasterResult} from "../subscriber/BroadcasterResult";
-import {EntitySchema} from "../";
+import {EntitySchema} from "../entity-schema/EntitySchema";
 import {OracleDriver} from "../driver/oracle/OracleDriver";
 import {AuroraDataApiDriver} from "../driver/aurora-data-api/AuroraDataApiDriver";
 
@@ -68,21 +68,33 @@ export class InsertQueryBuilder<Entity> extends QueryBuilder<Entity> {
                 if (broadcastResult.promises.length > 0) await Promise.all(broadcastResult.promises);
             }
 
+            let declareSql: string | null = null;
+            let selectOutputSql: string | null = null;
+
             // if update entity mode is enabled we may need extra columns for the returning statement
             // console.time(".prepare returning statement");
             const returningResultsEntityUpdator = new ReturningResultsEntityUpdator(queryRunner, this.expressionMap);
             if (this.expressionMap.updateEntity === true && this.expressionMap.mainAlias!.hasMetadata) {
                 this.expressionMap.extraReturningColumns = returningResultsEntityUpdator.getInsertionReturningColumns();
+
+                if (this.expressionMap.extraReturningColumns.length > 0 && this.connection.driver instanceof SqlServerDriver) {
+                    declareSql = this.connection.driver.buildTableVariableDeclaration("@OutputTable", this.expressionMap.extraReturningColumns);
+                    selectOutputSql = `SELECT * FROM @OutputTable`;
+                }
             }
             // console.timeEnd(".prepare returning statement");
 
             // execute query
             // console.time(".getting query and parameters");
-            const [sql, parameters] = this.getQueryAndParameters();
+            const [insertSql, parameters] = this.getQueryAndParameters();
             // console.timeEnd(".getting query and parameters");
             const insertResult = new InsertResult();
             // console.time(".query execution by database");
-            insertResult.raw = await queryRunner.query(sql, parameters);
+            const statements = [declareSql, insertSql, selectOutputSql];
+            insertResult.raw = await queryRunner.query(
+                statements.filter(sql => sql != null).join(";\n\n"),
+                parameters,
+            );
             // console.timeEnd(".query execution by database");
 
             // load returning results and set them to the entity if entity updation is enabled

src/query-builder/QueryBuilder.ts

@@ -635,6 +635,13 @@ export abstract class QueryBuilder<Entity> {
                     return this.connection.driver.createParameter(parameterName, Object.keys(this.expressionMap.nativeParameters).length);
                 }).join(", ");
             }
+
+            if (driver instanceof SqlServerDriver) {
+                if (this.expressionMap.queryType === "insert" || this.expressionMap.queryType === "update") {
+                    columnsExpression += " INTO @OutputTable";
+                }
+            }
+
             return columnsExpression;
 
         } else if (typeof this.expressionMap.returning === "string") {

src/query-builder/UpdateQueryBuilder.ts

@@ -74,21 +74,32 @@ export class UpdateQueryBuilder<Entity> extends QueryBuilder<Entity> implements
                 if (broadcastResult.promises.length > 0) await Promise.all(broadcastResult.promises);
             }
 
+            let declareSql: string | null = null;
+            let selectOutputSql: string | null = null;
+
             // if update entity mode is enabled we may need extra columns for the returning statement
             const returningResultsEntityUpdator = new ReturningResultsEntityUpdator(queryRunner, this.expressionMap);
             if (this.expressionMap.updateEntity === true &&
                 this.expressionMap.mainAlias!.hasMetadata &&
                 this.expressionMap.whereEntities.length > 0) {
                 this.expressionMap.extraReturningColumns = returningResultsEntityUpdator.getUpdationReturningColumns();
+
+                if (this.expressionMap.extraReturningColumns.length > 0 && this.connection.driver instanceof SqlServerDriver) {
+                    declareSql = this.connection.driver.buildTableVariableDeclaration("@OutputTable", this.expressionMap.extraReturningColumns);
+                    selectOutputSql = `SELECT * FROM @OutputTable`;
+                }
             }
 
             // execute update query
-            const [sql, parameters] = this.getQueryAndParameters();
+            const [updateSql, parameters] = this.getQueryAndParameters();
             const updateResult = new UpdateResult();
-            const result = await queryRunner.query(sql, parameters);
+            const statements = [declareSql, updateSql, selectOutputSql];
+            const result = await queryRunner.query(
+                statements.filter(sql => sql != null).join(";\n\n"),
+                parameters,
+            );
 
-            const driver = queryRunner.connection.driver;
-            if (driver instanceof PostgresDriver) {
+            if (this.connection.driver instanceof PostgresDriver) {
                 updateResult.raw = result[0];
                 updateResult.affected = result[1];
             }

test/functional/query-builder/entity-updation/entity-updation.ts

@@ -99,7 +99,7 @@ describe("query builder > entity updation", () => {
         post.version.should.be.equal(2);
     })));
 
-    it("should not update special entity properties after entity updation if updateEntity is set to true", () => Promise.all(connections.map(async connection => {
+    it("should not update special entity properties after entity updation if updateEntity is set to false", () => Promise.all(connections.map(async connection => {
 
         const post = new Post();
         post.title = "about entity updation in query builder";

test/github-issues/5160/entity/Post.ts

@@ -0,0 +1,32 @@
+import {Entity} from "../../../../src/decorator/entity/Entity";
+import {PrimaryGeneratedColumn} from "../../../../src/decorator/columns/PrimaryGeneratedColumn";
+import {Column} from "../../../../src/decorator/columns/Column";
+import {CreateDateColumn} from "../../../../src/decorator/columns/CreateDateColumn";
+import {UpdateDateColumn} from "../../../../src/decorator/columns/UpdateDateColumn";
+import {VersionColumn} from "../../../../src/decorator/columns/VersionColumn";
+
+@Entity()
+export class Post {
+
+    @PrimaryGeneratedColumn()
+    id: number;
+
+    @Column()
+    title: string;
+
+    @CreateDateColumn()
+    createDate: string;
+
+    @UpdateDateColumn()
+    updateDate: string;
+
+    @Column({ default: 100 })
+    order: number;
+
+    @VersionColumn()
+    version: number;
+
+    @Column({ default: 0 })
+    triggerValue: number;
+
+}

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

@@ -0,0 +1,96 @@
+import "reflect-metadata";
+import {expect} from "chai";
+import {Connection} from "../../../src";
+import {Post} from "./entity/Post";
+import {createTestingConnections, reloadTestingDatabases, closeTestingConnections} from "../../utils/test-utils";
+import {SqlServerDriver} from "../../../src/driver/sqlserver/SqlServerDriver";
+
+describe("github issues > #5160 (MSSQL) DML statement cannot have any enabled triggers if the statement contains an OUTPUT clause without INTO clause", () => {
+
+    let connections: Connection[];
+
+    before(async () => {
+        connections = await createTestingConnections({
+            entities: [Post],
+            schemaCreate: true,
+            dropSchema: true
+        });
+    });
+    beforeEach(async () => {
+        await reloadTestingDatabases(connections);
+
+        return Promise.all(connections.map(async connection => {
+            if (!(connection.driver instanceof SqlServerDriver)) {
+                return;
+            }
+
+            return connection.query(`
+                CREATE OR ALTER TRIGGER issue5160_post
+                ON post AFTER INSERT, UPDATE AS
+                BEGIN
+                    UPDATE post
+                    SET triggerValue = 1
+                    WHERE id IN (SELECT id FROM inserted);
+                END`
+            );
+        }));
+    });
+    after(() => closeTestingConnections(connections));
+
+    it("should update entity model after insertion to MSSQL table with trigger", () => Promise.all(connections.map(async connection => {
+        if (!(connection.driver instanceof SqlServerDriver)) {
+            return;
+        }
+
+        const post = new Post();
+        post.title = "about entity updation in query builder";
+
+        await connection.createQueryBuilder()
+            .insert()
+            .into(Post)
+            .values(post)
+            .execute();
+
+        post.id.should.be.a("number");
+        post.id.should.be.greaterThan(0);
+        post.title.should.be.equal("about entity updation in query builder");
+        post.order.should.be.equal(100);
+        post.createDate.should.be.instanceof(Date);
+        post.updateDate.should.be.instanceof(Date);
+        post.triggerValue.should.be.equal(0, "Returned values from INSERT...OUTPUT will not reflect data modified by triggers");
+
+        // for additional safety, re-fetch entity and check that the trigger fired and updated the field as expected
+        const updatedPost = await connection.createQueryBuilder(Post, "post")
+            .where({ id: post.id })
+            .getOne();
+
+        expect(updatedPost).is.not.undefined;
+        updatedPost!.id.should.be.equal(post.id);
+        updatedPost!.triggerValue.should.be.equal(1);
+    })));
+
+    it("should update entity model after save to MSSQL table with trigger", () => Promise.all(connections.map(async connection => {
+        if (!(connection.driver instanceof SqlServerDriver)) {
+            return;
+        }
+
+        const post = new Post();
+        post.title = "about entity updation in query builder";
+        await connection.manager.save(post);
+        post.version.should.be.equal(1);
+
+        post.title = "changed title";
+        await connection.manager.save(post);
+        post.version.should.be.equal(2);
+        post.triggerValue.should.be.equal(0, "Returned values from UPDATE...OUTPUT will not reflect data modified by triggers");
+
+        // for additional safety, re-fetch entity and check that the trigger fired and updated the field as expected
+        const updatedPost = await connection.createQueryBuilder(Post, "post")
+            .where({ id: post.id })
+            .getOne();
+
+        expect(updatedPost).is.not.undefined;
+        updatedPost!.id.should.be.equal(post.id);
+        updatedPost!.triggerValue.should.be.equal(1);
+    })));
+});