Postgres transaction extension for Vapor4 and Fluent4

Postgres+Transaction.swift

//
//  Postgres+Transaction.swift
//
//  Created by Mihael Isaev on 14.01.2020.
//

import Vapor
import FluentKit
import PostgresKit

public protocol PostgresTransactionable {
    var psql: EventLoopFuture<PostgresDatabase> { get }
}

extension Database {
    var psql: EventLoopFuture<PostgresDatabase> {
        eventLoop.future().flatMapThrowing { db -> PostgresDatabase in
            guard let db = self as? PostgresDatabase else {
                throw Abort(.internalServerError, reason: "Unable to reach database for transaction")
            }
            return db
        }
    }
}

extension Application: PostgresTransactionable {
    public var psql: EventLoopFuture<PostgresDatabase> {
        db(.psql).psql
    }
}

extension Request: PostgresTransactionable {
    public var psql: EventLoopFuture<PostgresDatabase> {
        db(.psql).psql
    }
}

extension PostgresTransactionable {
    public func transaction<V>(callback: @escaping (PostgresDatabase) -> EventLoopFuture<V>) -> EventLoopFuture<V> {
        psql.flatMap {
            $0.query("BEGIN;").transform(to: $0)
        }.flatMap { db in
            callback(db).flatMapError { error in
                db.query("ROLLBACK;").flatMapThrowing { _ in
                    throw error
                }
            }.flatMap { v in
                db.query("COMMIT;").transform(to: v)
            }
        }
    }
}

// MARK: - Additional Fluent conformances

extension Model {
    public static func query(on database: PostgresDatabase) -> FluentKit.QueryBuilder<Self> {
        query(on: database as! Database)
    }
    
    public static func find(_ id: Self.IDValue?, on database: PostgresDatabase) -> NIO.EventLoopFuture<Self?> {
        find(id, on: database as! Database)
    }

    public func save(on database: PostgresDatabase) -> NIO.EventLoopFuture<Void> {
        save(on: database as! Database)
    }

    public func create(on database: PostgresDatabase) -> NIO.EventLoopFuture<Void> {
        create(on: database as! Database)
    }

    public func update(on database: PostgresDatabase) -> NIO.EventLoopFuture<Void> {
        update(on: database as! Database)
    }

    public func delete(force: Bool = false, on database: PostgresDatabase) -> NIO.EventLoopFuture<Void> {
        delete(on: database as! Database)
    }

    public func restore(on database: PostgresDatabase) -> NIO.EventLoopFuture<Void> {
        restore(on: database as! Database)
    }
}

/// HOW TO USE:
/// ```swift
/// app.transaction { db in
///     User(email: "[email protected]").create(on: db)
/// }
/// ```