fix: enable returning additional columns with MSSQL (#7864)

typeorm · Jul 8, 2021 · e1db48d · e1db48d
1 parent 8937fd2
commit e1db48d
Show file tree

Hide file tree

Showing 4 changed files with 103 additions and 8 deletions.
diff --git a/src/query-builder/InsertQueryBuilder.ts b/src/query-builder/InsertQueryBuilder.ts
@@ -85,14 +85,30 @@ export class InsertQueryBuilder<Entity> extends QueryBuilder<Entity> {
             // 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);
+
+            const returningColumns: ColumnMetadata[] = [];
+
+            if (Array.isArray(this.expressionMap.returning) && this.expressionMap.mainAlias!.hasMetadata) {
+                for (const columnPath of this.expressionMap.returning) {
+                    returningColumns.push(
+                        ...this.expressionMap.mainAlias!.metadata.findColumnsWithPropertyPath(columnPath)
+                    );
+                }
+            }
+
             if (this.expressionMap.updateEntity === true && this.expressionMap.mainAlias!.hasMetadata) {
                 if (!(valueSets.length > 1 && this.connection.driver instanceof OracleDriver)) {
                     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`;
-                }
+
+                returningColumns.push(...this.expressionMap.extraReturningColumns.filter(
+                    c => !returningColumns.includes(c)
+                ));
+            }
+
+            if (returningColumns.length > 0 && this.connection.driver instanceof SqlServerDriver) {
+                declareSql = this.connection.driver.buildTableVariableDeclaration("@OutputTable", returningColumns);
+                selectOutputSql = `SELECT * FROM @OutputTable`;
             }
             // console.timeEnd(".prepare returning statement");
 

diff --git a/src/query-builder/UpdateQueryBuilder.ts b/src/query-builder/UpdateQueryBuilder.ts
@@ -79,15 +79,30 @@ export class UpdateQueryBuilder<Entity> extends QueryBuilder<Entity> implements
 
             // if update entity mode is enabled we may need extra columns for the returning statement
             const returningResultsEntityUpdator = new ReturningResultsEntityUpdator(queryRunner, this.expressionMap);
+
+            const returningColumns: ColumnMetadata[] = [];
+
+            if (Array.isArray(this.expressionMap.returning) && this.expressionMap.mainAlias!.hasMetadata) {
+                for (const columnPath of this.expressionMap.returning) {
+                    returningColumns.push(
+                        ...this.expressionMap.mainAlias!.metadata.findColumnsWithPropertyPath(columnPath)
+                    );
+                }
+            }
+
             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`;
-                }
+                returningColumns.push(...this.expressionMap.extraReturningColumns.filter(
+                    c => !returningColumns.includes(c)
+                ));
+            }
+
+            if (returningColumns.length > 0 && this.connection.driver instanceof SqlServerDriver) {
+                declareSql = this.connection.driver.buildTableVariableDeclaration("@OutputTable", returningColumns);
+                selectOutputSql = `SELECT * FROM @OutputTable`;
             }
 
             // execute update query

diff --git a/test/github-issues/7100/entity/Post.ts b/test/github-issues/7100/entity/Post.ts
@@ -0,0 +1,17 @@
+import {Entity} from "../../../../src/decorator/entity/Entity";
+import {PrimaryGeneratedColumn} from "../../../../src/decorator/columns/PrimaryGeneratedColumn";
+import {Column} from "../../../../src/decorator/columns/Column";
+
+@Entity()
+export class Post {
+
+    @PrimaryGeneratedColumn()
+    id: number;
+
+    @Column()
+    title: string;
+
+    @Column()
+    text: string;
+
+}
diff --git a/test/github-issues/7100/issue-7100.ts b/test/github-issues/7100/issue-7100.ts
@@ -0,0 +1,47 @@
+import "reflect-metadata";
+import {expect} from "chai";
+import {Connection} from "../../../src";
+import {Post} from "./entity/Post";
+import {createTestingConnections, reloadTestingDatabases, closeTestingConnections} from "../../utils/test-utils";
+
+describe("github issues > #7100 MSSQL error when user requests additional columns to be returned", () => {
+
+    let connections: Connection[];
+
+    before(async () => {
+        connections = await createTestingConnections({
+            entities: [Post],
+            schemaCreate: true,
+            dropSchema: true,
+            enabledDrivers: [ "mssql" ]
+        });
+    });
+    beforeEach(() => reloadTestingDatabases(connections));
+    after(() => closeTestingConnections(connections));
+
+    it("should return user requested columns", () => Promise.all(connections.map(async connection => {
+        const post = new Post();
+        post.title = "title";
+        post.text = "text"
+
+        await connection.createQueryBuilder()
+            .insert()
+            .into(Post)
+            .values(post)
+            .returning(["text"])
+            .execute();
+
+        // Locally we have forgotten what text was set to, must re-fetch
+        post.text = "";
+        await connection.createQueryBuilder(Post, "post")
+            .update()
+            .set({ title: "TITLE" })
+            .returning(["title", "text"])
+            .whereEntity(post)
+            .updateEntity(true)
+            .execute();
+
+        expect(post.title).to.be.equal("TITLE")
+        expect(post.text).to.be.equal("text");
+    })));
+});