Skip to content

Commit

Permalink
Add handler for LISTEN/NOTIFY notifications
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
jordanlewis committed May 3, 2020
1 parent 1172019 commit 04f58ea
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 1 deletion.
11 changes: 10 additions & 1 deletion conn.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
}
Expand All @@ -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))
Expand Down
11 changes: 11 additions & 0 deletions notify.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ package pq
// This module contains support for Postgres LISTEN/NOTIFY.

import (
"database/sql/driver"
"errors"
"fmt"
"sync"
Expand All @@ -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
Expand Down

0 comments on commit 04f58ea

Please sign in to comment.