Add ability to control transactions (#199)

* Conform Database to TransactionControlDatabase

* Fix manifest for CI

* Fix CI

* Use release of FluentKit

Author: 0xTim

Package.swift

@@ -11,7 +11,7 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/vapor/async-kit.git", from: "1.2.0"),
-        .package(url: "https://github.com/vapor/fluent-kit.git", from: "1.27.0"),
+        .package(url: "https://github.com/vapor/fluent-kit.git", from: "1.31.0"),
         .package(url: "https://github.com/vapor/postgres-kit.git", from: "2.5.1"),
     ],
     targets: [

Sources/FluentPostgresDriver/FluentPostgresDatabase.swift

@@ -124,6 +124,29 @@ extension _FluentPostgresDatabase: Database {
     }
 }
 
+extension _FluentPostgresDatabase: TransactionControlDatabase {
+    func beginTransaction() -> EventLoopFuture<Void> {
+        self.database.withConnection { conn in
+            self.logger.log(level: self.sqlLogLevel, "BEGIN")
+            return conn.simpleQuery("BEGIN").map { _ in }
+        }
+    }
+    
+    func commitTransaction() -> NIOCore.EventLoopFuture<Void> {
+        self.database.withConnection { conn in
+            self.logger.log(level: self.sqlLogLevel, "COMMIT")
+            return conn.simpleQuery("COMMIT").map { _ in }
+        }
+    }
+    
+    func rollbackTransaction() -> NIOCore.EventLoopFuture<Void> {
+        self.database.withConnection { conn in
+            self.logger.log(level: self.sqlLogLevel, "ROLLBACK")
+            return conn.simpleQuery("ROLLBACK").map { _ in }
+        }
+    }
+}
+
 extension _FluentPostgresDatabase: SQLDatabase {
     var dialect: SQLDialect {
         PostgresDialect()

Tests/FluentPostgresDriverTests/FluentPostgresTransactionControlTests.swift

@@ -0,0 +1,119 @@
+import Logging
+import FluentKit
+import FluentBenchmark
+import FluentPostgresDriver
+import XCTest
+import PostgresKit
+
+final class FluentPostgresTransactionControlTests: XCTestCase {
+    
+    func testTransactionControl() throws {
+        try (self.db as! TransactionControlDatabase).beginTransaction().wait()
+        
+        let todo1 = Todo(title: "Test")
+        let todo2 = Todo(title: "Test2")
+        try todo1.save(on: self.db).wait()
+        try todo2.save(on: self.db).wait()
+        
+        try (self.db as! TransactionControlDatabase).commitTransaction().wait()
+        
+        let count = try Todo.query(on: self.db).count().wait()
+        XCTAssertEqual(count, 2)
+    }
+    
+    func testRollback() throws {
+        try (self.db as! TransactionControlDatabase).beginTransaction().wait()
+        
+        let todo1 = Todo(title: "Test")
+        
+        try todo1.save(on: self.db).wait()
+        
+        let duplicate = Todo(title: "Test")
+        var errorCaught = false
+        
+        do {
+            try duplicate.create(on: self.db).wait()
+        } catch {
+            errorCaught = true
+            try (self.db as! TransactionControlDatabase).rollbackTransaction().wait()
+        }
+        
+        if !errorCaught {
+            try (self.db as! TransactionControlDatabase).commitTransaction().wait()
+        }
+        
+        XCTAssertTrue(errorCaught)
+        let count2 = try Todo.query(on: self.db).count().wait()
+        XCTAssertEqual(count2, 0)
+    }
+    
+    var benchmarker: FluentBenchmarker {
+        return .init(databases: self.dbs)
+    }
+    var eventLoopGroup: EventLoopGroup!
+    var threadPool: NIOThreadPool!
+    var dbs: Databases!
+    var db: Database {
+        self.benchmarker.database
+    }
+    var postgres: PostgresDatabase {
+        self.db as! PostgresDatabase
+    }
+    
+    override func setUpWithError() throws {
+        try super.setUpWithError()
+        
+        XCTAssert(isLoggingConfigured)
+        self.eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        self.threadPool = NIOThreadPool(numberOfThreads: 1)
+        self.dbs = Databases(threadPool: threadPool, on: self.eventLoopGroup)
+
+        self.dbs.use(.testPostgres(subconfig: "A"), as: .a)
+        self.dbs.use(.testPostgres(subconfig: "B"), as: .b)
+
+        let a = self.dbs.database(.a, logger: Logger(label: "test.fluent.a"), on: self.eventLoopGroup.next())
+        _ = try (a as! PostgresDatabase).query("drop schema public cascade").wait()
+        _ = try (a as! PostgresDatabase).query("create schema public").wait()
+
+        let b = self.dbs.database(.b, logger: Logger(label: "test.fluent.b"), on: self.eventLoopGroup.next())
+        _ = try (b as! PostgresDatabase).query("drop schema public cascade").wait()
+        _ = try (b as! PostgresDatabase).query("create schema public").wait()
+        
+        try CreateTodo().prepare(on: self.db).wait()
+    }
+
+    override func tearDownWithError() throws {
+        try CreateTodo().revert(on: self.db).wait()
+        self.dbs.shutdown()
+        try self.threadPool.syncShutdownGracefully()
+        try self.eventLoopGroup.syncShutdownGracefully()
+        try super.tearDownWithError()
+    }
+    
+    final class Todo: Model {
+        static let schema = "todos"
+
+        @ID
+        var id: UUID?
+
+        @Field(key: "title")
+        var title: String
+
+        init() { }
+        init(title: String) { self.title = title; id = nil }
+    }
+    
+    struct CreateTodo: Migration {
+        func prepare(on database: Database) -> EventLoopFuture<Void> {
+            return database.schema("todos")
+                .id()
+                .field("title", .string, .required)
+                .unique(on: "title")
+                .create()
+        }
+
+        func revert(on database: Database) -> EventLoopFuture<Void> {
+            return database.schema("todos").delete()
+        }
+    }
+}