From efbbde0487889b0ab88369faea24fa15a14327f4 Mon Sep 17 00:00:00 2001
From: Alex Gaynor <alex@alloy.us>
Date: Thu, 23 Jan 2020 14:19:40 -0600
Subject: [PATCH] Fixes #664 -- implement DriverContext for stdlib

---
 stdlib/sql.go | 56 +++++++++++++++++++++++++++++++++++----------------
 1 file changed, 39 insertions(+), 17 deletions(-)

diff --git a/stdlib/sql.go b/stdlib/sql.go
index 5803dd24..f33a835e 100644
--- a/stdlib/sql.go
+++ b/stdlib/sql.go
@@ -126,29 +126,18 @@ type Driver struct {
 }
 
 func (d *Driver) Open(name string) (driver.Conn, error) {
-	var connConfig *pgx.ConnConfig
-
-	d.configMutex.Lock()
-	connConfig = d.configs[name]
-	d.configMutex.Unlock()
-
-	if connConfig == nil {
-		var err error
-		connConfig, err = pgx.ParseConfig(name)
-		if err != nil {
-			return nil, err
-		}
-	}
-
 	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) // Ensure eventual timeout
 	defer cancel()
-	conn, err := pgx.ConnectConfig(ctx, connConfig)
+
+	connector, err := d.OpenConnector(name)
 	if err != nil {
 		return nil, err
 	}
+	return connector.Connect(ctx)
+}
 
-	c := &Conn{conn: conn, driver: d, connConfig: *connConfig}
-	return c, nil
+func (d *Driver) OpenConnector(name string) (driver.Connector, error) {
+	return &driverConnector{driver: d, name: name}, nil
 }
 
 func (d *Driver) registerConnConfig(c *pgx.ConnConfig) string {
@@ -165,6 +154,39 @@ func (d *Driver) unregisterConnConfig(connStr string) {
 	d.configMutex.Unlock()
 }
 
+type driverConnector struct {
+	driver *Driver
+	name   string
+}
+
+func (dc *driverConnector) Connect(ctx context.Context) (driver.Conn, error) {
+	var connConfig *pgx.ConnConfig
+
+	dc.driver.configMutex.Lock()
+	connConfig = dc.driver.configs[dc.name]
+	dc.driver.configMutex.Unlock()
+
+	if connConfig == nil {
+		var err error
+		connConfig, err = pgx.ParseConfig(dc.name)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	conn, err := pgx.ConnectConfig(ctx, connConfig)
+	if err != nil {
+		return nil, err
+	}
+
+	c := &Conn{conn: conn, driver: dc.driver, connConfig: *connConfig}
+	return c, nil
+}
+
+func (dc *driverConnector) Driver() driver.Driver {
+	return dc.driver
+}
+
 // RegisterConnConfig registers a ConnConfig and returns the connection string to use with Open.
 func RegisterConnConfig(c *pgx.ConnConfig) string {
 	return pgxDriver.registerConnConfig(c)