feat(postgres): use faster sql.DB instead of docker exec psql for snapshot/restore [rebased for main]

docs/modules/postgres.md

@@ -92,6 +92,15 @@ It's possible to use the Postgres container with PGVector, Timescale or Postgis,
 
 ## Examples
 
+### Wait Strategies
+
+The postgres module works best with these wait strategies.
+No default is supplied, so you need to set it explicitly.
+
+<!--codeinclude-->
+[Example Wait Strategies](../../modules/postgres/wait_strategies.go) inside_block:waitStrategy
+<!--/codeinclude-->
+
 ### Using Snapshots
 This example shows the usage of the postgres module's Snapshot feature to give each test a clean database without having
 to recreate the database container on every test or run heavy scripts to clean your database. This makes the individual
@@ -102,7 +111,35 @@ tests very modular, since they always run on a brand-new database.
     The Snapshot logic requires dropping the connected database and using the system database to run commands, which will
     not work if the database for the container is set to `"postgres"`.
 
-
 <!--codeinclude-->
 [Test with a reusable Postgres container](../../modules/postgres/postgres_test.go) inside_block:snapshotAndReset
 <!--/codeinclude-->
+
+### Snapshot/Restore with custom driver
+
+- Not available until the next release of testcontainers-go <a href="https://github.com/testcontainers/testcontainers-go"><span class="tc-version">:material-tag: main</span></a>
+
+The snapshot/restore feature tries to use the `postgres` driver with go's included `sql.DB` package to perform database operations.
+If the `postgres` driver is not installed, it will fall back to using `docker exec`, which works, but is slower.
+
+You can tell the module to use the database driver you have imported in your test package by setting `postgres.WithSQLDriver("name")` to your driver name.
+
+For example, if you use pgx, see the example below.
+
+```go
+package my_test
+
+import (
+	"testing"
+
+	_ "github.com/jackc/pgx/v5/stdlib"
+
+	"github.com/testcontainers/testcontainers-go/modules/postgres"
+)
+```
+
+The above code snippet is importing the `pgx` driver and the _Testcontainers for Go_ Postgres module.
+
+<!--codeinclude-->
+[Snapshot/Restore with custom driver](../../modules/postgres/postgres_test.go) inside_block:snapshotAndReset
+<!--/codeinclude-->

modules/postgres/go.mod

@@ -34,6 +34,7 @@ require (
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
 	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
+	github.com/jackc/puddle/v2 v2.2.1 // indirect
 	github.com/klauspost/compress v1.17.4 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/lufia/plan9stats v0.0.0-20211012122336-39d0f177ccd0 // indirect
@@ -60,6 +61,7 @@ require (
 	go.opentelemetry.io/otel/metric v1.24.0 // indirect
 	go.opentelemetry.io/otel/trace v1.24.0 // indirect
 	golang.org/x/crypto v0.22.0 // indirect
+	golang.org/x/sync v0.3.0 // indirect
 	golang.org/x/sys v0.19.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20231016165738-49dd2c1f3d0b // indirect

modules/postgres/options.go

@@ -0,0 +1,37 @@
+package postgres
+
+import (
+	"github.com/testcontainers/testcontainers-go"
+)
+
+type options struct {
+	// SQLDriverName is the name of the SQL driver to use.
+	SQLDriverName string
+}
+
+func defaultOptions() options {
+	return options{
+		SQLDriverName: "postgres",
+	}
+}
+
+// Compiler check to ensure that Option implements the testcontainers.ContainerCustomizer interface.
+var _ testcontainers.ContainerCustomizer = (Option)(nil)
+
+// Option is an option for the Redpanda container.
+type Option func(*options)
+
+// Customize is a NOOP. It's defined to satisfy the testcontainers.ContainerCustomizer interface.
+func (o Option) Customize(*testcontainers.GenericContainerRequest) error {
+	// NOOP to satisfy interface.
+	return nil
+}
+
+// WithSQLDriver sets the SQL driver to use for the container.
+// It is passed to sql.Open() to connect to the database when making or restoring snapshots.
+// This can be set if your app imports a different postgres driver, f.ex. "pgx"
+func WithSQLDriver(driver string) Option {
+	return func(o *options) {
+		o.SQLDriverName = driver
+	}
+}

modules/postgres/postgres.go

@@ -2,6 +2,7 @@ package postgres
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"io"
 	"net"
@@ -25,6 +26,9 @@ type PostgresContainer struct {
 	user         string
 	password     string
 	snapshotName string
+	// sqlDriverName is passed to sql.Open() to connect to the database when making or restoring snapshots.
+	// This can be set if your app imports a different postgres driver, f.ex. "pgx"
+	sqlDriverName string
 }
 
 // MustConnectionString panics if the address cannot be determined.
@@ -148,7 +152,12 @@ func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomize
 		Started:          true,
 	}
 
+	// Gather all config options (defaults and then apply provided options)
+	settings := defaultOptions()
 	for _, opt := range opts {
+		if apply, ok := opt.(Option); ok {
+			apply(&settings)
+		}
 		if err := opt.Customize(&genericContainerReq); err != nil {
 			return nil, err
 		}
@@ -163,7 +172,7 @@ func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomize
 	password := req.Env["POSTGRES_PASSWORD"]
 	dbName := req.Env["POSTGRES_DB"]
 
-	return &PostgresContainer{Container: container, dbName: dbName, password: password, user: user}, nil
+	return &PostgresContainer{Container: container, dbName: dbName, password: password, user: user, sqlDriverName: settings.SQLDriverName}, nil
 }
 
 type snapshotConfig struct {
@@ -187,54 +196,45 @@ func WithSnapshotName(name string) SnapshotOption {
 // customize the snapshot name with the options.
 // If a snapshot already exists under the given/default name, it will be overwritten with the new snapshot.
 func (c *PostgresContainer) Snapshot(ctx context.Context, opts ...SnapshotOption) error {
-	config := &snapshotConfig{}
-	for _, opt := range opts {
-		config = opt(config)
-	}
-
-	snapshotName := defaultSnapshotName
-	if config.snapshotName != "" {
-		snapshotName = config.snapshotName
-	}
-
-	if c.dbName == "postgres" {
-		return fmt.Errorf("cannot snapshot the postgres system database as it cannot be dropped to be restored")
+	snapshotName, err := c.checkSnapshotConfig(opts)
+	if err != nil {
+		return err
 	}
 
 	// execute the commands to create the snapshot, in order
-	cmds := []string{
+	if err := c.execCommandsSQL(ctx,
 		// Drop the snapshot database if it already exists
 		fmt.Sprintf(`DROP DATABASE IF EXISTS "%s"`, snapshotName),
 		// Create a copy of the database to another database to use as a template now that it was fully migrated
 		fmt.Sprintf(`CREATE DATABASE "%s" WITH TEMPLATE "%s" OWNER "%s"`, snapshotName, c.dbName, c.user),
 		// Snapshot the template database so we can restore it onto our original database going forward
 		fmt.Sprintf(`ALTER DATABASE "%s" WITH is_template = TRUE`, snapshotName),
-	}
-
-	for _, cmd := range cmds {
-		exitCode, reader, err := c.Exec(ctx, []string{"psql", "-U", c.user, "-d", c.dbName, "-c", cmd})
-		if err != nil {
-			return err
-		}
-		if exitCode != 0 {
-			buf := new(strings.Builder)
-			_, err := io.Copy(buf, reader)
-			if err != nil {
-				return fmt.Errorf("non-zero exit code for snapshot command, could not read command output: %w", err)
-			}
-
-			return fmt.Errorf("non-zero exit code for snapshot command: %s", buf.String())
-		}
+	); err != nil {
+		return err
 	}
 
 	c.snapshotName = snapshotName
-
 	return nil
 }
 
 // Restore will restore the database to a specific snapshot. By default, it will restore the last snapshot taken on the
 // database by the Snapshot method. If a snapshot name is provided, it will instead try to restore the snapshot by name.
 func (c *PostgresContainer) Restore(ctx context.Context, opts ...SnapshotOption) error {
+	snapshotName, err := c.checkSnapshotConfig(opts)
+	if err != nil {
+		return err
+	}
+
+	// execute the commands to restore the snapshot, in order
+	return c.execCommandsSQL(ctx,
+		// Drop the entire database by connecting to the postgres global database
+		fmt.Sprintf(`DROP DATABASE "%s" with (FORCE)`, c.dbName),
+		// Then restore the previous snapshot
+		fmt.Sprintf(`CREATE DATABASE "%s" WITH TEMPLATE "%s" OWNER "%s"`, c.dbName, snapshotName, c.user),
+	)
+}
+
+func (c *PostgresContainer) checkSnapshotConfig(opts []SnapshotOption) (string, error) {
 	config := &snapshotConfig{}
 	for _, opt := range opts {
 		config = opt(config)
@@ -246,17 +246,64 @@ func (c *PostgresContainer) Restore(ctx context.Context, opts ...SnapshotOption)
 	}
 
 	if c.dbName == "postgres" {
-		return fmt.Errorf("cannot restore the postgres system database as it cannot be dropped to be restored")
+		return "", fmt.Errorf("cannot restore the postgres system database as it cannot be dropped to be restored")
 	}
+	return snapshotName, nil
+}
 
-	// execute the commands to restore the snapshot, in order
-	cmds := []string{
-		// Drop the entire database by connecting to the postgres global database
-		fmt.Sprintf(`DROP DATABASE "%s" with (FORCE)`, c.dbName),
-		// Then restore the previous snapshot
-		fmt.Sprintf(`CREATE DATABASE "%s" WITH TEMPLATE "%s" OWNER "%s"`, c.dbName, snapshotName, c.user),
+func (c *PostgresContainer) execCommandsSQL(ctx context.Context, cmds ...string) error {
+	conn, cleanup, err := c.snapshotConnection(ctx)
+	if err != nil {
+		testcontainers.Logger.Printf("Could not connect to database to restore snapshot, falling back to `docker exec psql`: %v", err)
+		return c.execCommandsFallback(ctx, cmds)
+	}
+	if cleanup != nil {
+		defer cleanup()
+	}
+	for _, cmd := range cmds {
+		if _, err := conn.ExecContext(ctx, cmd); err != nil {
+			return fmt.Errorf("could not execute restore command %s: %w", cmd, err)
+		}
+	}
+	return nil
+}
+
+// snapshotConnection connects to the actual database using the "postgres" sql.DB driver, if it exists.
+// The returned function should be called as a defer() to close the pool.
+// No need to close the individual connection, that is done as part of the pool close.
+// Also, no need to cache the connection pool, since it is a single connection which is very fast to establish.
+func (c *PostgresContainer) snapshotConnection(ctx context.Context) (*sql.Conn, func(), error) {
+	// Connect to the database "postgres" instead of the app one
+	c2 := &PostgresContainer{
+		Container:     c.Container,
+		dbName:        "postgres",
+		user:          c.user,
+		password:      c.password,
+		sqlDriverName: c.sqlDriverName,
 	}
 
+	// Try to use an actual postgres connection, if the driver is loaded
+	connStr := c2.MustConnectionString(ctx, "sslmode=disable")
+	pool, err := sql.Open(c.sqlDriverName, connStr)
+	if err != nil {
+		return nil, nil, fmt.Errorf("sql.Open for snapshot connection failed: %w", err)
+	}
+
+	cleanupPool := func() {
+		if err := pool.Close(); err != nil {
+			testcontainers.Logger.Printf("Could not close database connection pool after restoring snapshot: %v", err)
+		}
+	}
+
+	conn, err := pool.Conn(ctx)
+	if err != nil {
+		cleanupPool()
+		return nil, nil, fmt.Errorf("DB.Conn for snapshot connection failed: %w", err)
+	}
+	return conn, cleanupPool, nil
+}
+
+func (c *PostgresContainer) execCommandsFallback(ctx context.Context, cmds []string) error {
 	for _, cmd := range cmds {
 		exitCode, reader, err := c.Exec(ctx, []string{"psql", "-v", "ON_ERROR_STOP=1", "-U", c.user, "-d", "postgres", "-c", cmd})
 		if err != nil {
@@ -272,6 +319,5 @@ func (c *PostgresContainer) Restore(ctx context.Context, opts ...SnapshotOption)
 			return fmt.Errorf("non-zero exit code for restore command: %s", buf.String())
 		}
 	}
-
 	return nil
 }