From 9b31034fdd66b0a937191ebd253276322ac7b8d6 Mon Sep 17 00:00:00 2001
From: smaher-edb <suraj.maher@enterprisedb.com>
Date: Thu, 13 Jul 2023 11:05:50 +0530
Subject: [PATCH] connect_timeout is not obeyed for sslmode=allow|prefer

connect_timeout given in conn string was not obeyed if sslmode is not specified or equals sslmode=allow|prefer. It took twice the amount of time specified by connect_timeout in conn string. While this behavior is correct if multi-host is provided in conn string, it doesn't look correct in case of single host. This behavior was also not matching with libpq.

Ref: [1672](https://github.com/jackc/pgx/issues/1672)
---
 pgconn.go | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/pgconn.go b/pgconn.go
index 5886c02..e531303 100644
--- a/pgconn.go
+++ b/pgconn.go
@@ -156,12 +156,15 @@ func ConnectConfig(octx context.Context, config *Config) (pgConn *PgConn, err er
 
 	foundBestServer := false
 	var fallbackConfig *FallbackConfig
-	for _, fc := range fallbackConfigs {
+	for i, fc := range fallbackConfigs {
 		// ConnectTimeout restricts the whole connection process.
 		if config.ConnectTimeout != 0 {
-			var cancel context.CancelFunc
-			ctx, cancel = context.WithTimeout(octx, config.ConnectTimeout)
-			defer cancel()
+			// create new context first time or when previous host was different
+			if i == 0 || (fallbackConfigs[i].Host != fallbackConfigs[i-1].Host) {
+				var cancel context.CancelFunc
+				ctx, cancel = context.WithTimeout(octx, config.ConnectTimeout)
+				defer cancel()
+			}
 		} else {
 			ctx = octx
 		}
@@ -599,9 +602,10 @@ func (pgConn *PgConn) PID() uint32 {
 // TxStatus returns the current TxStatus as reported by the server in the ReadyForQuery message.
 //
 // Possible return values:
-//   'I' - idle / not in transaction
-//   'T' - in a transaction
-//   'E' - in a failed transaction
+//
+//	'I' - idle / not in transaction
+//	'T' - in a transaction
+//	'E' - in a failed transaction
 //
 // See https://www.postgresql.org/docs/current/protocol-message-formats.html.
 func (pgConn *PgConn) TxStatus() byte {