From 3111ba2c6117fec45165fcb62187647ddacf2854 Mon Sep 17 00:00:00 2001
From: Daniel Cormier <dcormier@users.noreply.github.com>
Date: Fri, 27 Sep 2019 15:28:24 -0400
Subject: [PATCH] Implement database/sql/driver.DriverContext

In order to fully instrument SQL operations (including new conns)
by wrapping this driver with something like
`github.com/luna-duclos/instrumentedsql`, this driver needs to
implement `database/sql/driver.DriverContext`.

Now it does.
---
 conn.go           |  9 +++++++++
 connector_test.go | 21 +++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/conn.go b/conn.go
index 55152b12..4f518cf7 100644
--- a/conn.go
+++ b/conn.go
@@ -48,6 +48,15 @@ func (d *Driver) Open(name string) (driver.Conn, error) {
 	return Open(name)
 }
 
+// OpenConnector sets up a new Connector, ready to create connections. name is
+// a connection string. Most users should only use it through database/sql
+// package from the standard library.
+//
+// Implements `database/sql/driver.DriverContext`.
+func (d *Driver) OpenConnector(name string) (driver.Connector, error) {
+	return NewConnector(name)
+}
+
 func init() {
 	sql.Register("postgres", &Driver{})
 }
diff --git a/connector_test.go b/connector_test.go
index 3d2c67b0..57601174 100644
--- a/connector_test.go
+++ b/connector_test.go
@@ -65,3 +65,24 @@ func TestNewConnector_Driver(t *testing.T) {
 	}
 	txn.Rollback()
 }
+
+func TestNewConnector_DriverContext(t *testing.T) {
+	name := ""
+	d := &Driver{}
+	c, err := d.OpenConnector(name)
+	if err != nil {
+		t.Fatal(err)
+	}
+	db, err := c.Connect(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer db.Close()
+	// database/sql might not call our Open at all unless we do something with
+	// the connection
+	txn, err := db.(driver.ConnBeginTx).BeginTx(context.Background(), driver.TxOptions{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	txn.Rollback()
+}