From a3d6587a64ff95cc3d5830d2d08dbe0638753716 Mon Sep 17 00:00:00 2001
From: Clay Ellis <clayellis@me.com>
Date: Fri, 12 Jun 2020 20:49:46 -0600
Subject: [PATCH] Fix transaction rollback with top-level throw (#90)

* Fix GH89 by adding catchFlatMap to catch transaction closure throws. Add test.

* Add tests to allTests

* Update the contribute script to use docker-compose instead of docker-machine.

Co-authored-by: Tanner <me@tanner.xyz>
---
 ...greSQLDatabase+TransactionSupporting.swift | 12 +++---
 .../FluentPostgreSQLTests.swift               | 43 ++++++++++++++++++-
 contribute_boostrap.sh                        | 23 ++--------
 3 files changed, 52 insertions(+), 26 deletions(-)

diff --git a/Sources/FluentPostgreSQL/PostgreSQLDatabase+TransactionSupporting.swift b/Sources/FluentPostgreSQL/PostgreSQLDatabase+TransactionSupporting.swift
index 86a7a6a..3d0c7bb 100644
--- a/Sources/FluentPostgreSQL/PostgreSQLDatabase+TransactionSupporting.swift
+++ b/Sources/FluentPostgreSQL/PostgreSQLDatabase+TransactionSupporting.swift
@@ -1,14 +1,14 @@
 extension PostgreSQLDatabase: TransactionSupporting {
     /// See `TransactionSupporting`.
     public static func transactionExecute<T>(_ transaction: @escaping (PostgreSQLConnection) throws -> Future<T>, on connection: PostgreSQLConnection) -> Future<T> {
+        func rollback(error: Error) -> Future<T> {
+            return connection.simpleQuery("ROLLBACK").map { throw error }
+        }
+
         return connection.simpleQuery("BEGIN TRANSACTION").flatMap { results in
             return try transaction(connection).flatMap { res in
                 return connection.simpleQuery("END TRANSACTION").transform(to: res)
-            }.catchFlatMap { error in
-                return connection.simpleQuery("ROLLBACK").map { results in
-                    throw error
-                }
-            }
-        }
+            }.catchFlatMap(rollback)
+        }.catchFlatMap(rollback)
     }
 }
diff --git a/Tests/FluentPostgreSQLTests/FluentPostgreSQLTests.swift b/Tests/FluentPostgreSQLTests/FluentPostgreSQLTests.swift
index ec44c86..4497d07 100644
--- a/Tests/FluentPostgreSQLTests/FluentPostgreSQLTests.swift
+++ b/Tests/FluentPostgreSQLTests/FluentPostgreSQLTests.swift
@@ -458,6 +458,45 @@ class FluentPostgreSQLTests: XCTestCase {
         try C.prepare(on: conn).wait()
         defer { try? C.revert(on: conn).wait() }
     }
+
+    // https://github.com/vapor/fluent-postgresql/issues/89
+    func testGH89() throws {
+        let conn = try benchmarker.pool.requestConnection().wait()
+        conn.logger = DatabaseLogger(database: .psql, handler: PrintLogHandler())
+        defer { benchmarker.pool.releaseConnection(conn) }
+
+        try Planet.prepare(on: conn).wait()
+        defer { try? Planet.revert(on: conn).wait() }
+
+        enum SomeError: Error {
+            case error
+        }
+
+        func alwaysThrows() throws {
+            throw SomeError.error
+        }
+
+        var a = Planet(name: "Pluto")
+        a = try a.save(on: conn).wait()
+
+        do {
+            _ = try conn.transaction(on: .psql) { transaction -> Future<Planet> in
+                a.name = "No Longer A Planet"
+                let save = a.save(on: transaction)
+                try alwaysThrows()
+                return save
+            }.wait()
+        } catch {
+            // No-op
+        }
+
+        a = try Planet.query(on: conn)
+            .filter(\.id == a.requireID())
+            .first()
+            .wait()!
+
+        XCTAssertEqual(a.name, "Pluto")
+    }
     
     // https://github.com/vapor/fluent-postgresql/issues/85
     func testGH85() throws {
@@ -555,8 +594,10 @@ class FluentPostgreSQLTests: XCTestCase {
         ("testCustomFilter", testCustomFilter),
         ("testCreateOrUpdate", testCreateOrUpdate),
         ("testEnumArray", testEnumArray),
+        ("testAlterDrop", testAlterDrop),
+        ("testGH89", testGH89),
         ("testGH85", testGH85),
-        ("testGH35", testGH35),
+        ("testGH35", testGH35)
     ]
 }
 
diff --git a/contribute_boostrap.sh b/contribute_boostrap.sh
index 8c84d37..2ee0796 100755
--- a/contribute_boostrap.sh
+++ b/contribute_boostrap.sh
@@ -1,23 +1,8 @@
-echo "💧  starting docker..."
-docker-machine start default
-
-echo "💧  exporting docker machine environment..."
-eval $(docker-machine env default)
-
-echo "💧  cleaning previous vapor-psql dev db..."
-docker stop vapor-psql
-docker rm vapor-psql
-
-echo "💧  creating vapor-psql dev db..."
-docker run --name vapor-psql -e POSTGRES_USER=vapor_username -e POSTGRES_DB=vapor_database -p 5432:5432 -d postgres:latest
-
 echo "💧  generating xcode proj..."
 swift package generate-xcodeproj
 
-echo "💧  add the following env variable to Xcode test scheme:"
-echo ""
-echo "    PSQL_HOSTNAME: `docker-machine ip`"
-echo ""
-
 echo "💧  opening xcode..."
-open *.xcodeproj
\ No newline at end of file
+open *.xcodeproj
+
+echo "💧  starting docker..."
+docker-compose up psql-10