publicize listen/notify methods, fixes vapor#40

Author: tanner0101

Sources/PostgreSQL/Connection/PostgreSQLConnection+NotifyAndListen.swift

@@ -1,29 +1,57 @@
 extension PostgreSQLConnection {
-    internal func listen(_ channelName: String, handler: @escaping (String) throws -> ()) throws -> Future<Void> {
-        closeHandlers.append({ conn in
-            return conn.send([.query(.init(query: "UNLISTEN \"\(channelName)\";"))], onResponse: { _ in })
-        })
-
-        notificationHandlers[channelName] = { message in
-            try handler(message)
-        }
-        return queue.enqueue([.query(.init(query: "LISTEN \"\(channelName)\";"))], onInput: { message in
+    /// Begins listening for notifications on a channel.
+    ///
+    ///     LISTEN "<channel name>"
+    ///
+    /// To subscribe to a channel, call `listen(...)` and provide the channel name.
+    ///
+    ///     conn.listen("foo") { message in
+    ///         print(message)
+    ///         return true
+    ///     }
+    ///
+    /// Once a connection is listening, it may not be used to send further queries until `UNLISTEN` is sent.
+    /// To unlisten, return `true` in the callback handler. Returning `false` will continue the subscription.
+    ///
+    /// See `notify(...)` to send messages.
+    ///
+    /// - parameters:
+    ///     - channelName: String identifier for the channel to subscribe to.
+    ///     - handler: Handles incoming String messages. Returning `true` here will end the subscription
+    ///                sending an `UNLISTEN` command.
+    /// - returns: A future that signals completion of the `UNLISTEN` command.
+    public func listen(_ channelName: String, handler: @escaping (String) throws -> (Bool)) -> Future<Void> {
+        let promise = eventLoop.newPromise(Void.self)
+        return queue.enqueue([.query(.init(query: "LISTEN \"\(channelName)\";"))]) { message in
             switch message {
-            case let .notificationResponse(notification):
-                try self.notificationHandlers[notification.channel]?(notification.message)
-            default:
-                break
+            case .close: return false
+            case .readyForQuery: return false
+            case .notification(let notif):
+                if try handler(notif.message) {
+                    self.simpleQuery("UNLISTEN \"\(channelName)\"").transform(to: ()).cascade(promise: promise)
+                    return true
+                } else {
+                    return false
+                }
+            default: throw PostgreSQLError(identifier: "listen", reason: "Unexpected message during listen: \(message).")
             }
-            return false
-        })
-    }
-
-    internal func notify(_ channelName: String, message: String) throws -> Future<Void> {
-        return send([.query(.init(query: "NOTIFY \"\(channelName)\", '\(message)';"))]).map(to: Void.self, { _ in })
+        }.flatMap { promise.futureResult }
     }
 
-    internal func unlisten(_ channelName: String, unlistenHandler: (() -> Void)? = nil) throws -> Future<Void> {
-        notificationHandlers.removeValue(forKey: channelName)
-        return send([.query(.init(query: "UNLISTEN \"\(channelName)\";"))], onResponse: { _ in unlistenHandler?() })
+    /// Sends a notification to a listening connection. Use in conjunction with `listen(...)`.
+    ///
+    ///     NOTIFY "foo" 'hello'
+    ///
+    /// A single connection can be used to send notifications to as many channels as desired.
+    ///
+    ///     conn.notify("foo", message: "hello")
+    ///
+    /// - parameters:
+    ///     - channelName: String identifier for the channel to send to.
+    ///     - message: String message to send to subscribers.
+    /// - returns: A future that signals completion of the send.
+    public func notify(_ channelName: String, message: String) -> Future<Void> {
+        return simpleQuery("NOTIFY \"\(channelName)\", '\(message)'").transform(to: ())
     }
 }
+

Sources/PostgreSQL/Connection/PostgreSQLConnection.swift

@@ -23,6 +23,7 @@ public final class PostgreSQLConnection: DatabaseConnection, BasicWorker {
     /// The channel
     internal let channel: Channel
 
+    /// Previously fetched table name cache
     internal var tableNameCache: TableNameCache?
 
     /// In-flight `send(...)` futures.
@@ -31,15 +32,6 @@ 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 PostgreSQL client on the provided data source and sink.
     init(queue: QueueHandler<PostgreSQLMessage, PostgreSQLMessage>, channel: Channel) {
         self.queue = queue
@@ -122,13 +114,7 @@ public final class PostgreSQLConnection: DatabaseConnection, BasicWorker {
 
     /// Executes close handlers before closing.
     private  func executeCloseHandlersThenClose() -> Future<Void> {
-        if let beforeClose = closeHandlers.popLast() {
-            return beforeClose(self).then { _ in
-                self.executeCloseHandlersThenClose()
-            }
-        } else {
-            return channel.close(mode: .all)
-        }
+        return channel.close(mode: .all)
     }
 
     /// Called when this class deinitializes.

Sources/PostgreSQL/Message/PostgreSQLMessage+0.swift

@@ -37,7 +37,7 @@ enum PostgreSQLMessage {
     case notice(ErrorResponse)
 
     /// Identifies the message as a notification response.
-    case notificationResponse(NotificationResponse)
+    case notification(Notification)
 
     /// Identifies the message as a parameter description.
     case parameterDescription(ParameterDescription)

Sources/PostgreSQL/Message/PostgreSQLMessage+NotificationResponse.swift

@@ -1,5 +1,5 @@
 extension PostgreSQLMessage {
-    struct NotificationResponse {
+    struct Notification {
         /// The message coming from PSQL
         let processID: Int32
 
@@ -13,9 +13,9 @@ extension PostgreSQLMessage {
 
 // MARK: Parse
 
-extension PostgreSQLMessage.NotificationResponse {
+extension PostgreSQLMessage.Notification {
     /// Parses an instance of this message type from a byte buffer.
-    static func parse(from buffer: inout ByteBuffer) throws -> PostgreSQLMessage.NotificationResponse {
+    static func parse(from buffer: inout ByteBuffer) throws -> PostgreSQLMessage.Notification {
         guard let processID = buffer.readInteger(as: Int32.self) else {
             throw PostgreSQLError.protocol(reason: "Could not read process ID from notification response.")
         }

Sources/PostgreSQL/Pipeline/PostgreSQLMessageDecoder.swift

@@ -69,7 +69,7 @@ final class PostgreSQLMessageDecoder: ByteToMessageDecoder {
         case .E: message = try .error(.parse(from: &buffer))
         case .n: message = .noData
         case .N: message = try .notice(.parse(from: &buffer))
-        case .A: message = try .notificationResponse(.parse(from: &buffer))
+        case .A: message = try .notification(.parse(from: &buffer))
         case .t: message = try .parameterDescription(.parse(from: &buffer))
         case .S: message = try .parameterStatus(.parse(from: &buffer))
         case .one: message = .parseComplete

Tests/PostgreSQLTests/PostgreSQLConnectionTests.swift

@@ -523,6 +523,19 @@ class PostgreSQLConnectionTests: XCTestCase {
         default: XCTFail("invalid row count")
         }
     }
+    
+    func testListen() throws {
+        let conn = try PostgreSQLConnection.makeTest(transport: .cleartext)
+        let done = conn.listen("foo") { message in
+            XCTAssertEqual(message, "hi")
+            return true
+        }
+        do {
+            let conn = try PostgreSQLConnection.makeTest(transport: .cleartext)
+            _ = try conn.notify("foo", message: "hi").wait()
+        }
+        try done.wait()
+    }
 
     static var allTests = [
         ("testVersion", testVersion),