Merge pull request #52 from vapor/dbkit-gm

dbkit 1.0.0 gm

Package.swift

@@ -8,22 +8,25 @@ let package = Package(
     ],
     dependencies: [
         // 🌎 Utility package containing tools for byte manipulation, Codable, OS APIs, and debugging.
-        .package(url: "https://github.com/vapor/core.git", from: "3.0.0-rc.2"),
+        .package(url: "https://github.com/vapor/core.git", from: "3.0.0"),
 
         // 🔑 Hashing (BCrypt, SHA, HMAC, etc), encryption, and randomness.
-        .package(url: "https://github.com/vapor/crypto.git", from: "3.0.0-rc.2"),
+        .package(url: "https://github.com/vapor/crypto.git", from: "3.0.0"),
 
         // 🗄 Core services for creating database integrations.
-        .package(url: "https://github.com/vapor/database-kit.git", from: "1.0.0-rc.2"),
+        .package(url: "https://github.com/vapor/database-kit.git", from: "1.0.0"),
 
         // 📦 Dependency injection / inversion of control framework.
-        .package(url: "https://github.com/vapor/service.git", from: "1.0.0-rc.2"),
+        .package(url: "https://github.com/vapor/service.git", from: "1.0.0"),
+
+        // *️⃣ Build SQL queries in Swift.
+        .package(url: "https://github.com/vapor/sql.git", from: "1.0.0"),
 
         // Event-driven network application framework for high performance protocol servers & clients, non-blocking.
         .package(url: "https://github.com/apple/swift-nio.git", from: "1.0.0"),
     ],
     targets: [
-        .target(name: "PostgreSQL", dependencies: ["Async", "Bits", "Core", "Crypto", "DatabaseKit", "NIO", "Service"]),
+        .target(name: "PostgreSQL", dependencies: ["Async", "Bits", "Core", "Crypto", "DatabaseKit", "NIO", "Service", "SQL"]),
         .testTarget(name: "PostgreSQLTests", dependencies: ["Core", "PostgreSQL"]),
     ]
 )

Sources/PostgreSQL/Connection/PostgreSQLConnection+NotifyAndListen.swift

@@ -0,0 +1,38 @@
+import Async
+
+
+extension PostgreSQLConnection {
+    /// Note: after calling `listen'` on a connection, it can no longer handle other database operations. Do not try to send other SQL commands through this connection afterwards.
+    /// IAlso, notifications will only be sent for as long as this connection remains open; you are responsible for opening a new connection to listen on when this one closes.
+    internal func listen(_ channelName: String, handler: @escaping (String) throws -> ()) throws -> Future<Void> {
+        closeHandlers.append({ conn in
+            let query = PostgreSQLQuery(query: "UNLISTEN \"\(channelName)\";")
+            return conn.send([.query(query)], onResponse: { _ in })
+        })
+
+        notificationHandlers[channelName] = { message in
+            try handler(message)
+        }
+        let query = PostgreSQLQuery(query: "LISTEN \"\(channelName)\";")
+        return queue.enqueue([.query(query)], onInput: { message in
+            switch message {
+            case let .notificationResponse(notification):
+                try self.notificationHandlers[notification.channel]?(notification.message)
+            default:
+                break
+            }
+            return false
+        })
+    }
+
+    internal func notify(_ channelName: String, message: String) throws -> Future<Void> {
+        let query = PostgreSQLQuery(query: "NOTIFY \"\(channelName)\", '\(message)';")
+        return send([.query(query)]).map(to: Void.self, { _ in })
+    }
+
+    internal func unlisten(_ channelName: String, unlistenHandler: (() -> Void)? = nil) throws -> Future<Void> {
+        notificationHandlers.removeValue(forKey: channelName)
+        let query = PostgreSQLQuery(query: "UNLISTEN \"\(channelName)\";")
+        return send([.query(query)], onResponse: { _ in unlistenHandler?() })
+    }
+}

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,8 +21,26 @@ 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() }
+        logger?.record(query: string, values: parameters.map { $0.description })
+
         let parse = PostgreSQLParseRequest(
             statementName: "",
             query: string,

Sources/PostgreSQL/Connection/PostgreSQLConnection+SimpleQuery.swift

@@ -10,11 +10,15 @@ 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> {
-        logger?.log(query: string, parameters: [])
+        return operation { self._simpleQuery(string, onRow: onRow) }
+    }
+
+    /// Non-operation bounded simple query.
+    private func _simpleQuery(_ string: String, onRow: @escaping ([PostgreSQLColumn: PostgreSQLData]) -> ()) -> Future<Void> {
+        logger?.record(query: string)
         var currentRow: PostgreSQLRowDescription?
         let query = PostgreSQLQuery(query: string)
         return send([.query(query)]) { message in

Sources/PostgreSQL/Connection/PostgreSQLConnection.swift

@@ -1,21 +1,27 @@
-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
     }
 
     /// Handles enqueued redis commands and responses.
-    private let queue: QueueHandler<PostgreSQLMessage, PostgreSQLMessage>
+    internal let queue: QueueHandler<PostgreSQLMessage, PostgreSQLMessage>
 
     /// The channel
     private let channel: Channel
 
     /// If non-nil, will log queries.
-    public var logger: PostgreSQLLogger?
+    public var logger: DatabaseLogger?
+
+    /// See `DatabaseConnection`.
+    public var isClosed: Bool
+
+    /// See `Extendable`.
+    public var extend: Extend
 
     /// Returns a new unique portal name.
     internal var nextPortalName: String {
@@ -32,21 +38,63 @@ 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>
+
+    /// Block type to be called on close of connection
+    internal typealias CloseHandler = ((PostgreSQLConnection) -> Future<Void>)
+    /// Called on close of the connection
+    internal var closeHandlers = [CloseHandler]()
+    /// Handler type for Notifications
+    internal typealias NotificationHandler = (String) throws -> Void
+    /// Handlers to be stored by channel name
+    internal var notificationHandlers: [String: NotificationHandler] = [:]
+
     /// 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 +104,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.
@@ -134,8 +193,29 @@ public final class PostgreSQLConnection {
         }
     }
 
+
     /// Closes this client.
     public func close() {
-        channel.close(promise: nil)
+        _ = executeCloseHandlersThenClose()
     }
+
+
+    private  func executeCloseHandlersThenClose() -> Future<Void> {
+        if let beforeClose = closeHandlers.popLast() {
+            return beforeClose(self).then { _ in
+                self.executeCloseHandlersThenClose()
+            }
+        } else {
+            return channel.close(mode: .all)
+        }
+    }
+
+
+    /// Called when this class deinitializes.
+    deinit {
+        close()
+    }
+
 }
+
+private let closeError = PostgreSQLError(identifier: "closed", reason: "Connection is closed.", source: .capture())

Sources/PostgreSQL/Database/PostgreSQLDatabase.swift

@@ -1,26 +1,27 @@
 import Async
 
 /// Creates connections to an identified PostgreSQL database.
-public final class PostgreSQLDatabase: Database {
+public final class PostgreSQLDatabase: Database, LogSupporting {
+    /// See `LogSupporting`
+    public static func enableLogging(_ logger: DatabaseLogger, on conn: PostgreSQLConnection) {
+        conn.logger = logger
+    }
+
     /// This database's configuration.
     public let config: PostgreSQLDatabaseConfig
 
-    /// If non-nil, will log queries.
-    public var logger: PostgreSQLLogger?
-
     /// Creates a new `PostgreSQLDatabase`.
     public init(config: PostgreSQLDatabaseConfig) {
         self.config = config
     }
 
     /// 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
                 print("[PostgreSQL] \(error)")
             }.flatMap(to: PostgreSQLConnection.self) { client in
-                client.logger = self.logger
                 return client.authenticate(
                     username: config.username,
                     database: config.database,
@@ -31,9 +32,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> {

Sources/PostgreSQL/Message+Parse/PostgreSQLMessageDecoder.swift

@@ -49,6 +49,7 @@ final class PostgreSQLMessageDecoder: ByteToMessageDecoder {
         let decoder = _PostgreSQLMessageDecoder(data: messageData)
         let message: PostgreSQLMessage
         switch messageType {
+        case .A: message = try .notificationResponse(decoder.decode())
         case .E: message = try .error(decoder.decode())
         case .N: message = try .notice(decoder.decode())
         case .R: message = try .authenticationRequest(decoder.decode())

Sources/PostgreSQL/Message/PostgreSQLMessage.swift

@@ -7,6 +7,8 @@ enum PostgreSQLMessage {
     case error(PostgreSQLDiagnosticResponse)
     /// Identifies the message as a notice.
     case notice(PostgreSQLDiagnosticResponse)
+    /// Identifies the message as a notification response.
+    case notificationResponse(PostgreSQLNotificationResponse)
     /// One of the various authentication request message formats.
     case authenticationRequest(PostgreSQLAuthenticationRequest)
     /// Identifies the message as a password response.

Sources/PostgreSQL/Message/PostgreSQLNotificationResponse.swift

@@ -0,0 +1,13 @@
+import Foundation
+
+struct PostgreSQLNotificationResponse: Decodable {
+    /// The message coming from PSQL
+    let channel: String
+    let message: String
+    init(from decoder: Decoder) throws {
+        let container = try decoder.singleValueContainer()
+        _ = try container.decode(Int32.self)
+        channel = try container.decode(String.self)
+        message = try container.decode(String.self)
+    }
+}

Sources/PostgreSQL/PostgreSQLProvider.swift

@@ -14,7 +14,7 @@ public final class PostgreSQLProvider: Provider {
         try services.register(DatabaseKitProvider())
         services.register(PostgreSQLDatabaseConfig.self)
         services.register(PostgreSQLDatabase.self)
-        var databases = DatabaseConfig()
+        var databases = DatabasesConfig()
         databases.add(database: PostgreSQLDatabase.self, as: .psql)
         services.register(databases)
     }