Add handler for LISTEN/NOTIFY notifications

lib/pq currently has a dedicated set of functionality for subscribing to
LISTEN/NOTIFY messages, but you have to use it on a separate connection.

This commit adds a simple way to register a handler for notifications
that is on an ordinary database connection. This is useful if a user
wants to LISTEN to a channel on the same connection as their ordinary
database connection.

Author: jordanlewis

conn.go

@@ -152,6 +152,9 @@ type conn struct {
 
 	// If not nil, notices will be synchronously sent here
 	noticeHandler func(*Error)
+
+	// If not nil, notifications will be synchronously sent here
+	notificationHandler func(*Notification)
 }
 
 // Handle driver-side settings in parsed connection string.
@@ -977,6 +980,10 @@ func (cn *conn) recv() (t byte, r *readBuf) {
 			if n := cn.noticeHandler; n != nil {
 				n(parseError(r))
 			}
+		case 'A':
+			if n := cn.notificationHandler; n != nil {
+				n(recvNotification(r))
+			}
 		default:
 			return
 		}
@@ -994,7 +1001,9 @@ func (cn *conn) recv1Buf(r *readBuf) byte {
 
 		switch t {
 		case 'A':
-			// ignore
+			if n := cn.notificationHandler; n != nil {
+				n(recvNotification(r))
+			}
 		case 'N':
 			if n := cn.noticeHandler; n != nil {
 				n(parseError(r))

notify.go

@@ -4,6 +4,7 @@ package pq
 // This module contains support for Postgres LISTEN/NOTIFY.
 
 import (
+	"database/sql/driver"
 	"errors"
 	"fmt"
 	"sync"
@@ -29,6 +30,16 @@ func recvNotification(r *readBuf) *Notification {
 	return &Notification{bePid, channel, extra}
 }
 
+// SetNotificationHandler sets the given notification handler on the given
+// connection. A runtime panic occurs if c is not a pq connection. A nil handler
+// may be used to unset it.
+//
+// Note: Notification handlers are executed synchronously by pq meaning commands
+// won't continue to be processed until the handler returns.
+func SetNotificationHandler(c driver.Conn, handler func(*Notification)) {
+	c.(*conn).notificationHandler = handler
+}
+
 const (
 	connStateIdle int32 = iota
 	connStateExpectResponse