From 460946d66256f11b10bbd3b3bac99def937143f2 Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Wed, 2 Jan 2019 13:14:34 -0600
Subject: [PATCH] Move notice handling to pgconn

---
 config.go      |  2 ++
 pgconn.go      | 19 +++++++++++++++++++
 pgconn_test.go | 23 +++++++++++++++++++++++
 3 files changed, 44 insertions(+)

diff --git a/config.go b/config.go
index d8872f66..bd1fec9b 100644
--- a/config.go
+++ b/config.go
@@ -40,6 +40,8 @@ type Config struct {
 	// server is acceptable. If this returns an error the connection is closed and the next fallback config is tried. This
 	// allows implementing high availability behavior such as libpq does with target_session_attrs.
 	AfterConnectFunc AfterConnectFunc
+
+	OnNotice NoticeHandler // Callback function called when a notice response is received.
 }
 
 // FallbackConfig is additional settings to attempt a connection with when the primary Config fails to establish a
diff --git a/pgconn.go b/pgconn.go
index b3abe8e0..6b6330dc 100644
--- a/pgconn.go
+++ b/pgconn.go
@@ -48,9 +48,19 @@ func (pe *PgError) Error() string {
 	return pe.Severity + ": " + pe.Message + " (SQLSTATE " + pe.Code + ")"
 }
 
+// Notice represents a notice response message reported by the PostgreSQL server. Be aware that this is distinct from
+// LISTEN/NOTIFY notification.
+type Notice PgError
+
 // DialFunc is a function that can be used to connect to a PostgreSQL server
 type DialFunc func(ctx context.Context, network, addr string) (net.Conn, error)
 
+// NoticeHandler is a function that can handle notices received from the PostgreSQL server. Notices can be received at
+// any time, usually during handling of a query response. The *PgConn is provided so the handler is aware of the origin
+// of the notice, but it must not invoke any query method. Be aware that this is distinct from LISTEN/NOTIFY
+// notification.
+type NoticeHandler func(*PgConn, *Notice)
+
 // ErrTLSRefused occurs when the connection attempt requires TLS and the
 // PostgreSQL server refuses to use TLS
 var ErrTLSRefused = errors.New("server refused TLS connection")
@@ -277,6 +287,10 @@ func (pgConn *PgConn) ReceiveMessage() (pgproto3.BackendMessage, error) {
 			// TODO - close pgConn
 			return nil, errorResponseToPgError(msg)
 		}
+	case *pgproto3.NoticeResponse:
+		if pgConn.Config.OnNotice != nil {
+			pgConn.Config.OnNotice(pgConn, noticeResponseToNotice(msg))
+		}
 	}
 
 	return msg, nil
@@ -858,6 +872,11 @@ func errorResponseToPgError(msg *pgproto3.ErrorResponse) *PgError {
 	}
 }
 
+func noticeResponseToNotice(msg *pgproto3.NoticeResponse) *Notice {
+	pgerr := errorResponseToPgError((*pgproto3.ErrorResponse)(msg))
+	return (*Notice)(pgerr)
+}
+
 // CancelRequest sends a cancel request to the PostgreSQL server. It returns an error if unable to deliver the cancel
 // request, but lack of an error does not ensure that the query was canceled. As specified in the documentation, there
 // is no way to be sure a query was canceled. See https://www.postgresql.org/docs/11/protocol-flow.html#id-1.10.5.7.9
diff --git a/pgconn_test.go b/pgconn_test.go
index 8d6b606a..98ec9664 100644
--- a/pgconn_test.go
+++ b/pgconn_test.go
@@ -551,3 +551,26 @@ func TestCommandTag(t *testing.T) {
 		assert.Equalf(t, tt.rowsAffected, actual, "%d. %v", i, tt.commandTag)
 	}
 }
+
+func TestConnOnNotice(t *testing.T) {
+	t.Parallel()
+
+	config, err := pgconn.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.Nil(t, err)
+
+	var msg string
+	config.OnNotice = func(c *pgconn.PgConn, notice *pgconn.Notice) {
+		msg = notice.Message
+	}
+
+	pgConn, err := pgconn.ConnectConfig(context.Background(), config)
+	require.Nil(t, err)
+	defer closeConn(t, pgConn)
+
+	_, err = pgConn.Exec(context.Background(), `do $$
+begin
+  raise notice 'hello, world';
+end$$;`)
+	require.Nil(t, err)
+	assert.Equal(t, "hello, world", msg)
+}