From c7cfd9485156742925128aa58be2c1bf8ea303d2 Mon Sep 17 00:00:00 2001
From: Braden Groom <braden.groom@reddit.com>
Date: Tue, 17 Oct 2023 21:43:08 -0500
Subject: [PATCH] Run postgres datastore tests with pgbouncer

---
 internal/testserver/datastore/postgres.go | 82 +++++++++++++++++------
 1 file changed, 63 insertions(+), 19 deletions(-)

diff --git a/internal/testserver/datastore/postgres.go b/internal/testserver/datastore/postgres.go
index 64001c70a8..3ed4535458 100644
--- a/internal/testserver/datastore/postgres.go
+++ b/internal/testserver/datastore/postgres.go
@@ -11,6 +11,7 @@ import (
 	"github.com/google/uuid"
 	"github.com/jackc/pgx/v5"
 	"github.com/ory/dockertest/v3"
+	"github.com/ory/dockertest/v3/docker"
 	"github.com/stretchr/testify/require"
 
 	pgmigrations "github.com/authzed/spicedb/internal/datastore/postgres/migrations"
@@ -36,42 +37,85 @@ func RunPostgresForTestingWithCommitTimestamps(t testing.TB, bridgeNetworkName s
 	pool, err := dockertest.NewPool("")
 	require.NoError(t, err)
 
-	name := fmt.Sprintf("postgres-%s", uuid.New().String())
+	if bridgeNetworkName == "" {
+		network, err := pool.Client.CreateNetwork(docker.CreateNetworkOptions{Name: "pgbouncer_postgres_network"})
+		if err != nil {
+			require.NoError(t, err)
+		}
+		bridgeNetworkName = network.ID
+		t.Cleanup(func() {
+			pool.Client.RemoveNetwork(network.ID)
+		})
+	}
 
-	cmd := []string{"-N", "500"} // Max Connections
+	uniqueID := uuid.New().String()
+	postgresName := fmt.Sprintf("postgres-%s", uniqueID)
+	pgbouncerName := fmt.Sprintf("pgbouncer-%s", uniqueID)
+
+	maxConnections := "500"
+	dbUser := "postgres"
+	dbPassword := "secret"
+	postgresPort := "5432"
+	pgbouncerPort := "6432"
+
+	cmd := []string{"-N", maxConnections}
 	if withCommitTimestamps {
 		cmd = append(cmd, "-c", "track_commit_timestamp=1")
 	}
 
-	resource, err := pool.RunWithOptions(&dockertest.RunOptions{
-		Name:         name,
-		Repository:   "postgres",
-		Tag:          pgVersion,
-		Env:          []string{"POSTGRES_PASSWORD=secret", "POSTGRES_DB=defaultdb"},
-		ExposedPorts: []string{"5432/tcp"},
+	postgres, err := pool.RunWithOptions(&dockertest.RunOptions{
+		Name:       postgresName,
+		Repository: "postgres",
+		Tag:        pgVersion,
+		Env: []string{
+			"POSTGRES_USER=" + dbUser,
+			"POSTGRES_PASSWORD=" + dbPassword,
+			// use md5 auth to align postgres and pgbouncer auth methods
+			"POSTGRES_HOST_AUTH_METHOD=md5",
+			"POSTGRES_INITDB_ARGS=--auth=md5",
+		},
+		ExposedPorts: []string{postgresPort},
 		NetworkID:    bridgeNetworkName,
 		Cmd:          cmd,
 	})
 	require.NoError(t, err)
 
+	pgbouncer, err := pool.RunWithOptions(&dockertest.RunOptions{
+		Name:       pgbouncerName,
+		Repository: "edoburu/pgbouncer",
+		Tag:        "latest",
+		Env: []string{
+			"DB_USER=" + dbUser,
+			"DB_PASSWORD=" + dbPassword,
+			"DB_HOST=" + postgresName,
+			"DB_PORT=" + postgresPort,
+			"LISTEN_PORT=" + pgbouncerPort,
+			"DB_NAME=*",     // Needed to make pgbouncer okay with the randomly named databases generated by the test suite
+			"AUTH_TYPE=md5", // use the same auth type as postgres
+			"MAX_CLIENT_CONN=" + maxConnections,
+			// params needed for spicedb
+			"POOL_MODE=session", // https://github.com/authzed/spicedb/issues/1217
+			// Indicate to pgbouncer that this startup param is handled by postgres and should be ignored by pgbouncer.
+			"IGNORE_STARTUP_PARAMETERS=plan_cache_mode",
+		},
+		ExposedPorts: []string{pgbouncerPort + "/tcp"},
+		NetworkID:    bridgeNetworkName,
+	})
+	require.NoError(t, err)
+
 	builder := &postgresTester{
 		hostname:        "localhost",
-		creds:           "postgres:secret",
+		creds:           dbUser + ":" + dbPassword,
 		targetMigration: targetMigration,
+		port:            pgbouncer.GetPort(pgbouncerPort + "/tcp"),
 	}
 	t.Cleanup(func() {
-		require.NoError(t, pool.Purge(resource))
+		require.NoError(t, pool.Purge(pgbouncer))
+		require.NoError(t, pool.Purge(postgres))
 	})
 
-	port := resource.GetPort(fmt.Sprintf("%d/tcp", 5432))
-	if bridgeNetworkName != "" {
-		builder.hostname = name
-		builder.port = "5432"
-	} else {
-		builder.port = port
-	}
-
-	uri := fmt.Sprintf("postgres://%s@localhost:%s/defaultdb?sslmode=disable", builder.creds, port)
+	uri := fmt.Sprintf("postgresql://%s@localhost:%s/?sslmode=disable", builder.creds, builder.port)
+	fmt.Println(uri)
 	require.NoError(t, pool.Retry(func() error {
 		var err error
 		ctx, cancelConnect := context.WithTimeout(context.Background(), dockerBootTimeout)