Merge branch 'notify-listen' into dbkit-gm

Author: pedantix

Sources/PostgreSQL/Connection/PostgreSQLConnection+NotifyAndListen.swift

@@ -0,0 +1,29 @@
+import Async
+
+extension PostgreSQLConnection {
+    public func listen(
+        _ channelName: String,
+        handler: @escaping (String) throws -> ()
+        ) throws -> Future<Void> {
+        beforeClose = { conn in
+            let query = PostgreSQLQuery(query: "UNLISTEN \"\(channelName)\";")
+            return conn.send([.query(query)], onResponse: { _ in })
+        }
+        let query = PostgreSQLQuery(query: "LISTEN \"\(channelName)\";")
+        return queue.enqueue([.query(query)], onInput: { message in
+            switch message {
+            case let .notificationResponse(notification):
+                try handler(notification.message)
+            default:
+                break
+            }
+            return false
+        })
+    }
+
+    public 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 })
+    }
+}

Sources/PostgreSQL/Connection/PostgreSQLConnection.swift

@@ -9,7 +9,7 @@ public final class PostgreSQLConnection: DatabaseConnection, BasicWorker {
     }
 
     /// Handles enqueued redis commands and responses.
-    private let queue: QueueHandler<PostgreSQLMessage, PostgreSQLMessage>
+    internal let queue: QueueHandler<PostgreSQLMessage, PostgreSQLMessage>
 
     /// The channel
     private let channel: Channel
@@ -184,10 +184,17 @@ public final class PostgreSQLConnection: DatabaseConnection, BasicWorker {
         }
     }
 
+    internal var beforeClose: ((PostgreSQLConnection) -> Future<Void>)?
+
     /// Closes this client.
     public func close() {
-        isClosed = true
-        channel.close(promise: nil)
+        if let beforeClose = beforeClose {
+            _ = beforeClose(self).then { _ in
+                self.channel.close(mode: CloseMode.all)
+            }
+        } else {
+            channel.close(promise: nil)
+        }
     }
 
     /// Called when this class deinitializes.

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,14 @@
+import Foundation
+
+struct PostgreSQLNotificationResponse: Decodable {
+    /// The message coming from PSQL
+    let message: String
+    init(from decoder: Decoder) throws {
+        let container = try decoder.singleValueContainer()
+        _ = try container.decode(Int32.self) // message length
+        _ = try container.decode(Int32.self) // process id of message
+        let channelId = try container.decode(String.self)
+        let message = try? container.decode(String.self)
+        self.message = message ?? channelId
+    }
+}

Tests/PostgreSQLTests/PostgreSQLConnectionTests.swift

@@ -5,6 +5,7 @@ import PostgreSQL
 import Core
 
 class PostgreSQLConnectionTests: XCTestCase {
+    let defaultTimeout = 5.0
     func testVersion() throws {
         let client = try PostgreSQLConnection.makeTest()
         let results = try client.simpleQuery("SELECT version();").wait()
@@ -342,6 +343,32 @@ class PostgreSQLConnectionTests: XCTestCase {
         _ = try categories.wait()
     }
 
+    func testNotifyAndListen() throws {
+        let completionHandlerExpectation1 = expectation(description: "first completion handler called")
+        let completionHandlerExpectation2 = expectation(description: "final completion handler called")
+        let notifyConn = try PostgreSQLConnection.makeTest()
+        let listenConn = try PostgreSQLConnection.makeTest()
+        let channelName = "Foo"
+        let messageText = "Bar"
+        let finalMessageText = "Baz"
+
+        try listenConn.listen(channelName) { text in
+            if text == messageText {
+                completionHandlerExpectation1.fulfill()
+            } else if text == finalMessageText {
+                completionHandlerExpectation2.fulfill()
+            }
+        }.catch({ err in XCTFail("error \(err)") })
+
+        try notifyConn.notify(channelName, message: messageText).wait()
+        try notifyConn.notify(channelName, message: finalMessageText).wait()
+
+        notifyConn.close()
+        listenConn.close()
+        waitForExpectations(timeout: defaultTimeout)
+    }
+
+
     func testURLParsing() throws {
         let databaseURL = "postgres://username:password@hostname.com:5432/database"
         let config = try PostgreSQLDatabaseConfig(url: databaseURL)
@@ -361,6 +388,7 @@ class PostgreSQLConnectionTests: XCTestCase {
         ("testStruct", testStruct),
         ("testNull", testNull),
         ("testGH24", testGH24),
+        ("testNotifyAndListen", testNotifyAndListen),
         ("testURLParsing", testURLParsing),
     ]
 }
@@ -370,7 +398,8 @@ extension PostgreSQLConnection {
     static func makeTest() throws -> PostgreSQLConnection {
         let hostname: String
         #if Xcode
-        hostname = (try? Process.execute("docker-machine", "ip")) ?? "192.168.99.100"
+        //hostname = (try? Process.execute("docker-machine", "ip")) ?? "192.168.99.100"
+        hostname = "localhost"
         #else
         hostname = "localhost"
         #endif