From 2647eff5675f7a45d02b82b633580357b11e05ad Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Mon, 25 May 2020 11:49:37 -0500
Subject: [PATCH] Fix ValidateConnect with cancelable context

fixes #40
---
 pgconn.go      | 7 +++++++
 pgconn_test.go | 7 +++++--
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/pgconn.go b/pgconn.go
index 43edbb6b..5644904a 100644
--- a/pgconn.go
+++ b/pgconn.go
@@ -288,6 +288,13 @@ func connect(ctx context.Context, config *Config, fallbackConfig *FallbackConfig
 		case *pgproto3.ReadyForQuery:
 			pgConn.status = connStatusIdle
 			if config.ValidateConnect != nil {
+				// ValidateConnect may execute commands that cause the context to be watched again. Unwatch first to avoid
+				// the watch already in progress panic. This is that last thing done by this method so there is no need to
+				// restart the watch after ValidateConnect returns.
+				//
+				// See https://github.com/jackc/pgconn/issues/40.
+				pgConn.contextWatcher.Unwatch()
+
 				err := config.ValidateConnect(ctx, pgConn)
 				if err != nil {
 					pgConn.conn.Close()
diff --git a/pgconn_test.go b/pgconn_test.go
index 9a75dede..6362c51b 100644
--- a/pgconn_test.go
+++ b/pgconn_test.go
@@ -346,9 +346,12 @@ func TestConnectWithValidateConnectTargetSessionAttrsReadWrite(t *testing.T) {
 	config.ValidateConnect = pgconn.ValidateConnectTargetSessionAttrsReadWrite
 	config.RuntimeParams["default_transaction_read_only"] = "on"
 
-	conn, err := pgconn.ConnectConfig(context.Background(), config)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	conn, err := pgconn.ConnectConfig(ctx, config)
 	if !assert.NotNil(t, err) {
-		conn.Close(context.Background())
+		conn.Close(ctx)
 	}
 }