Repository: jackc/pgx
Pull request #1835 (2023-12-13 21:24:46)

pgconn: normalize more connection errors
Normalize the error that is returned by startTLS in pgconn.connect. This makes it possible to determine if the error was a context error. Before this change the returned error for a context that is canceled or hits the deadline would be `tls error (read tcp 127.0.0.1:55914->127.0.0.1:5432: i/o timeout)`.

This change also uses errors.As to unwrap errors in normalizeTimeoutError. This allows the function to handle the `writeError` type that shows up as `failed to write startup message (read tcp 127.0.0.1:55914->127.0.0.1:5432: i/o timeout)`.


Patch:
diff --git a/pgconn/config.go b/pgconn/config.go
index db0170e0..942a864a 100644
--- a/pgconn/config.go
+++ b/pgconn/config.go
@@ -60,6 +60,11 @@ type Config struct {
 	// OnNotification is a callback function called when a notification from the LISTEN/NOTIFY system is received.
 	OnNotification NotificationHandler
 
+	// OnPgError is a callback function called when a Postgres error is received by the server. The default handler will close
+	// the connection on any FATAL errors. If you override this handler you should call the previously set handler or ensure
+	// that you close on FATAL errors by returning false.
+	OnPgError PgErrorHandler
+
 	createdByParseConfig bool // Used to enforce created by ParseConfig rule.
 }
 
@@ -261,6 +266,13 @@ func ParseConfigWithOptions(connString string, options ParseConfigOptions) (*Con
 		BuildFrontend: func(r io.Reader, w io.Writer) *pgproto3.Frontend {
 			return pgproto3.NewFrontend(r, w)
 		},
+		OnPgError: func(_ *PgConn, pgErr *PgError) bool {
+			// we want to automatically close any fatal errors
+			if strings.EqualFold(pgErr.Severity, "FATAL") {
+				return false
+			}
+			return true
+		},
 	}
 
 	if connectTimeoutSetting, present := settings["connect_timeout"]; present {
diff --git a/pgconn/errors.go b/pgconn/errors.go
index 3c54bbec..61957829 100644
--- a/pgconn/errors.go
+++ b/pgconn/errors.go
@@ -107,14 +107,15 @@ func (e *parseConfigError) Unwrap() error {
 }
 
 func normalizeTimeoutError(ctx context.Context, err error) error {
-	if err, ok := err.(net.Error); ok && err.Timeout() {
+	var netErr net.Error
+	if errors.As(err, &netErr) && netErr.Timeout() {
 		if ctx.Err() == context.Canceled {
 			// Since the timeout was caused by a context cancellation, the actual error is context.Canceled not the timeout error.
 			return context.Canceled
 		} else if ctx.Err() == context.DeadlineExceeded {
 			return &errTimeout{err: ctx.Err()}
 		} else {
-			return &errTimeout{err: err}
+			return &errTimeout{err: netErr}
 		}
 	}
 	return err
diff --git a/pgconn/pgconn.go b/pgconn/pgconn.go
index 1ccdc4db..803b41d1 100644
--- a/pgconn/pgconn.go
+++ b/pgconn/pgconn.go
@@ -52,6 +52,12 @@ type LookupFunc func(ctx context.Context, host string) (addrs []string, err erro
 // BuildFrontendFunc is a function that can be used to create Frontend implementation for connection.
 type BuildFrontendFunc func(r io.Reader, w io.Writer) *pgproto3.Frontend
 
+// PgErrorHandler is a function that handles errors returned from Postgres. This function must return true to keep
+// the connection open. Returning false will cause the connection to be closed immediately. You should return
+// false on any FATAL-severity errors. This will not receive network errors. The *PgConn is provided so the handler is
+// aware of the origin of the error, but it must not invoke any query method.
+type PgErrorHandler func(*PgConn, *PgError) bool
+
 // 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
@@ -289,7 +295,7 @@ func connect(ctx context.Context, config *Config, fallbackConfig *FallbackConfig
 		pgConn.contextWatcher.Unwatch() // Always unwatch `netConn` after TLS.
 		if err != nil {
 			netConn.Close()
-			return nil, &connectError{config: config, msg: "tls error", err: err}
+			return nil, &connectError{config: config, msg: "tls error", err: normalizeTimeoutError(ctx, err)}
 		}
 
 		pgConn.conn = nbTLSConn
@@ -547,11 +553,12 @@ func (pgConn *PgConn) receiveMessage() (pgproto3.BackendMessage, error) {
 	case *pgproto3.ParameterStatus:
 		pgConn.parameterStatuses[msg.Name] = msg.Value
 	case *pgproto3.ErrorResponse:
-		if msg.Severity == "FATAL" {
+		err := ErrorResponseToPgError(msg)
+		if pgConn.config.OnPgError != nil && !pgConn.config.OnPgError(pgConn, err) {
 			pgConn.status = connStatusClosed
 			pgConn.conn.Close() // Ignore error as the connection is already broken and there is already an error to return.
 			close(pgConn.cleanupDone)
-			return nil, ErrorResponseToPgError(msg)
+			return nil, err
 		}
 	case *pgproto3.NoticeResponse:
 		if pgConn.config.OnNotice != nil {


Test patch:
diff --git a/pgconn/pgconn_test.go b/pgconn/pgconn_test.go
index 80621930..7ca3992e 100644
--- a/pgconn/pgconn_test.go
+++ b/pgconn/pgconn_test.go
@@ -3148,6 +3148,41 @@ func TestPipelineCloseDetectsUnsyncedRequests(t *testing.T) {
 	require.EqualError(t, err, "pipeline has unsynced requests")
 }
 
+func TestConnOnPgError(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+
+	config, err := pgconn.ParseConfig(os.Getenv("PGX_TEST_DATABASE"))
+	require.NoError(t, err)
+	config.OnPgError = func(c *pgconn.PgConn, pgErr *pgconn.PgError) bool {
+		require.NotNil(t, c)
+		require.NotNil(t, pgErr)
+		// close connection on undefined tables only
+		if pgErr.Code == "42P01" {
+			return false
+		}
+		return true
+	}
+
+	pgConn, err := pgconn.ConnectConfig(ctx, config)
+	require.NoError(t, err)
+	defer closeConn(t, pgConn)
+
+	_, err = pgConn.Exec(ctx, "select 'Hello, world'").ReadAll()
+	assert.NoError(t, err)
+	assert.False(t, pgConn.IsClosed())
+
+	_, err = pgConn.Exec(ctx, "select 1/0").ReadAll()
+	assert.Error(t, err)
+	assert.False(t, pgConn.IsClosed())
+
+	_, err = pgConn.Exec(ctx, "select * from non_existant_table").ReadAll()
+	assert.Error(t, err)
+	assert.True(t, pgConn.IsClosed())
+}
+
 func Example() {
 	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
 	defer cancel()