From 0b62f832b064fae8018a522ea8d5665033749a02 Mon Sep 17 00:00:00 2001
From: fzerorubigd <fzerorubigd@gmail.com>
Date: Thu, 28 Mar 2019 16:31:55 +0100
Subject: [PATCH] [stdlib] Add support for creating a DB from pgx.Pool

Also the configuration used in the Conn structure (used to implement the
driver.Conn interface) stores a ConnConfig which is used only for determining
if the Connection should be used with Simple Protocol or not.
---
 stdlib/opendbpool.go         | 60 ++++++++++++++++++++++++++++++++++++
 stdlib/stdlibutil110_test.go |  9 +++++-
 2 files changed, 68 insertions(+), 1 deletion(-)
 create mode 100644 stdlib/opendbpool.go

diff --git a/stdlib/opendbpool.go b/stdlib/opendbpool.go
new file mode 100644
index 00000000..dd7596ea
--- /dev/null
+++ b/stdlib/opendbpool.go
@@ -0,0 +1,60 @@
+// +build go1.10
+
+package stdlib
+
+import (
+	"context"
+	"database/sql"
+	"database/sql/driver"
+
+	"github.com/jackc/pgx"
+)
+
+// OptionOpenDB options for configuring the driver when opening a new db pool.
+type OptionOpenDBFromPool func(*poolConnector)
+
+// OptionAfterConnect provide a callback for after connect.
+func OptionPreferSimpleProtocol(preferSimpleProtocol bool) OptionOpenDBFromPool {
+	return func(dc *poolConnector) {
+		dc.preferSimpleProtocol = preferSimpleProtocol
+	}
+}
+
+// OpenDBFromPool create a sql.DB connection from a pgx.ConnPool
+func OpenDBFromPool(pool *pgx.ConnPool, opts ...OptionOpenDBFromPool) *sql.DB {
+	c := poolConnector{
+		pool:   pool,
+		driver: pgxDriver,
+	}
+
+	for _, opt := range opts {
+		opt(&c)
+	}
+
+	return sql.OpenDB(c)
+}
+
+type poolConnector struct {
+	pool                 *pgx.ConnPool
+	driver               *Driver
+	preferSimpleProtocol bool
+}
+
+// Connect implement driver.Connector interface
+func (pc poolConnector) Connect(ctx context.Context) (driver.Conn, error) {
+	var (
+		err  error
+		conn *pgx.Conn
+	)
+
+	if conn, err = pc.pool.Acquire(); err != nil {
+		return nil, err
+	}
+
+	return &Conn{conn: conn, driver: pc.driver, connConfig: pgx.ConnConfig{PreferSimpleProtocol: pc.preferSimpleProtocol}}, nil
+}
+
+// Driver implement driver.Connector interface
+func (pc poolConnector) Driver() driver.Driver {
+	return pc.driver
+}
diff --git a/stdlib/stdlibutil110_test.go b/stdlib/stdlibutil110_test.go
index c83b645b..8a5f209c 100644
--- a/stdlib/stdlibutil110_test.go
+++ b/stdlib/stdlibutil110_test.go
@@ -16,5 +16,12 @@ func openDB(t *testing.T) *sql.DB {
 		t.Fatalf("pgx.ParseConnectionString failed: %v", err)
 	}
 
-	return stdlib.OpenDB(config)
+	pool, err := pgx.NewConnPool(pgx.ConnPoolConfig{
+		ConnConfig: config,
+	})
+
+	if err != nil {
+		t.Fatalf("pgx.ParseConnectionString failed: %v", err)
+	}
+	return stdlib.OpenDBFromPool(pool)
 }