Move notice handling to pgconn
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user