Run postgres datastore tests with pgbouncer

.github/workflows/build-test.yaml

@@ -94,7 +94,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        datastore: ["crdb", "mysql", "postgres", "spanner"]
+        datastore: ["crdb", "mysql", "postgres", "pgbouncer", "spanner"]
     steps:
       - uses: "actions/checkout@v3"
       - uses: "authzed/actions/setup-go@main"

internal/testserver/datastore/postgres.go

@@ -5,12 +5,14 @@ package datastore
 
 import (
 	"context"
+	"flag"
 	"fmt"
 	"testing"
 
 	"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"
@@ -19,12 +21,34 @@ import (
 	"github.com/authzed/spicedb/pkg/secrets"
 )
 
+const (
+	POSTGRES_TEST_USER            = "postgres"
+	POSTGRES_TEST_PASSWORD        = "secret"
+	POSTGRES_TEST_PORT            = "5432"
+	POSTGRES_TEST_MAX_CONNECTIONS = "500"
+	PGBOUNCER_TEST_PORT           = "6432"
+)
+
+var enablePgbouncer = flag.Bool("enablePgbouncer", false, "run pgbouncer in front of postgres datastore tests")
+
+func TestMain(m *testing.M) {
+	flag.Parse()
+}
+
+type container struct {
+	hostHostname      string
+	hostPort          string
+	containerHostname string
+	containerPort     string
+}
+
 type postgresTester struct {
-	conn            *pgx.Conn
-	hostname        string
-	port            string
-	creds           string
-	targetMigration string
+	container
+	hostConn             *pgx.Conn
+	creds                string
+	targetMigration      string
+	pgbouncerProxy       *container
+	useContainerHostname bool
 }
 
 // RunPostgresForTesting returns a RunningEngineForTest for postgres
@@ -36,52 +60,58 @@ func RunPostgresForTestingWithCommitTimestamps(t testing.TB, bridgeNetworkName s
 	pool, err := dockertest.NewPool("")
 	require.NoError(t, err)
 
-	name := fmt.Sprintf("postgres-%s", uuid.New().String())
+	bridgeSupplied := bridgeNetworkName != ""
+	if *enablePgbouncer && !bridgeSupplied {
+		// We will need a network bridge if we're running pgbouncer
+		bridgeNetworkName = createNetworkBridge(t, pool)
+	}
 
-	cmd := []string{"-N", "500"} // Max Connections
+	postgresContainerHostname := fmt.Sprintf("postgres-%s", uuid.New().String())
+
+	cmd := []string{"-N", POSTGRES_TEST_MAX_CONNECTIONS}
 	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:       postgresContainerHostname,
+		Repository: "postgres",
+		Tag:        pgVersion,
+		Env: []string{
+			"POSTGRES_USER=" + POSTGRES_TEST_USER,
+			"POSTGRES_PASSWORD=" + POSTGRES_TEST_PASSWORD,
+			// use md5 auth to align postgres and pgbouncer auth methods
+			"POSTGRES_HOST_AUTH_METHOD=md5",
+			"POSTGRES_INITDB_ARGS=--auth=md5",
+		},
+		ExposedPorts: []string{POSTGRES_TEST_PORT + "/tcp"},
 		NetworkID:    bridgeNetworkName,
 		Cmd:          cmd,
 	})
 	require.NoError(t, err)
+	t.Cleanup(func() {
+		require.NoError(t, pool.Purge(postgres))
+	})
 
 	builder := &postgresTester{
-		hostname:        "localhost",
-		creds:           "postgres:secret",
-		targetMigration: targetMigration,
+		container: container{
+			hostHostname:      "localhost",
+			hostPort:          postgres.GetPort(POSTGRES_TEST_PORT + "/tcp"),
+			containerHostname: postgresContainerHostname,
+			containerPort:     POSTGRES_TEST_PORT,
+		},
+		creds:                POSTGRES_TEST_USER + ":" + POSTGRES_TEST_PASSWORD,
+		targetMigration:      targetMigration,
+		useContainerHostname: bridgeSupplied,
 	}
-	t.Cleanup(func() {
-		require.NoError(t, pool.Purge(resource))
-	})
 
-	port := resource.GetPort(fmt.Sprintf("%d/tcp", 5432))
-	if bridgeNetworkName != "" {
-		builder.hostname = name
-		builder.port = "5432"
-	} else {
-		builder.port = port
+	if *enablePgbouncer {
+		// if we are running with pgbouncer enabled then set it up
+		builder.runPgbouncerForTesting(t, pool, bridgeNetworkName)
 	}
 
-	uri := fmt.Sprintf("postgres://%s@localhost:%s/defaultdb?sslmode=disable", builder.creds, port)
-	require.NoError(t, pool.Retry(func() error {
-		var err error
-		ctx, cancelConnect := context.WithTimeout(context.Background(), dockerBootTimeout)
-		defer cancelConnect()
-		builder.conn, err = pgx.Connect(ctx, uri)
-		if err != nil {
-			return err
-		}
-		return nil
-	}))
+	builder.hostConn = builder.initializeHostConnection(t, pool)
+
 	return builder
 }
 
@@ -91,14 +121,15 @@ func (b *postgresTester) NewDatabase(t testing.TB) string {
 
 	newDBName := "db" + uniquePortion
 
-	_, err = b.conn.Exec(context.Background(), "CREATE DATABASE "+newDBName)
+	_, err = b.hostConn.Exec(context.Background(), "CREATE DATABASE "+newDBName)
 	require.NoError(t, err)
 
+	hostname, port := b.getHostnameAndPort()
 	return fmt.Sprintf(
 		"postgres://%s@%s:%s/%s?sslmode=disable",
 		b.creds,
-		b.hostname,
-		b.port,
+		hostname,
+		port,
 		newDBName,
 	)
 }
@@ -113,3 +144,94 @@ func (b *postgresTester) NewDatastore(t testing.TB, initFunc InitFunc) datastore
 
 	return initFunc("postgres", connectStr)
 }
+
+func createNetworkBridge(t testing.TB, pool *dockertest.Pool) string {
+	bridgeNetworkName := fmt.Sprintf("bridge-%s", uuid.New().String())
+	network, err := pool.Client.CreateNetwork(docker.CreateNetworkOptions{Name: bridgeNetworkName})
+
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		pool.Client.RemoveNetwork(network.ID)
+	})
+
+	return bridgeNetworkName
+}
+
+func (b *postgresTester) runPgbouncerForTesting(t testing.TB, pool *dockertest.Pool, bridgeNetworkName string) {
+	uniqueID := uuid.New().String()
+	pgbouncerContainerHostname := fmt.Sprintf("pgbouncer-%s", uniqueID)
+
+	pgbouncer, err := pool.RunWithOptions(&dockertest.RunOptions{
+		Name:       pgbouncerContainerHostname,
+		Repository: "edoburu/pgbouncer",
+		Tag:        "latest",
+		Env: []string{
+			"DB_USER=" + POSTGRES_TEST_USER,
+			"DB_PASSWORD=" + POSTGRES_TEST_PASSWORD,
+			"DB_HOST=" + b.containerHostname,
+			"DB_PORT=" + b.containerPort,
+			"LISTEN_PORT=" + PGBOUNCER_TEST_PORT,
+			"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=" + POSTGRES_TEST_MAX_CONNECTIONS,
+			// params needed for spicedb
+			"POOL_MODE=session",                         // https://github.com/authzed/spicedb/issues/1217
+			"IGNORE_STARTUP_PARAMETERS=plan_cache_mode", // Tell pgbouncer to pass this param thru to postgres.
+		},
+		ExposedPorts: []string{PGBOUNCER_TEST_PORT + "/tcp"},
+		NetworkID:    bridgeNetworkName,
+	})
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		require.NoError(t, pool.Purge(pgbouncer))
+	})
+
+	b.pgbouncerProxy = &container{
+		hostHostname:      "localhost",
+		hostPort:          pgbouncer.GetPort(PGBOUNCER_TEST_PORT + "/tcp"),
+		containerHostname: pgbouncerContainerHostname,
+		containerPort:     PGBOUNCER_TEST_PORT,
+	}
+}
+
+func (b *postgresTester) initializeHostConnection(t testing.TB, pool *dockertest.Pool) (conn *pgx.Conn) {
+	hostname, port := b.getHostHostnameAndPort()
+	uri := fmt.Sprintf("postgresql://%s@%s:%s/?sslmode=disable", b.creds, hostname, port)
+	err := pool.Retry(func() error {
+		var err error
+		ctx, cancelConnect := context.WithTimeout(context.Background(), dockerBootTimeout)
+		defer cancelConnect()
+		conn, err = pgx.Connect(ctx, uri)
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+	require.NoError(t, err)
+	return conn
+}
+
+func (b *postgresTester) getHostnameAndPort() (string, string) {
+	// If a bridgeNetworkName is supplied then we will return the container
+	// hostname and port that is resolvable from within the container network.
+	// If bridgeNetworkName is not supplied then the hostname and port will be
+	// resolvable from the host.
+	if b.useContainerHostname {
+		return b.getContainerHostnameAndPort()
+	}
+	return b.getHostHostnameAndPort()
+}
+
+func (b *postgresTester) getHostHostnameAndPort() (string, string) {
+	if b.pgbouncerProxy != nil {
+		return b.pgbouncerProxy.hostHostname, b.pgbouncerProxy.hostPort
+	}
+	return b.hostHostname, b.hostPort
+}
+
+func (b *postgresTester) getContainerHostnameAndPort() (string, string) {
+	if b.pgbouncerProxy != nil {
+		return b.pgbouncerProxy.containerHostname, b.pgbouncerProxy.containerPort
+	}
+	return b.containerHostname, b.containerPort
+}

magefiles/test.go

@@ -17,8 +17,8 @@ func (t Test) All() error {
 	ds := Testds{}
 	c := Testcons{}
 	mg.Deps(t.Unit, t.Integration, t.Image, t.Analyzers,
-		ds.Crdb, ds.Postgres, ds.Spanner, ds.Mysql,
-		c.Crdb, c.Spanner, c.Postgres, c.Mysql)
+		ds.Crdb, ds.Postgres, ds.Pgbouncer, ds.Spanner, ds.Mysql,
+		c.Crdb, c.Spanner, c.Postgres, c.Pgbouncer, c.Mysql)
 	return nil
 }
 
@@ -75,14 +75,20 @@ func (Testds) Postgres() error {
 	return datastoreTest("postgres")
 }
 
+// Pgbouncer runs datastore tests for postgres fronted with pgbouncer
+func (Testds) Pgbouncer() error {
+	return datastoreTest("postgres", "-enablePgbouncer")
+}
+
 // Mysql Run datastore tests for mysql
 func (Testds) Mysql() error {
 	return datastoreTest("mysql")
 }
 
-func datastoreTest(datastore string) error {
+func datastoreTest(datastore string, dsFlags ...string) error {
 	mg.Deps(checkDocker)
-	return goTest(fmt.Sprintf("./internal/datastore/%s/...", datastore), "-tags", "ci,docker", "-timeout", "10m")
+	args := append([]string{"-tags", "ci,docker", "-timeout", "10m"}, dsFlags...)
+	return goTest(fmt.Sprintf("./internal/datastore/%s/...", datastore), args...)
 }
 
 type Testcons mg.Namespace
@@ -102,15 +108,24 @@ func (Testcons) Postgres() error {
 	return consistencyTest("postgres")
 }
 
+// Pgbouncer runs consistency tests for postgres fronted with pgbouncer
+func (Testcons) Pgbouncer() error {
+	return consistencyTest("postgres", "-enablePgbouncer")
+}
+
 // Mysql Run consistency tests for mysql
 func (Testcons) Mysql() error {
 	return consistencyTest("mysql")
 }
 
-func consistencyTest(datastore string) error {
+func consistencyTest(datastore string, consFlags ...string) error {
 	mg.Deps(checkDocker)
-	return goTest("./internal/services/integrationtesting/...",
+	args := append([]string{
 		"-tags", "ci,docker,datastoreconsistency",
 		"-timeout", "10m",
-		"-run", fmt.Sprintf("TestConsistencyPerDatastore/%s", datastore))
+		"-run", fmt.Sprintf("TestConsistencyPerDatastore/%s", datastore),
+	},
+		consFlags...,
+	)
+	return goTest("./internal/services/integrationtesting/...", args...)
 }

magefiles/util.go

@@ -21,8 +21,8 @@ func goTest(path string, args ...string) error {
 
 // run go test in a directory
 func goDirTest(dir string, path string, args ...string) error {
-	testArgs := append([]string{"test", "-failfast", "-count=1"}, args...)
-	return RunSh(goCmdForTests(), WithV(), WithDir(dir), WithArgs(testArgs...))(path)
+	testArgs := append([]string{"test", path, "-failfast", "-count=1"}, args...)
+	return RunSh(goCmdForTests(), WithV(), WithDir(dir))(testArgs...)
 }
 
 // check if docker is installed and running