added the ability to listen to multiple channels

Author: Shaun Hubbard

Sources/PostgreSQL/Connection/PostgreSQLConnection+NotifyAndListen.swift

@@ -1,21 +1,26 @@
 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.
     public func listen(
         _ channelName: String,
         handler: @escaping (String) throws -> ()
         ) throws -> Future<Void> {
-        beforeClose = { conn in
+        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 handler(notification.message)
+                try self.notificationHandlers[notification.channel]?(notification.message)
             default:
                 break
             }
@@ -28,4 +33,10 @@ extension PostgreSQLConnection {
         let query = PostgreSQLQuery(query: "NOTIFY \"\(channelName)\", '\(message)';")
         return send([.query(query)]).map(to: Void.self, { _ in })
     }
+
+    public 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.swift

@@ -44,6 +44,15 @@ public final class PostgreSQLConnection: DatabaseConnection, BasicWorker {
     /// 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
@@ -184,19 +193,24 @@ public final class PostgreSQLConnection: DatabaseConnection, BasicWorker {
         }
     }
 
-    internal var beforeClose: ((PostgreSQLConnection) -> Future<Void>)?
 
     /// Closes this client.
     public func close() {
-        if let beforeClose = beforeClose {
-            _ = beforeClose(self).then { _ in
-                self.channel.close(mode: CloseMode.all)
+        _ = executeCloseHandlersThenClose()
+    }
+
+
+    private  func executeCloseHandlersThenClose() -> Future<Void> {
+        if let beforeClose = closeHandlers.popLast() {
+            return beforeClose(self).then { _ in
+                self.executeCloseHandlersThenClose()
             }
         } else {
-            channel.close(promise: nil)
+            return channel.close(mode: .all)
         }
     }
 
+
     /// Called when this class deinitializes.
     deinit {
         close()

Sources/PostgreSQL/Message/PostgreSQLNotificationResponse.swift

@@ -9,6 +9,5 @@ struct PostgreSQLNotificationResponse: Decodable {
         _ = try container.decode(Int32.self)
         channel = try container.decode(String.self)
         message = try container.decode(String.self)
-        NSLog("Found self \(channel) \(message)")
     }
 }

Tests/PostgreSQLTests/PostgreSQLConnectionTests.swift

@@ -368,6 +368,63 @@ class PostgreSQLConnectionTests: XCTestCase {
         listenConn.close()
     }
 
+    func testNotifyAndListenOnMultipleChannels() 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 = "Fooze"
+        let channelName2 = "Foozalz"
+        let messageText = "Bar"
+        let finalMessageText = "Baz"
+
+        try listenConn.listen(channelName) { text in
+            if text == messageText {
+                completionHandlerExpectation1.fulfill()
+            }
+
+            }.catch({ err in XCTFail("error \(err)") })
+
+        try listenConn.listen(channelName2) { text in
+            if text == finalMessageText {
+                completionHandlerExpectation2.fulfill()
+            }
+            }.catch({ err in XCTFail("error \(err)") })
+
+        try notifyConn.notify(channelName, message: messageText).wait()
+        try notifyConn.notify(channelName2, message: finalMessageText).wait()
+
+        waitForExpectations(timeout: defaultTimeout)
+        notifyConn.close()
+        listenConn.close()
+    }
+
+    func testUnlisten() throws {
+        let completionHandlerExpectation = expectation(description: "notify completion handler called")
+        completionHandlerExpectation.expectedFulfillmentCount = 2
+        completionHandlerExpectation.assertForOverFulfill = true
+
+        let notifyConn = try PostgreSQLConnection.makeTest()
+        let listenConn = try PostgreSQLConnection.makeTest()
+        let channelName = "Foozers"
+        let messageText = "Bar"
+
+        try listenConn.listen(channelName) { text in
+            if text == messageText {
+                completionHandlerExpectation.fulfill()
+            }
+        }.catch({ err in XCTFail("error \(err)") })
+
+        try notifyConn.notify(channelName, message: messageText).wait()
+        try notifyConn.unlisten(channelName, unlistenHandler: {
+            completionHandlerExpectation.fulfill()
+        }).wait()
+        waitForExpectations(timeout: defaultTimeout)
+
+        notifyConn.close()
+        listenConn.close()
+    }
+
     func testURLParsing() throws {
         let databaseURL = "postgres://username:password@hostname.com:5432/database"
         let config = try PostgreSQLDatabaseConfig(url: databaseURL)
@@ -388,6 +445,8 @@ class PostgreSQLConnectionTests: XCTestCase {
         ("testNull", testNull),
         ("testGH24", testGH24),
         ("testNotifyAndListen", testNotifyAndListen),
+        ("testNotifyAndListenOnMultipleChannels", testNotifyAndListenOnMultipleChannels),
+        ("testUnlisten", testUnlisten),
         ("testURLParsing", testURLParsing),
         ]
 }