Add Listen/Notify functionality to PostgreSQL

Author: Shaun Hubbard

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)
+                return true
+            default:
+                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 {
     }
 
     /// Handles enqueued redis commands and responses.
-    private let queue: QueueHandler<PostgreSQLMessage, PostgreSQLMessage>
+    internal let queue: QueueHandler<PostgreSQLMessage, PostgreSQLMessage>
 
     /// The channel
     private let channel: Channel
@@ -134,8 +134,17 @@ public final class PostgreSQLConnection {
         }
     }
 
+    internal var beforeClose: ((PostgreSQLConnection) -> Future<Void>)?
+
     /// Closes this client.
     public func close() {
-        channel.close(promise: nil)
+        if let beforeClose = beforeClose {
+            _ = beforeClose(self).then { _ in
+                self.channel.close(mode: CloseMode.all)
+            }
+        } else {
+            channel.close(promise: nil)
+        }
+
     }
 }

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,12 @@
+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
+        message = try container.decode(String.self)
+    }
+}

Tests/PostgreSQLTests/PostgreSQLConnectionTests.swift

@@ -342,6 +342,25 @@ class PostgreSQLConnectionTests: XCTestCase {
         _ = try categories.wait()
     }
 
+    func testNotifyAndListen() throws {
+        let notifyConn = try PostgreSQLConnection.makeTest()
+        let listenConn = try PostgreSQLConnection.makeTest()
+        var messageReceived = false
+        //listenConn
+        let channelName = "Foo"
+        let messageText = "Bar"
+        try listenConn.listen(channelName) { text in
+            messageReceived = text == messageText
+            }.catch({ err in XCTFail("error \(err)") })
+
+        try notifyConn.notify(channelName, message: messageText).wait()
+
+        sleep(1) // Wait for any delay in the message being received
+        notifyConn.close()
+        listenConn.close()
+        XCTAssert(messageReceived)
+    }
+
     static var allTests = [
         ("testVersion", testVersion),
         ("testSelectTypes", testSelectTypes),
@@ -351,6 +370,7 @@ class PostgreSQLConnectionTests: XCTestCase {
         ("testStruct", testStruct),
         ("testNull", testNull),
         ("testGH24", testGH24),
+        ("testNotifyAndListen", testNotifyAndListen)
     ]
 }
 
@@ -359,7 +379,10 @@ 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"
+        // TODO: Switch this back contrib_bootstrap wasnt working for me and I didnt feel like
+        // delaying
+        hostname = "localhost"
         #else
         hostname = "localhost"
         #endif