From 888e735492d25b1c42194213038f4458e4b96aaf Mon Sep 17 00:00:00 2001
From: Eno Compton <enocom@google.com>
Date: Mon, 24 Jul 2023 11:45:06 -0600
Subject: [PATCH] fix: avoid holding lock over IO (#333)

---
 driver/pgxv4/postgres.go | 31 ++++++++++++++-----------------
 1 file changed, 14 insertions(+), 17 deletions(-)

diff --git a/driver/pgxv4/postgres.go b/driver/pgxv4/postgres.go
index ef8a1dae..a6ce6745 100644
--- a/driver/pgxv4/postgres.go
+++ b/driver/pgxv4/postgres.go
@@ -59,32 +59,29 @@ type pgDriver struct {
 //
 // "host=projects/<PROJECT>/locations/<REGION>/clusters/<CLUSTER>/instances/<INSTANCE> user=myuser password=mypass"
 func (p *pgDriver) Open(name string) (driver.Conn, error) {
-	var (
-		dbURI string
-		ok    bool
-	)
-
-	p.mu.RLock()
-	dbURI, ok = p.dbURIs[name]
-	p.mu.RUnlock()
-
-	if ok {
-		return stdlib.GetDefaultDriver().Open(dbURI)
+	dbURI, err := p.dbURI(name)
+	if err != nil {
+		return nil, err
 	}
+	return stdlib.GetDefaultDriver().Open(dbURI)
+
+}
 
+// dbURI registers a driver using the provided DSN. If the name has already
+// been registered, dbURI returns the existing registration.
+func (p *pgDriver) dbURI(name string) (string, error) {
 	p.mu.Lock()
 	defer p.mu.Unlock()
-	// Recheck to ensure dbURI wasn't created between locks
-	dbURI, ok = p.dbURIs[name]
+	dbURI, ok := p.dbURIs[name]
 	if ok {
-		return stdlib.GetDefaultDriver().Open(dbURI)
+		return dbURI, nil
 	}
 
 	config, err := pgx.ParseConfig(name)
 	if err != nil {
-		return nil, err
+		return "", err
 	}
-	instConnName := config.Config.Host // Extract instance URI
+	instConnName := config.Config.Host // Extract instance connection name
 	config.Config.Host = "localhost"   // Replace it with a default value
 	config.DialFunc = func(ctx context.Context, _, _ string) (net.Conn, error) {
 		return p.d.Dial(ctx, instConnName)
@@ -93,5 +90,5 @@ func (p *pgDriver) Open(name string) (driver.Conn, error) {
 	dbURI = stdlib.RegisterConnConfig(config)
 	p.dbURIs[name] = dbURI
 
-	return stdlib.GetDefaultDriver().Open(dbURI)
+	return dbURI, nil
 }