dbkit 1.0.0 gm; connection pipelining

Author: tanner0101

Sources/PostgreSQL/Connection/PostgreSQLConnection+Query.swift

@@ -5,9 +5,9 @@ extension PostgreSQLConnection {
     public func query(
         _ string: String,
         _ parameters: [PostgreSQLDataConvertible] = []
-    ) throws -> Future<[[PostgreSQLColumn: PostgreSQLData]]> {
+    ) -> Future<[[PostgreSQLColumn: PostgreSQLData]]> {
         var rows: [[PostgreSQLColumn: PostgreSQLData]] = []
-        return try query(string, parameters) { row in
+        return query(string, parameters) { row in
             rows.append(row)
         }.map(to: [[PostgreSQLColumn: PostgreSQLData]].self) {
             return rows
@@ -21,6 +21,22 @@ extension PostgreSQLConnection {
         _ parameters: [PostgreSQLDataConvertible] = [],
         resultFormat: PostgreSQLResultFormat = .binary(),
         onRow: @escaping ([PostgreSQLColumn: PostgreSQLData]) throws -> ()
+    ) -> Future<Void> {
+        return operation {
+            do {
+                return try self._query(string, parameters, resultFormat: resultFormat, onRow: onRow)
+            } catch {
+                return self.eventLoop.newFailedFuture(error: error)
+            }
+        }
+    }
+
+    /// Non-operation bounded query.
+    private func _query(
+        _ string: String,
+        _ parameters: [PostgreSQLDataConvertible] = [],
+        resultFormat: PostgreSQLResultFormat = .binary(),
+        onRow: @escaping ([PostgreSQLColumn: PostgreSQLData]) throws -> ()
     ) throws -> Future<Void> {
         let parameters = try parameters.map { try $0.convertToPostgreSQLData() }
         let parse = PostgreSQLParseRequest(

Sources/PostgreSQL/Connection/PostgreSQLConnection+SimpleQuery.swift

@@ -10,10 +10,14 @@ extension PostgreSQLConnection {
             return rows
         }
     }
-
     /// Sends a simple PostgreSQL query command, returning the parsed results to
     /// the supplied closure.
     public func simpleQuery(_ string: String, onRow: @escaping ([PostgreSQLColumn: PostgreSQLData]) -> ()) -> Future<Void> {
+        return operation { self._simpleQuery(string, onRow: onRow) }
+    }
+
+    /// Non-operation bounded simple query.
+    private func _simpleQuery(_ string: String, onRow: @escaping ([PostgreSQLColumn: PostgreSQLData]) -> ()) -> Future<Void> {
         logger?.log(query: string, parameters: [])
         var currentRow: PostgreSQLRowDescription?
         let query = PostgreSQLQuery(query: string)

Sources/PostgreSQL/Connection/PostgreSQLConnection.swift

@@ -1,9 +1,9 @@
-import Async
 import Crypto
 import NIO
 
 /// A PostgreSQL frontend client.
-public final class PostgreSQLConnection {
+public final class PostgreSQLConnection: DatabaseConnection, BasicWorker {
+    /// See `BasicWorker`.
     public var eventLoop: EventLoop {
         return channel.eventLoop
     }
@@ -17,6 +17,12 @@ public final class PostgreSQLConnection {
     /// If non-nil, will log queries.
     public var logger: PostgreSQLLogger?
 
+    /// See `DatabaseConnection`.
+    public var isClosed: Bool
+
+    /// See `Extendable`.
+    public var extend: Extend
+
     /// Returns a new unique portal name.
     internal var nextPortalName: String {
         defer { uniqueNameCounter = uniqueNameCounter &+ 1 }
@@ -32,21 +38,54 @@ public final class PostgreSQLConnection {
     /// A unique identifier for this connection, used to generate statment and portal names
     private var uniqueNameCounter: UInt8
 
+    /// In-flight `send(...)` futures.
+    private var currentSend: Promise<Void>?
+
+    /// The current query running, if one exists.
+    private var pipeline: Future<Void>
+
     /// Creates a new Redis client on the provided data source and sink.
     init(queue: QueueHandler<PostgreSQLMessage, PostgreSQLMessage>, channel: Channel) {
         self.queue = queue
         self.channel = channel
         self.uniqueNameCounter = 0
+        self.isClosed = false
+        self.extend = [:]
+        self.pipeline = channel.eventLoop.newSucceededFuture(result: ())
+        channel.closeFuture.always {
+            self.isClosed = true
+            if let current = self.currentSend {
+                current.fail(error: closeError)
+            }
+        }
     }
-    
-    deinit {
-        close()
+    /// Sends `PostgreSQLMessage` to the server.
+    func send(_ message: [PostgreSQLMessage]) -> Future<[PostgreSQLMessage]> {
+        var responses: [PostgreSQLMessage] = []
+        return send(message) { response in
+            responses.append(response)
+        }.map(to: [PostgreSQLMessage].self) {
+            return responses
+        }
     }
 
     /// Sends `PostgreSQLMessage` to the server.
     func send(_ messages: [PostgreSQLMessage], onResponse: @escaping (PostgreSQLMessage) throws -> ()) -> Future<Void> {
+        // if currentSend is not nil, previous send has not completed
+        assert(currentSend == nil, "Attempting to call `send(...)` again before previous invocation has completed.")
+
+        // ensure the connection is not closed
+        guard !isClosed else {
+            return eventLoop.newFailedFuture(error: closeError)
+        }
+
+        // create a new promise and store it
+        let promise = eventLoop.newPromise(Void.self)
+        currentSend = promise
+
+        // cascade this enqueue to the newly created promise
         var error: Error?
-        return queue.enqueue(messages) { message in
+        queue.enqueue(messages) { message in
             switch message {
             case .readyForQuery:
                 if let e = error { throw e }
@@ -56,17 +95,28 @@ public final class PostgreSQLConnection {
             default: try onResponse(message)
             }
             return false // request until ready for query
-        }
+        }.cascade(promise: promise)
+
+        // when the promise completes, remove the reference to it
+        promise.futureResult.always { self.currentSend = nil }
+
+        // return the promise's future result (same as `queue.enqueue`)
+        return promise.futureResult
     }
 
-    /// Sends `PostgreSQLMessage` to the server.
-    func send(_ message: [PostgreSQLMessage]) -> Future<[PostgreSQLMessage]> {
-        var responses: [PostgreSQLMessage] = []
-        return send(message) { response in
-            responses.append(response)
-        }.map(to: [PostgreSQLMessage].self) {
-            return responses
+    /// Submits an async task to be pipelined.
+    internal func operation(_ work: @escaping () -> Future<Void>) -> Future<Void> {
+        /// perform this work when the current pipeline future is completed
+        let new = pipeline.then(work)
+
+        /// append this work to the pipeline, discarding errors as the pipeline
+        //// does not care about them
+        pipeline = new.catchMap { err in
+            return ()
         }
+
+        /// return the newly enqueued work's future result
+        return new
     }
 
     /// Authenticates the `PostgreSQLClient` using a username with no password.
@@ -136,6 +186,15 @@ public final class PostgreSQLConnection {
 
     /// Closes this client.
     public func close() {
+        isClosed = true
         channel.close(promise: nil)
     }
+
+    /// Called when this class deinitializes.
+    deinit {
+        close()
+    }
+
 }
+
+private let closeError = PostgreSQLError(identifier: "closed", reason: "Connection is closed.", source: .capture())

Sources/PostgreSQL/Database/PostgreSQLDatabase.swift

@@ -14,7 +14,7 @@ public final class PostgreSQLDatabase: Database {
     }
 
     /// See `Database.makeConnection()`
-    public func makeConnection(on worker: Worker) -> Future<PostgreSQLConnection> {
+    public func newConnection(on worker: Worker) -> Future<PostgreSQLConnection> {
         let config = self.config
         return Future.flatMap(on: worker) {
             return try PostgreSQLConnection.connect(hostname: config.hostname, port: config.port, on: worker) { error in
@@ -31,9 +31,6 @@ public final class PostgreSQLDatabase: Database {
     }
 }
 
-/// A connection created by a `PostgreSQLDatabase`.
-extension PostgreSQLConnection: DatabaseConnection, BasicWorker { }
-
 extension DatabaseIdentifier {
     /// Default identifier for `PostgreSQLDatabase`.
     public static var psql: DatabaseIdentifier<PostgreSQLDatabase> {

Tests/PostgreSQLTests/PostgreSQLConnectionTests.swift

@@ -331,10 +331,10 @@ class PostgreSQLConnectionTests: XCTestCase {
 
 
         /// SELECT
-        let acronyms = try client.query("""
+        let acronyms = client.query("""
         SELECT "acronyms".* FROM "acronyms" WHERE ("acronyms"."id" = $1) LIMIT 1 OFFSET 0
         """, [1])
-        let categories = try client.query("""
+        let categories = client.query("""
         SELECT "categories".* FROM "categories" WHERE ("categories"."id" = $1) LIMIT 1 OFFSET 0
         """, [1])