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/postgres_test.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,37 @@ 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.SQLDriverName` 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"
+)
+
+func init() {
+	postgres.SQLDriverName = "pgx"
+}
+
+func TestSomething(t *testing.T) {
+    // Your test code here, using postgres.RunContainer()
+}
+```

modules/postgres/postgres.go

@@ -2,6 +2,7 @@ package postgres
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"io"
 	"net"
@@ -18,6 +19,10 @@ const (
 	defaultSnapshotName  = "migrated_template"
 )
 
+// 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"
+var SQLDriverName = "postgres"
+
 // PostgresContainer represents the postgres container type used in the module
 type PostgresContainer struct {
 	testcontainers.Container
@@ -187,54 +192,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 +242,63 @@ 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,
 	}
 
+	// Try to use an actual postgres connection, if the driver is loaded
+	connStr := c2.MustConnectionString(ctx, "sslmode=disable")
+	pool, err := sql.Open(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 +314,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
 }