From 61f0710101e0636fef09d85d6683363b475b89e1 Mon Sep 17 00:00:00 2001
From: Jack Christensen <jack@jackchristensen.com>
Date: Sat, 24 Aug 2019 22:12:55 -0500
Subject: [PATCH] Add docs and tests for PgBouncer

---
 README.md         |  2 ++
 doc.go            |  5 +++
 pgbouncer_test.go | 80 +++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+)
 create mode 100644 pgbouncer_test.go

diff --git a/README.md b/README.md
index 6b2cc634..6b09643a 100644
--- a/README.md
+++ b/README.md
@@ -146,6 +146,8 @@ Now you can run the tests:
 PGX_TEST_DATABASE="host=/var/run/postgresql database=pgx_test" go test ./...
 ```
 
+In addition there are tests specific for PgBouncer that will be run if the `PGX_TEST_PGBOUNCER_CONN_STRING` is set.
+
 ## Version Policy
 
 pgx follows semantic versioning for the documented public API on stable releases. This is the prerelease of `v4`. Branch `v3` is the latest stable release. `master` can contain new features or behavior that will change or be removed before being merged to the stable `v3` branch (in practice, this occurs very rarely).
diff --git a/doc.go b/doc.go
index 01d08e40..d2cc314d 100644
--- a/doc.go
+++ b/doc.go
@@ -206,5 +206,10 @@ Logging
 pgx defines a simple logger interface. Connections optionally accept a logger that satisfies this interface. Set
 LogLevel to control logging verbosity. Adapters for github.com/inconshreveable/log15, github.com/sirupsen/logrus, and
 the testing log are provided in the log directory.
+
+PgBouncer
+
+pgx is compatible with PgBouncer in two modes. One is when the connection has a statement cache in "describe" mode. The
+other is when the connection is using the simple protocol. This can be set with the PreferSimpleProtocol config option.
 */
 package pgx
diff --git a/pgbouncer_test.go b/pgbouncer_test.go
new file mode 100644
index 00000000..7ecc2c00
--- /dev/null
+++ b/pgbouncer_test.go
@@ -0,0 +1,80 @@
+package pgx_test
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/jackc/pgconn"
+	"github.com/jackc/pgconn/stmtcache"
+	"github.com/jackc/pgx/v4"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPgbouncerStatementCacheDescribe(t *testing.T) {
+	connString := os.Getenv("PGX_TEST_PGBOUNCER_CONN_STRING")
+	if connString == "" {
+		t.Skipf("Skipping due to missing environment variable %v", "PGX_TEST_PGBOUNCER_CONN_STRING")
+	}
+
+	config := mustParseConfig(t, connString)
+	config.BuildPreparedStatementCache = func(conn *pgconn.PgConn) stmtcache.Cache {
+		return stmtcache.New(conn, stmtcache.ModeDescribe, 1024)
+	}
+
+	testPgbouncer(t, config, 10, 100)
+}
+
+func TestPgbouncerSimpleProtocol(t *testing.T) {
+	connString := os.Getenv("PGX_TEST_PGBOUNCER_CONN_STRING")
+	if connString == "" {
+		t.Skipf("Skipping due to missing environment variable %v", "PGX_TEST_PGBOUNCER_CONN_STRING")
+	}
+
+	config := mustParseConfig(t, connString)
+	config.BuildPreparedStatementCache = nil
+	config.PreferSimpleProtocol = true
+
+	testPgbouncer(t, config, 10, 100)
+}
+
+func testPgbouncer(t *testing.T, config *pgx.ConnConfig, workers, iterations int) {
+	doneChan := make(chan struct{})
+
+	for i := 0; i < workers; i++ {
+		go func() {
+			defer func() { doneChan <- struct{}{} }()
+			conn, err := pgx.ConnectConfig(context.Background(), config)
+			require.Nil(t, err)
+			defer closeConn(t, conn)
+
+			for i := 0; i < iterations; i++ {
+				var i32 int32
+				var i64 int64
+				var f32 float32
+				var s string
+				var s2 string
+				err = conn.QueryRow(context.Background(), "select 1::int4, 2::int8, 3::float4, 'hi'::text").Scan(&i32, &i64, &f32, &s)
+				require.NoError(t, err)
+				assert.Equal(t, int32(1), i32)
+				assert.Equal(t, int64(2), i64)
+				assert.Equal(t, float32(3), f32)
+				assert.Equal(t, "hi", s)
+
+				err = conn.QueryRow(context.Background(), "select 1::int8, 2::float4, 'bye'::text, 4::int4, 'whatever'::text").Scan(&i64, &f32, &s, &i32, &s2)
+				require.NoError(t, err)
+				assert.Equal(t, int64(1), i64)
+				assert.Equal(t, float32(2), f32)
+				assert.Equal(t, "bye", s)
+				assert.Equal(t, int32(4), i32)
+				assert.Equal(t, "whatever", s2)
+			}
+		}()
+	}
+
+	for i := 0; i < workers; i++ {
+		<-doneChan
+	}
+
+}