From 51f83fffd9dad86c8e49cf13f047dacc0feaba66 Mon Sep 17 00:00:00 2001
From: Claus Strasburger <1156864+cfstras@users.noreply.github.com>
Date: Thu, 20 Jun 2024 14:37:04 +0200
Subject: [PATCH 01/12] feat(postgres): use faster sql.DB instead of `docker
 exec psql` for creating & restoring snapshots

---
 modules/postgres/postgres.go | 111 +++++++++++++++++++++++------------
 1 file changed, 73 insertions(+), 38 deletions(-)

diff --git a/modules/postgres/postgres.go b/modules/postgres/postgres.go
index f8a3a1c630..43d86cbe19 100644
--- a/modules/postgres/postgres.go
+++ b/modules/postgres/postgres.go
@@ -2,6 +2,7 @@ package postgres
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"io"
 	"net"
@@ -187,54 +188,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 +238,61 @@ 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)
+	}
+	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("postgres", 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 +308,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
 }

From 6b54ac0dc738029e2c12c3f858b5c8701980d90a Mon Sep 17 00:00:00 2001
From: Claus Strasburger <1156864+cfstras@users.noreply.github.com>
Date: Thu, 20 Jun 2024 15:13:37 +0200
Subject: [PATCH 02/12] test(postgres): Make tests work reliably on slower
 computers

---
 modules/postgres/postgres_test.go | 35 +++++++++++++++++--------------
 1 file changed, 19 insertions(+), 16 deletions(-)

diff --git a/modules/postgres/postgres_test.go b/modules/postgres/postgres_test.go
index 9dc2194685..54a1f59b48 100644
--- a/modules/postgres/postgres_test.go
+++ b/modules/postgres/postgres_test.go
@@ -26,38 +26,47 @@ const (
 	password = "password"
 )
 
+var (
+	// BasicWaitStrategies is a simple but reliable way to wait for postgres to start.
+	BasicWaitStrategies = testcontainers.WithWaitStrategy(
+		// First, we wait for the container to log readiness twice.
+		// This is because it will restart itself after the first startup.
+		wait.ForLog("database system is ready to accept connections").WithOccurrence(2),
+		// Then, we wait for docker to actually serve the port on localhost.
+		// For non-linux OSes like Mac and Windows, Docker or Rancher Desktop will have to start a separate proxy
+		// process that sometimes takes a little longer to initialize.
+		// Without this, the tests will be flaky on those OSes!
+		wait.ForListeningPort("5432/tcp"),
+	)
+)
+
 func TestPostgres(t *testing.T) {
 	ctx := context.Background()
 
 	tests := []struct {
 		name  string
 		image string
-		wait  wait.Strategy
 	}{
 		{
 			name:  "Postgres",
 			image: "docker.io/postgres:15.2-alpine",
-			wait:  wait.ForLog("database system is ready to accept connections").WithOccurrence(2).WithStartupTimeout(5 * time.Second),
 		},
 		{
 			name: "Timescale",
 			// timescale {
 			image: "docker.io/timescale/timescaledb:2.1.0-pg11",
-			wait:  wait.ForLog("database system is ready to accept connections").WithOccurrence(2).WithStartupTimeout(5 * time.Second),
 			// }
 		},
 		{
 			name: "Postgis",
 			// postgis {
 			image: "docker.io/postgis/postgis:12-3.0",
-			wait:  wait.ForLog("database system is ready to accept connections").WithOccurrence(2).WithStartupTimeout(30 * time.Second),
 			// }
 		},
 		{
 			name: "Pgvector",
 			// pgvector {
 			image: "docker.io/pgvector/pgvector:pg16",
-			wait:  wait.ForLog("database system is ready to accept connections").WithOccurrence(2).WithStartupTimeout(30 * time.Second),
 			// }
 		},
 	}
@@ -69,7 +78,7 @@ func TestPostgres(t *testing.T) {
 				postgres.WithDatabase(dbname),
 				postgres.WithUsername(user),
 				postgres.WithPassword(password),
-				testcontainers.WithWaitStrategy(tt.wait),
+				BasicWaitStrategies,
 			)
 			if err != nil {
 				t.Fatal(err)
@@ -166,7 +175,7 @@ func TestWithConfigFile(t *testing.T) {
 		postgres.WithDatabase(dbname),
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
-		testcontainers.WithWaitStrategy(wait.ForLog("database system is ready to accept connections").WithOccurrence(2).WithStartupTimeout(5*time.Second)),
+		BasicWaitStrategies,
 	)
 	if err != nil {
 		t.Fatal(err)
@@ -197,7 +206,7 @@ func TestWithInitScript(t *testing.T) {
 		postgres.WithDatabase(dbname),
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
-		testcontainers.WithWaitStrategy(wait.ForLog("database system is ready to accept connections").WithOccurrence(2).WithStartupTimeout(5*time.Second)),
+		BasicWaitStrategies,
 	)
 	if err != nil {
 		t.Fatal(err)
@@ -235,10 +244,7 @@ func TestSnapshot(t *testing.T) {
 		postgres.WithDatabase(dbname),
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
-		testcontainers.WithWaitStrategy(
-			wait.ForLog("database system is ready to accept connections").
-				WithOccurrence(2).
-				WithStartupTimeout(5*time.Second)),
+		BasicWaitStrategies,
 	)
 	if err != nil {
 		t.Fatal(err)
@@ -341,10 +347,7 @@ func TestSnapshotWithOverrides(t *testing.T) {
 		postgres.WithDatabase(dbname),
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
-		testcontainers.WithWaitStrategy(
-			wait.ForLog("database system is ready to accept connections").
-				WithOccurrence(2).
-				WithStartupTimeout(5*time.Second)),
+		BasicWaitStrategies,
 	)
 	if err != nil {
 		t.Fatal(err)

From c827da4e898f74dfb6b8e5aa2ac6b683d337b338 Mon Sep 17 00:00:00 2001
From: Claus Strasburger <1156864+cfstras@users.noreply.github.com>
Date: Thu, 20 Jun 2024 15:15:03 +0200
Subject: [PATCH 03/12] chore(postgres): simplify testdata

Remove all commented lines
---
 modules/postgres/testdata/my-postgres.conf | 690 ---------------------
 1 file changed, 690 deletions(-)

diff --git a/modules/postgres/testdata/my-postgres.conf b/modules/postgres/testdata/my-postgres.conf
index 27dd8b527e..128ef1aa8c 100644
--- a/modules/postgres/testdata/my-postgres.conf
+++ b/modules/postgres/testdata/my-postgres.conf
@@ -1,691 +1 @@
-# -----------------------------
-# PostgreSQL configuration file
-# -----------------------------
-#
-# This file consists of lines of the form:
-#
-#   name = value
-#
-# (The "=" is optional.)  Whitespace may be used.  Comments are introduced with
-# "#" anywhere on a line.  The complete list of parameter names and allowed
-# values can be found in the PostgreSQL documentation.
-#
-# The commented-out settings shown in this file represent the default values.
-# Re-commenting a setting is NOT sufficient to revert it to the default value;
-# you need to reload the server.
-#
-# This file is read on server startup and when the server receives a SIGHUP
-# signal.  If you edit the file on a running system, you have to SIGHUP the
-# server for the changes to take effect, run "pg_ctl reload", or execute
-# "SELECT pg_reload_conf()".  Some parameters, which are marked below,
-# require a server shutdown and restart to take effect.
-#
-# Any parameter can also be given as a command-line option to the server, e.g.,
-# "postgres -c log_connections=on".  Some parameters can be changed at run time
-# with the "SET" SQL command.
-#
-# Memory units:   B  = bytes           Time units:  ms  = milliseconds
-#                kB = kilobytes                     s   = seconds
-#                MB = megabytes                     min = minutes
-#                GB = gigabytes                     h   = hours
-#                TB = terabytes                     d   = days
-
-
-#------------------------------------------------------------------------------
-# FILE LOCATIONS
-#------------------------------------------------------------------------------
-
-# The default values of these variables are driven from the -D command-line
-# option or PGDATA environment variable, represented here as ConfigDir.
-
-#data_directory = 'ConfigDir'		# use data in another directory
-					# (change requires restart)
-#hba_file = 'ConfigDir/pg_hba.conf'	# host-based authentication file
-					# (change requires restart)
-#ident_file = 'ConfigDir/pg_ident.conf'	# ident configuration file
-					# (change requires restart)
-
-# If external_pid_file is not explicitly set, no extra PID file is written.
-#external_pid_file = ''			# write an extra PID file
-					# (change requires restart)
-
-
-#------------------------------------------------------------------------------
-# CONNECTIONS AND AUTHENTICATION
-#------------------------------------------------------------------------------
-
-# - Connection Settings -
-
 listen_addresses = '*'
-					# comma-separated list of addresses;
-					# defaults to 'localhost'; use '*' for all
-					# (change requires restart)
-#port = 5432				# (change requires restart)
-#max_connections = 100			# (change requires restart)
-#superuser_reserved_connections = 3	# (change requires restart)
-#unix_socket_directories = '/tmp'	# comma-separated list of directories
-					# (change requires restart)
-#unix_socket_group = ''			# (change requires restart)
-#unix_socket_permissions = 0777		# begin with 0 to use octal notation
-					# (change requires restart)
-#bonjour = off				# advertise server via Bonjour
-					# (change requires restart)
-#bonjour_name = ''			# defaults to the computer name
-					# (change requires restart)
-
-# - TCP Keepalives -
-# see "man 7 tcp" for details
-
-#tcp_keepalives_idle = 0		# TCP_KEEPIDLE, in seconds;
-					# 0 selects the system default
-#tcp_keepalives_interval = 0		# TCP_KEEPINTVL, in seconds;
-					# 0 selects the system default
-#tcp_keepalives_count = 0		# TCP_KEEPCNT;
-					# 0 selects the system default
-
-# - Authentication -
-
-#authentication_timeout = 1min		# 1s-600s
-#password_encryption = md5		# md5 or scram-sha-256
-#db_user_namespace = off
-
-# GSSAPI using Kerberos
-#krb_server_keyfile = ''
-#krb_caseins_users = off
-
-# - SSL -
-
-#ssl = off
-#ssl_ca_file = ''
-#ssl_cert_file = 'server.crt'
-#ssl_crl_file = ''
-#ssl_key_file = 'server.key'
-#ssl_ciphers = 'HIGH:MEDIUM:+3DES:!aNULL' # allowed SSL ciphers
-#ssl_prefer_server_ciphers = on
-#ssl_ecdh_curve = 'prime256v1'
-#ssl_dh_params_file = ''
-#ssl_passphrase_command = ''
-#ssl_passphrase_command_supports_reload = off
-
-
-#------------------------------------------------------------------------------
-# RESOURCE USAGE (except WAL)
-#------------------------------------------------------------------------------
-
-# - Memory -
-
-#shared_buffers = 32MB			# min 128kB
-					# (change requires restart)
-#huge_pages = try			# on, off, or try
-					# (change requires restart)
-#temp_buffers = 8MB			# min 800kB
-#max_prepared_transactions = 0		# zero disables the feature
-					# (change requires restart)
-# Caution: it is not advisable to set max_prepared_transactions nonzero unless
-# you actively intend to use prepared transactions.
-#work_mem = 4MB				# min 64kB
-#maintenance_work_mem = 64MB		# min 1MB
-#autovacuum_work_mem = -1		# min 1MB, or -1 to use maintenance_work_mem
-#max_stack_depth = 2MB			# min 100kB
-#dynamic_shared_memory_type = posix	# the default is the first option
-					# supported by the operating system:
-					#   posix
-					#   sysv
-					#   windows
-					#   mmap
-					# use none to disable dynamic shared memory
-					# (change requires restart)
-
-# - Disk -
-
-#temp_file_limit = -1			# limits per-process temp file space
-					# in kB, or -1 for no limit
-
-# - Kernel Resources -
-
-#max_files_per_process = 1000		# min 25
-					# (change requires restart)
-
-# - Cost-Based Vacuum Delay -
-
-#vacuum_cost_delay = 0			# 0-100 milliseconds
-#vacuum_cost_page_hit = 1		# 0-10000 credits
-#vacuum_cost_page_miss = 10		# 0-10000 credits
-#vacuum_cost_page_dirty = 20		# 0-10000 credits
-#vacuum_cost_limit = 200		# 1-10000 credits
-
-# - Background Writer -
-
-#bgwriter_delay = 200ms			# 10-10000ms between rounds
-#bgwriter_lru_maxpages = 100		# max buffers written/round, 0 disables
-#bgwriter_lru_multiplier = 2.0		# 0-10.0 multiplier on buffers scanned/round
-#bgwriter_flush_after = 0		# measured in pages, 0 disables
-
-# - Asynchronous Behavior -
-
-#effective_io_concurrency = 1		# 1-1000; 0 disables prefetching
-#max_worker_processes = 8		# (change requires restart)
-#max_parallel_maintenance_workers = 2	# taken from max_parallel_workers
-#max_parallel_workers_per_gather = 2	# taken from max_parallel_workers
-#parallel_leader_participation = on
-#max_parallel_workers = 8		# maximum number of max_worker_processes that
-					# can be used in parallel operations
-#old_snapshot_threshold = -1		# 1min-60d; -1 disables; 0 is immediate
-					# (change requires restart)
-#backend_flush_after = 0		# measured in pages, 0 disables
-
-
-#------------------------------------------------------------------------------
-# WRITE-AHEAD LOG
-#------------------------------------------------------------------------------
-
-# - Settings -
-
-#wal_level = replica			# minimal, replica, or logical
-					# (change requires restart)
-#fsync = on				# flush data to disk for crash safety
-					# (turning this off can cause
-					# unrecoverable data corruption)
-#synchronous_commit = on		# synchronization level;
-					# off, local, remote_write, remote_apply, or on
-#wal_sync_method = fsync		# the default is the first option
-					# supported by the operating system:
-					#   open_datasync
-					#   fdatasync (default on Linux and FreeBSD)
-					#   fsync
-					#   fsync_writethrough
-					#   open_sync
-#full_page_writes = on			# recover from partial page writes
-#wal_compression = off			# enable compression of full-page writes
-#wal_log_hints = off			# also do full page writes of non-critical updates
-					# (change requires restart)
-#wal_buffers = -1			# min 32kB, -1 sets based on shared_buffers
-					# (change requires restart)
-#wal_writer_delay = 200ms		# 1-10000 milliseconds
-#wal_writer_flush_after = 1MB		# measured in pages, 0 disables
-
-#commit_delay = 0			# range 0-100000, in microseconds
-#commit_siblings = 5			# range 1-1000
-
-# - Checkpoints -
-
-#checkpoint_timeout = 5min		# range 30s-1d
-#max_wal_size = 1GB
-#min_wal_size = 80MB
-#checkpoint_completion_target = 0.5	# checkpoint target duration, 0.0 - 1.0
-#checkpoint_flush_after = 0		# measured in pages, 0 disables
-#checkpoint_warning = 30s		# 0 disables
-
-# - Archiving -
-
-#archive_mode = off		# enables archiving; off, on, or always
-				# (change requires restart)
-#archive_command = ''		# command to use to archive a logfile segment
-				# placeholders: %p = path of file to archive
-				#               %f = file name only
-				# e.g. 'test ! -f /mnt/server/archivedir/%f && cp %p /mnt/server/archivedir/%f'
-#archive_timeout = 0		# force a logfile segment switch after this
-				# number of seconds; 0 disables
-
-
-#------------------------------------------------------------------------------
-# REPLICATION
-#------------------------------------------------------------------------------
-
-# - Sending Servers -
-
-# Set these on the master and on any standby that will send replication data.
-
-#max_wal_senders = 10		# max number of walsender processes
-				# (change requires restart)
-#wal_keep_segments = 0		# in logfile segments; 0 disables
-#wal_sender_timeout = 60s	# in milliseconds; 0 disables
-
-#max_replication_slots = 10	# max number of replication slots
-				# (change requires restart)
-#track_commit_timestamp = off	# collect timestamp of transaction commit
-				# (change requires restart)
-
-# - Master Server -
-
-# These settings are ignored on a standby server.
-
-#synchronous_standby_names = ''	# standby servers that provide sync rep
-				# method to choose sync standbys, number of sync standbys,
-				# and comma-separated list of application_name
-				# from standby(s); '*' = all
-#vacuum_defer_cleanup_age = 0	# number of xacts by which cleanup is delayed
-
-# - Standby Servers -
-
-# These settings are ignored on a master server.
-
-#hot_standby = on			# "off" disallows queries during recovery
-					# (change requires restart)
-#max_standby_archive_delay = 30s	# max delay before canceling queries
-					# when reading WAL from archive;
-					# -1 allows indefinite delay
-#max_standby_streaming_delay = 30s	# max delay before canceling queries
-					# when reading streaming WAL;
-					# -1 allows indefinite delay
-#wal_receiver_status_interval = 10s	# send replies at least this often
-					# 0 disables
-#hot_standby_feedback = off		# send info from standby to prevent
-					# query conflicts
-#wal_receiver_timeout = 60s		# time that receiver waits for
-					# communication from master
-					# in milliseconds; 0 disables
-#wal_retrieve_retry_interval = 5s	# time to wait before retrying to
-					# retrieve WAL after a failed attempt
-
-# - Subscribers -
-
-# These settings are ignored on a publisher.
-
-#max_logical_replication_workers = 4	# taken from max_worker_processes
-					# (change requires restart)
-#max_sync_workers_per_subscription = 2	# taken from max_logical_replication_workers
-
-
-#------------------------------------------------------------------------------
-# QUERY TUNING
-#------------------------------------------------------------------------------
-
-# - Planner Method Configuration -
-
-#enable_bitmapscan = on
-#enable_hashagg = on
-#enable_hashjoin = on
-#enable_indexscan = on
-#enable_indexonlyscan = on
-#enable_material = on
-#enable_mergejoin = on
-#enable_nestloop = on
-#enable_parallel_append = on
-#enable_seqscan = on
-#enable_sort = on
-#enable_tidscan = on
-#enable_partitionwise_join = off
-#enable_partitionwise_aggregate = off
-#enable_parallel_hash = on
-#enable_partition_pruning = on
-
-# - Planner Cost Constants -
-
-#seq_page_cost = 1.0			# measured on an arbitrary scale
-#random_page_cost = 4.0			# same scale as above
-#cpu_tuple_cost = 0.01			# same scale as above
-#cpu_index_tuple_cost = 0.005		# same scale as above
-#cpu_operator_cost = 0.0025		# same scale as above
-#parallel_tuple_cost = 0.1		# same scale as above
-#parallel_setup_cost = 1000.0	# same scale as above
-
-#jit_above_cost = 100000		# perform JIT compilation if available
-					# and query more expensive than this;
-					# -1 disables
-#jit_inline_above_cost = 500000		# inline small functions if query is
-					# more expensive than this; -1 disables
-#jit_optimize_above_cost = 500000	# use expensive JIT optimizations if
-					# query is more expensive than this;
-					# -1 disables
-
-#min_parallel_table_scan_size = 8MB
-#min_parallel_index_scan_size = 512kB
-#effective_cache_size = 4GB
-
-# - Genetic Query Optimizer -
-
-#geqo = on
-#geqo_threshold = 12
-#geqo_effort = 5			# range 1-10
-#geqo_pool_size = 0			# selects default based on effort
-#geqo_generations = 0			# selects default based on effort
-#geqo_selection_bias = 2.0		# range 1.5-2.0
-#geqo_seed = 0.0			# range 0.0-1.0
-
-# - Other Planner Options -
-
-#default_statistics_target = 100	# range 1-10000
-#constraint_exclusion = partition	# on, off, or partition
-#cursor_tuple_fraction = 0.1		# range 0.0-1.0
-#from_collapse_limit = 8
-#join_collapse_limit = 8		# 1 disables collapsing of explicit
-					# JOIN clauses
-#force_parallel_mode = off
-#jit = off				# allow JIT compilation
-
-
-#------------------------------------------------------------------------------
-# REPORTING AND LOGGING
-#------------------------------------------------------------------------------
-
-# - Where to Log -
-
-#log_destination = 'stderr'		# Valid values are combinations of
-					# stderr, csvlog, syslog, and eventlog,
-					# depending on platform.  csvlog
-					# requires logging_collector to be on.
-
-# This is used when logging to stderr:
-#logging_collector = off		# Enable capturing of stderr and csvlog
-					# into log files. Required to be on for
-					# csvlogs.
-					# (change requires restart)
-
-# These are only used if logging_collector is on:
-#log_directory = 'log'			# directory where log files are written,
-					# can be absolute or relative to PGDATA
-#log_filename = 'postgresql-%Y-%m-%d_%H%M%S.log'	# log file name pattern,
-					# can include strftime() escapes
-#log_file_mode = 0600			# creation mode for log files,
-					# begin with 0 to use octal notation
-#log_truncate_on_rotation = off		# If on, an existing log file with the
-					# same name as the new log file will be
-					# truncated rather than appended to.
-					# But such truncation only occurs on
-					# time-driven rotation, not on restarts
-					# or size-driven rotation.  Default is
-					# off, meaning append to existing files
-					# in all cases.
-#log_rotation_age = 1d			# Automatic rotation of logfiles will
-					# happen after that time.  0 disables.
-#log_rotation_size = 10MB		# Automatic rotation of logfiles will
-					# happen after that much log output.
-					# 0 disables.
-
-# These are relevant when logging to syslog:
-#syslog_facility = 'LOCAL0'
-#syslog_ident = 'postgres'
-#syslog_sequence_numbers = on
-#syslog_split_messages = on
-
-# This is only relevant when logging to eventlog (win32):
-# (change requires restart)
-#event_source = 'PostgreSQL'
-
-# - When to Log -
-
-#log_min_messages = warning		# values in order of decreasing detail:
-					#   debug5
-					#   debug4
-					#   debug3
-					#   debug2
-					#   debug1
-					#   info
-					#   notice
-					#   warning
-					#   error
-					#   log
-					#   fatal
-					#   panic
-
-#log_min_error_statement = error	# values in order of decreasing detail:
-					#   debug5
-					#   debug4
-					#   debug3
-					#   debug2
-					#   debug1
-					#   info
-					#   notice
-					#   warning
-					#   error
-					#   log
-					#   fatal
-					#   panic (effectively off)
-
-#log_min_duration_statement = -1	# -1 is disabled, 0 logs all statements
-					# and their durations, > 0 logs only
-					# statements running at least this number
-					# of milliseconds
-
-
-# - What to Log -
-
-#debug_print_parse = off
-#debug_print_rewritten = off
-#debug_print_plan = off
-#debug_pretty_print = on
-#log_checkpoints = off
-#log_connections = off
-#log_disconnections = off
-#log_duration = off
-#log_error_verbosity = default		# terse, default, or verbose messages
-#log_hostname = off
-#log_line_prefix = '%m [%p] '		# special values:
-					#   %a = application name
-					#   %u = user name
-					#   %d = database name
-					#   %r = remote host and port
-					#   %h = remote host
-					#   %p = process ID
-					#   %t = timestamp without milliseconds
-					#   %m = timestamp with milliseconds
-					#   %n = timestamp with milliseconds (as a Unix epoch)
-					#   %i = command tag
-					#   %e = SQL state
-					#   %c = session ID
-					#   %l = session line number
-					#   %s = session start timestamp
-					#   %v = virtual transaction ID
-					#   %x = transaction ID (0 if none)
-					#   %q = stop here in non-session
-					#        processes
-					#   %% = '%'
-					# e.g. '<%u%%%d> '
-#log_lock_waits = off			# log lock waits >= deadlock_timeout
-#log_statement = 'none'			# none, ddl, mod, all
-#log_replication_commands = off
-#log_temp_files = -1			# log temporary files equal or larger
-					# than the specified size in kilobytes;
-					# -1 disables, 0 logs all temp files
-#log_timezone = 'GMT'
-
-#------------------------------------------------------------------------------
-# PROCESS TITLE
-#------------------------------------------------------------------------------
-
-#cluster_name = ''			# added to process titles if nonempty
-					# (change requires restart)
-#update_process_title = on
-
-
-#------------------------------------------------------------------------------
-# STATISTICS
-#------------------------------------------------------------------------------
-
-# - Query and Index Statistics Collector -
-
-#track_activities = on
-#track_counts = on
-#track_io_timing = off
-#track_functions = none			# none, pl, all
-#track_activity_query_size = 1024	# (change requires restart)
-#stats_temp_directory = 'pg_stat_tmp'
-
-
-# - Monitoring -
-
-#log_parser_stats = off
-#log_planner_stats = off
-#log_executor_stats = off
-#log_statement_stats = off
-
-
-#------------------------------------------------------------------------------
-# AUTOVACUUM
-#------------------------------------------------------------------------------
-
-#autovacuum = on			# Enable autovacuum subprocess?  'on'
-					# requires track_counts to also be on.
-#log_autovacuum_min_duration = -1	# -1 disables, 0 logs all actions and
-					# their durations, > 0 logs only
-					# actions running at least this number
-					# of milliseconds.
-#autovacuum_max_workers = 3		# max number of autovacuum subprocesses
-					# (change requires restart)
-#autovacuum_naptime = 1min		# time between autovacuum runs
-#autovacuum_vacuum_threshold = 50	# min number of row updates before
-					# vacuum
-#autovacuum_analyze_threshold = 50	# min number of row updates before
-					# analyze
-#autovacuum_vacuum_scale_factor = 0.2	# fraction of table size before vacuum
-#autovacuum_analyze_scale_factor = 0.1	# fraction of table size before analyze
-#autovacuum_freeze_max_age = 200000000	# maximum XID age before forced vacuum
-					# (change requires restart)
-#autovacuum_multixact_freeze_max_age = 400000000	# maximum multixact age
-					# before forced vacuum
-					# (change requires restart)
-#autovacuum_vacuum_cost_delay = 20ms	# default vacuum cost delay for
-					# autovacuum, in milliseconds;
-					# -1 means use vacuum_cost_delay
-#autovacuum_vacuum_cost_limit = -1	# default vacuum cost limit for
-					# autovacuum, -1 means use
-					# vacuum_cost_limit
-
-
-#------------------------------------------------------------------------------
-# CLIENT CONNECTION DEFAULTS
-#------------------------------------------------------------------------------
-
-# - Statement Behavior -
-
-#client_min_messages = notice		# values in order of decreasing detail:
-					#   debug5
-					#   debug4
-					#   debug3
-					#   debug2
-					#   debug1
-					#   log
-					#   notice
-					#   warning
-					#   error
-#search_path = '"$user", public'	# schema names
-#row_security = on
-#default_tablespace = ''		# a tablespace name, '' uses the default
-#temp_tablespaces = ''			# a list of tablespace names, '' uses
-					# only default tablespace
-#check_function_bodies = on
-#default_transaction_isolation = 'read committed'
-#default_transaction_read_only = off
-#default_transaction_deferrable = off
-#session_replication_role = 'origin'
-#statement_timeout = 0			# in milliseconds, 0 is disabled
-#lock_timeout = 0			# in milliseconds, 0 is disabled
-#idle_in_transaction_session_timeout = 0	# in milliseconds, 0 is disabled
-#vacuum_freeze_min_age = 50000000
-#vacuum_freeze_table_age = 150000000
-#vacuum_multixact_freeze_min_age = 5000000
-#vacuum_multixact_freeze_table_age = 150000000
-#vacuum_cleanup_index_scale_factor = 0.1	# fraction of total number of tuples
-						# before index cleanup, 0 always performs
-						# index cleanup
-#bytea_output = 'hex'			# hex, escape
-#xmlbinary = 'base64'
-#xmloption = 'content'
-#gin_fuzzy_search_limit = 0
-#gin_pending_list_limit = 4MB
-
-# - Locale and Formatting -
-
-#datestyle = 'iso, mdy'
-#intervalstyle = 'postgres'
-#timezone = 'GMT'
-#timezone_abbreviations = 'Default'     # Select the set of available time zone
-					# abbreviations.  Currently, there are
-					#   Default
-					#   Australia (historical usage)
-					#   India
-					# You can create your own file in
-					# share/timezonesets/.
-#extra_float_digits = 0			# min -15, max 3
-#client_encoding = sql_ascii		# actually, defaults to database
-					# encoding
-
-# These settings are initialized by initdb, but they can be changed.
-#lc_messages = 'C'			# locale for system error message
-					# strings
-#lc_monetary = 'C'			# locale for monetary formatting
-#lc_numeric = 'C'			# locale for number formatting
-#lc_time = 'C'				# locale for time formatting
-
-# default configuration for text search
-#default_text_search_config = 'pg_catalog.simple'
-
-# - Shared Library Preloading -
-
-#shared_preload_libraries = ''	# (change requires restart)
-#local_preload_libraries = ''
-#session_preload_libraries = ''
-#jit_provider = 'llvmjit'		# JIT library to use
-
-# - Other Defaults -
-
-#dynamic_library_path = '$libdir'
-
-
-#------------------------------------------------------------------------------
-# LOCK MANAGEMENT
-#------------------------------------------------------------------------------
-
-#deadlock_timeout = 1s
-#max_locks_per_transaction = 64		# min 10
-					# (change requires restart)
-#max_pred_locks_per_transaction = 64	# min 10
-					# (change requires restart)
-#max_pred_locks_per_relation = -2	# negative values mean
-					# (max_pred_locks_per_transaction
-					#  / -max_pred_locks_per_relation) - 1
-#max_pred_locks_per_page = 2            # min 0
-
-
-#------------------------------------------------------------------------------
-# VERSION AND PLATFORM COMPATIBILITY
-#------------------------------------------------------------------------------
-
-# - Previous PostgreSQL Versions -
-
-#array_nulls = on
-#backslash_quote = safe_encoding	# on, off, or safe_encoding
-#default_with_oids = off
-#escape_string_warning = on
-#lo_compat_privileges = off
-#operator_precedence_warning = off
-#quote_all_identifiers = off
-#standard_conforming_strings = on
-#synchronize_seqscans = on
-
-# - Other Platforms and Clients -
-
-#transform_null_equals = off
-
-
-#------------------------------------------------------------------------------
-# ERROR HANDLING
-#------------------------------------------------------------------------------
-
-#exit_on_error = off			# terminate session on any error?
-#restart_after_crash = on		# reinitialize after backend crash?
-#data_sync_retry = off			# retry or panic on failure to fsync
-					# data?
-					# (change requires restart)
-
-
-#------------------------------------------------------------------------------
-# CONFIG FILE INCLUDES
-#------------------------------------------------------------------------------
-
-# These options allow settings to be loaded from files other than the
-# default postgresql.conf.  Note that these are directives, not variable
-# assignments, so they can usefully be given more than once.
-
-#include_dir = '...'			# include files ending in '.conf' from
-					# a directory, e.g., 'conf.d'
-#include_if_exists = '...'		# include file only if it exists
-#include = '...'			# include file
-
-
-#------------------------------------------------------------------------------
-# CUSTOMIZED OPTIONS
-#------------------------------------------------------------------------------
-
-# Add settings for extensions here

From 097be95b828596bc62ff3c718831f3a9b1261d4d Mon Sep 17 00:00:00 2001
From: Claus Strasburger <1156864+cfstras@users.noreply.github.com>
Date: Thu, 20 Jun 2024 17:49:53 +0200
Subject: [PATCH 04/12] feat(postgres): add postgres.SQLDriverName for using go
 SQL drivers that don't register as "postgres"

---
 modules/postgres/postgres.go | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/modules/postgres/postgres.go b/modules/postgres/postgres.go
index 43d86cbe19..1bc789a534 100644
--- a/modules/postgres/postgres.go
+++ b/modules/postgres/postgres.go
@@ -19,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
@@ -273,7 +277,7 @@ func (c *PostgresContainer) snapshotConnection(ctx context.Context) (*sql.Conn,
 
 	// Try to use an actual postgres connection, if the driver is loaded
 	connStr := c2.MustConnectionString(ctx, "sslmode=disable")
-	pool, err := sql.Open("postgres", connStr)
+	pool, err := sql.Open(SQLDriverName, connStr)
 	if err != nil {
 		return nil, nil, fmt.Errorf("sql.Open for snapshot connection failed: %w", err)
 	}

From 3529658d5d81f1288ca80b4597cc0f080b89aaa8 Mon Sep 17 00:00:00 2001
From: Claus Strasburger <1156864+cfstras@users.noreply.github.com>
Date: Thu, 20 Jun 2024 17:52:29 +0200
Subject: [PATCH 05/12] test(postgres): Add test for snapshot fallback to
 docker-exec

---
 modules/postgres/postgres_test.go | 105 ++++++++++++++++++++++++++++++
 1 file changed, 105 insertions(+)

diff --git a/modules/postgres/postgres_test.go b/modules/postgres/postgres_test.go
index 54a1f59b48..5c99e52431 100644
--- a/modules/postgres/postgres_test.go
+++ b/modules/postgres/postgres_test.go
@@ -12,6 +12,7 @@ import (
 	"github.com/docker/go-connections/nat"
 	"github.com/jackc/pgx/v5"
 	_ "github.com/lib/pq"
+
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -403,3 +404,107 @@ func TestSnapshotWithOverrides(t *testing.T) {
 		}
 	})
 }
+
+func TestSnapshotWithDockerExecFallback(t *testing.T) {
+	// Tell the postgres module to use a driver that doesn't exist
+	// This will cause the module to fall back to using docker exec
+	postgres.SQLDriverName = "DoesNotExist"
+
+	ctx := context.Background()
+
+	// 1. Start the postgres container and run any migrations on it
+	ctr, err := postgres.RunContainer(
+		ctx,
+		testcontainers.WithImage("docker.io/postgres:16-alpine"),
+		postgres.WithDatabase(dbname),
+		postgres.WithUsername(user),
+		postgres.WithPassword(password),
+		BasicWaitStrategies,
+	)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Run any migrations on the database
+	_, _, err = ctr.Exec(ctx, []string{"psql", "-U", user, "-d", dbname, "-c", "CREATE TABLE users (id SERIAL, name TEXT NOT NULL, age INT NOT NULL)"})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// 2. Create a snapshot of the database to restore later
+	err = ctr.Snapshot(ctx, postgres.WithSnapshotName("test-snapshot"))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Clean up the container after the test is complete
+	t.Cleanup(func() {
+		if err := ctr.Terminate(ctx); err != nil {
+			t.Fatalf("failed to terminate container: %s", err)
+		}
+	})
+
+	dbURL, err := ctr.ConnectionString(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Run("Test inserting a user", func(t *testing.T) {
+		t.Cleanup(func() {
+			// 3. In each test, reset the DB to its snapshot state.
+			err := ctr.Restore(ctx)
+			if err != nil {
+				t.Fatal(err)
+			}
+		})
+
+		conn, err2 := pgx.Connect(context.Background(), dbURL)
+		if err2 != nil {
+			t.Fatal(err2)
+		}
+		defer conn.Close(context.Background())
+
+		_, err2 = conn.Exec(ctx, "INSERT INTO users(name, age) VALUES ($1, $2)", "test", 42)
+		if err2 != nil {
+			t.Fatal(err2)
+		}
+
+		var name string
+		var age int64
+		err2 = conn.QueryRow(context.Background(), "SELECT name, age FROM users LIMIT 1").Scan(&name, &age)
+		if err2 != nil {
+			t.Fatal(err2)
+		}
+
+		if name != "test" {
+			t.Fatalf("Expected %s to equal `test`", name)
+		}
+		if age != 42 {
+			t.Fatalf("Expected %d to equal `42`", age)
+		}
+	})
+
+	t.Run("Test querying empty DB", func(t *testing.T) {
+		// 4. Run as many tests as you need, they will each get a clean database
+		t.Cleanup(func() {
+			err := ctr.Restore(ctx)
+			if err != nil {
+				t.Fatal(err)
+			}
+		})
+
+		conn, err2 := pgx.Connect(context.Background(), dbURL)
+		if err2 != nil {
+			t.Fatal(err2)
+		}
+		defer conn.Close(context.Background())
+
+		var name string
+		var age int64
+		err2 = conn.QueryRow(context.Background(), "SELECT name, age FROM users LIMIT 1").Scan(&name, &age)
+		if !errors.Is(err2, pgx.ErrNoRows) {
+			t.Fatalf("Expected error to be a NoRows error, since the DB should be empty on every test. Got %s instead", err2)
+		}
+	})
+	// }
+}

From 9f8a9d0876a99bfa1fc68b7d997a55f834314aee Mon Sep 17 00:00:00 2001
From: Claus Strasburger <1156864+cfstras@users.noreply.github.com>
Date: Thu, 20 Jun 2024 18:08:40 +0200
Subject: [PATCH 06/12] docs(postgres): Update docs for wait strategies and
 custom SQL drivers

---
 docs/modules/postgres.md          | 39 ++++++++++++++++++++++++++++++-
 modules/postgres/postgres_test.go |  2 ++
 2 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/docs/modules/postgres.md b/docs/modules/postgres.md
index ebd1ed919e..e530bf955e 100644
--- a/docs/modules/postgres.md
+++ b/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,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
+
+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()
+}
+```
diff --git a/modules/postgres/postgres_test.go b/modules/postgres/postgres_test.go
index 5c99e52431..917326a852 100644
--- a/modules/postgres/postgres_test.go
+++ b/modules/postgres/postgres_test.go
@@ -28,6 +28,7 @@ const (
 )
 
 var (
+	// waitStrategy {
 	// BasicWaitStrategies is a simple but reliable way to wait for postgres to start.
 	BasicWaitStrategies = testcontainers.WithWaitStrategy(
 		// First, we wait for the container to log readiness twice.
@@ -39,6 +40,7 @@ var (
 		// Without this, the tests will be flaky on those OSes!
 		wait.ForListeningPort("5432/tcp"),
 	)
+	// }
 )
 
 func TestPostgres(t *testing.T) {

From 4e115834bd5c59c00694c53e6b46ec8999ae5494 Mon Sep 17 00:00:00 2001
From: Claus Strasburger <1156864+cfstras@users.noreply.github.com>
Date: Mon, 24 Jun 2024 13:26:38 +0200
Subject: [PATCH 07/12] feat(postgres): expose postgres.BasicWaitStrategies

---
 modules/postgres/postgres_test.go   | 28 ++++++----------------------
 modules/postgres/wait_strategies.go | 27 +++++++++++++++++++++++++++
 2 files changed, 33 insertions(+), 22 deletions(-)
 create mode 100644 modules/postgres/wait_strategies.go

diff --git a/modules/postgres/postgres_test.go b/modules/postgres/postgres_test.go
index 917326a852..7113bc3f56 100644
--- a/modules/postgres/postgres_test.go
+++ b/modules/postgres/postgres_test.go
@@ -27,22 +27,6 @@ const (
 	password = "password"
 )
 
-var (
-	// waitStrategy {
-	// BasicWaitStrategies is a simple but reliable way to wait for postgres to start.
-	BasicWaitStrategies = testcontainers.WithWaitStrategy(
-		// First, we wait for the container to log readiness twice.
-		// This is because it will restart itself after the first startup.
-		wait.ForLog("database system is ready to accept connections").WithOccurrence(2),
-		// Then, we wait for docker to actually serve the port on localhost.
-		// For non-linux OSes like Mac and Windows, Docker or Rancher Desktop will have to start a separate proxy
-		// process that sometimes takes a little longer to initialize.
-		// Without this, the tests will be flaky on those OSes!
-		wait.ForListeningPort("5432/tcp"),
-	)
-	// }
-)
-
 func TestPostgres(t *testing.T) {
 	ctx := context.Background()
 
@@ -81,7 +65,7 @@ func TestPostgres(t *testing.T) {
 				postgres.WithDatabase(dbname),
 				postgres.WithUsername(user),
 				postgres.WithPassword(password),
-				BasicWaitStrategies,
+				postgres.BasicWaitStrategies(),
 			)
 			if err != nil {
 				t.Fatal(err)
@@ -178,7 +162,7 @@ func TestWithConfigFile(t *testing.T) {
 		postgres.WithDatabase(dbname),
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
-		BasicWaitStrategies,
+		postgres.BasicWaitStrategies(),
 	)
 	if err != nil {
 		t.Fatal(err)
@@ -209,7 +193,7 @@ func TestWithInitScript(t *testing.T) {
 		postgres.WithDatabase(dbname),
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
-		BasicWaitStrategies,
+		postgres.BasicWaitStrategies(),
 	)
 	if err != nil {
 		t.Fatal(err)
@@ -247,7 +231,7 @@ func TestSnapshot(t *testing.T) {
 		postgres.WithDatabase(dbname),
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
-		BasicWaitStrategies,
+		postgres.BasicWaitStrategies(),
 	)
 	if err != nil {
 		t.Fatal(err)
@@ -350,7 +334,7 @@ func TestSnapshotWithOverrides(t *testing.T) {
 		postgres.WithDatabase(dbname),
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
-		BasicWaitStrategies,
+		postgres.BasicWaitStrategies(),
 	)
 	if err != nil {
 		t.Fatal(err)
@@ -421,7 +405,7 @@ func TestSnapshotWithDockerExecFallback(t *testing.T) {
 		postgres.WithDatabase(dbname),
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
-		BasicWaitStrategies,
+		postgres.BasicWaitStrategies(),
 	)
 	if err != nil {
 		t.Fatal(err)
diff --git a/modules/postgres/wait_strategies.go b/modules/postgres/wait_strategies.go
new file mode 100644
index 0000000000..e748f4cc3a
--- /dev/null
+++ b/modules/postgres/wait_strategies.go
@@ -0,0 +1,27 @@
+package postgres
+
+import (
+	"github.com/testcontainers/testcontainers-go"
+	"github.com/testcontainers/testcontainers-go/wait"
+)
+
+// BasicWaitStrategies is a simple but reliable way to wait for postgres to start.
+// It returns a two-step wait strategy:
+//
+//   - It will wait for the container to log `database system is ready to accept connections` twice, because it will restart itself after the first startup.
+//   - It will then wait for docker to actually serve the port on localhost.
+//     For non-linux OSes like Mac and Windows, Docker or Rancher Desktop will have to start a separate proxy.
+//     Without this, the tests will be flaky on those OSes!
+func BasicWaitStrategies() testcontainers.CustomizeRequestOption {
+	// waitStrategy {
+	return testcontainers.WithWaitStrategy(
+		// First, we wait for the container to log readiness twice.
+		// This is because it will restart itself after the first startup.
+		wait.ForLog("database system is ready to accept connections").WithOccurrence(2),
+		// Then, we wait for docker to actually serve the port on localhost.
+		// For non-linux OSes like Mac and Windows, Docker or Rancher Desktop will have to start a separate proxy.
+		// Without this, the tests will be flaky on those OSes!
+		wait.ForListeningPort("5432/tcp"),
+	)
+	// }
+}

From 8f6d922074793708009036ba127b20c9fcc76aa3 Mon Sep 17 00:00:00 2001
From: Claus Strasburger <1156864+cfstras@users.noreply.github.com>
Date: Mon, 24 Jun 2024 13:29:19 +0200
Subject: [PATCH 08/12] docs(postgres): include docs marker of release

---
 docs/modules/postgres.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/docs/modules/postgres.md b/docs/modules/postgres.md
index e530bf955e..18750f9032 100644
--- a/docs/modules/postgres.md
+++ b/docs/modules/postgres.md
@@ -117,6 +117,8 @@ tests very modular, since they always run on a brand-new database.
 
 ### 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.
 

From fbbf348579f9f485bc05a7d4598dfc498b03f98b Mon Sep 17 00:00:00 2001
From: Claus Strasburger <1156864+cfstras@users.noreply.github.com>
Date: Mon, 24 Jun 2024 13:32:57 +0200
Subject: [PATCH 09/12] chore(postgres): safeguard nil check

---
 modules/postgres/postgres.go | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/modules/postgres/postgres.go b/modules/postgres/postgres.go
index 1bc789a534..b3dcfafb6d 100644
--- a/modules/postgres/postgres.go
+++ b/modules/postgres/postgres.go
@@ -253,7 +253,9 @@ func (c *PostgresContainer) execCommandsSQL(ctx context.Context, cmds ...string)
 		testcontainers.Logger.Printf("Could not connect to database to restore snapshot, falling back to `docker exec psql`: %v", err)
 		return c.execCommandsFallback(ctx, cmds)
 	}
-	defer cleanup()
+	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)

From 4b62dc75181e598628b4f83f7ace2e9f44cbf7da Mon Sep 17 00:00:00 2001
From: Claus Strasburger <1156864+cfstras@users.noreply.github.com>
Date: Tue, 25 Jun 2024 14:45:09 +0200
Subject: [PATCH 10/12] style: small comments from review

---
 docs/modules/postgres.md          | 2 +-
 modules/postgres/postgres_test.go | 1 -
 2 files changed, 1 insertion(+), 2 deletions(-)

diff --git a/docs/modules/postgres.md b/docs/modules/postgres.md
index 18750f9032..8421f19bf0 100644
--- a/docs/modules/postgres.md
+++ b/docs/modules/postgres.md
@@ -98,7 +98,7 @@ 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
+[Example Wait Strategies](../../modules/postgres/wait_strategies.go) inside_block:waitStrategy
 <!--/codeinclude-->
 
 ### Using Snapshots
diff --git a/modules/postgres/postgres_test.go b/modules/postgres/postgres_test.go
index 7113bc3f56..476bef4b50 100644
--- a/modules/postgres/postgres_test.go
+++ b/modules/postgres/postgres_test.go
@@ -12,7 +12,6 @@ import (
 	"github.com/docker/go-connections/nat"
 	"github.com/jackc/pgx/v5"
 	_ "github.com/lib/pq"
-
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 

From d8468eca864b3662daef6cd5e200916f18b40b7b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20de=20la=20Pe=C3=B1a?= <mdelapenya@gmail.com>
Date: Fri, 28 Jun 2024 12:32:45 +0200
Subject: [PATCH 11/12] chore: convert SQLDriverName global into a container
 options

---
 docs/modules/postgres.md          | 14 +++++-------
 modules/postgres/go.mod           |  2 ++
 modules/postgres/options.go       | 37 +++++++++++++++++++++++++++++++
 modules/postgres/postgres.go      | 16 ++++++++-----
 modules/postgres/postgres_test.go | 11 +++++----
 5 files changed, 62 insertions(+), 18 deletions(-)
 create mode 100644 modules/postgres/options.go

diff --git a/docs/modules/postgres.md b/docs/modules/postgres.md
index 8421f19bf0..c9414daaf2 100644
--- a/docs/modules/postgres.md
+++ b/docs/modules/postgres.md
@@ -122,7 +122,7 @@ tests very modular, since they always run on a brand-new database.
 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.
+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.
 
@@ -136,12 +136,10 @@ import (
 
 	"github.com/testcontainers/testcontainers-go/modules/postgres"
 )
+```
 
-func init() {
-	postgres.SQLDriverName = "pgx"
-}
+The above code snippet is importing the `pgx` driver and the _Testcontainers for Go_ Postgres module.
 
-func TestSomething(t *testing.T) {
-    // Your test code here, using postgres.RunContainer()
-}
-```
+<!--codeinclude-->
+[Snapshot/Restore with custom driver](../../modules/postgres/postgres_test.go) inside_block:snapshotAndReset
+<!--/codeinclude-->
\ No newline at end of file
diff --git a/modules/postgres/go.mod b/modules/postgres/go.mod
index 3dfc8974be..92a189fc27 100644
--- a/modules/postgres/go.mod
+++ b/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
diff --git a/modules/postgres/options.go b/modules/postgres/options.go
new file mode 100644
index 0000000000..ad24c79fc3
--- /dev/null
+++ b/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
+	}
+}
diff --git a/modules/postgres/postgres.go b/modules/postgres/postgres.go
index b3dcfafb6d..3cb18e5df9 100644
--- a/modules/postgres/postgres.go
+++ b/modules/postgres/postgres.go
@@ -19,10 +19,6 @@ 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
@@ -30,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.
@@ -153,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
 		}
@@ -168,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 {
@@ -279,7 +283,7 @@ func (c *PostgresContainer) snapshotConnection(ctx context.Context) (*sql.Conn,
 
 	// Try to use an actual postgres connection, if the driver is loaded
 	connStr := c2.MustConnectionString(ctx, "sslmode=disable")
-	pool, err := sql.Open(SQLDriverName, connStr)
+	pool, err := sql.Open(c.sqlDriverName, connStr)
 	if err != nil {
 		return nil, nil, fmt.Errorf("sql.Open for snapshot connection failed: %w", err)
 	}
diff --git a/modules/postgres/postgres_test.go b/modules/postgres/postgres_test.go
index 476bef4b50..a645872991 100644
--- a/modules/postgres/postgres_test.go
+++ b/modules/postgres/postgres_test.go
@@ -11,6 +11,7 @@ import (
 
 	"github.com/docker/go-connections/nat"
 	"github.com/jackc/pgx/v5"
+	_ "github.com/jackc/pgx/v5/stdlib"
 	_ "github.com/lib/pq"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -231,6 +232,7 @@ func TestSnapshot(t *testing.T) {
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
 		postgres.BasicWaitStrategies(),
+		postgres.WithSQLDriver("pgx"),
 	)
 	if err != nil {
 		t.Fatal(err)
@@ -391,12 +393,9 @@ func TestSnapshotWithOverrides(t *testing.T) {
 }
 
 func TestSnapshotWithDockerExecFallback(t *testing.T) {
-	// Tell the postgres module to use a driver that doesn't exist
-	// This will cause the module to fall back to using docker exec
-	postgres.SQLDriverName = "DoesNotExist"
-
 	ctx := context.Background()
 
+	// postgresWithSQLDriver {
 	// 1. Start the postgres container and run any migrations on it
 	ctr, err := postgres.RunContainer(
 		ctx,
@@ -405,7 +404,11 @@ func TestSnapshotWithDockerExecFallback(t *testing.T) {
 		postgres.WithUsername(user),
 		postgres.WithPassword(password),
 		postgres.BasicWaitStrategies(),
+		// Tell the postgres module to use a driver that doesn't exist
+		// This will cause the module to fall back to using docker exec
+		postgres.WithSQLDriver("DoesNotExist"),
 	)
+	// }
 	if err != nil {
 		t.Fatal(err)
 	}

From 9330f0badfca726df79ee47350d7775b09b40ca0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Manuel=20de=20la=20Pe=C3=B1a?= <mdelapenya@gmail.com>
Date: Tue, 2 Jul 2024 16:09:37 +0200
Subject: [PATCH 12/12] chore: update container struct for snapshotting

---
 modules/postgres/postgres.go | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/modules/postgres/postgres.go b/modules/postgres/postgres.go
index 3cb18e5df9..8986935bb0 100644
--- a/modules/postgres/postgres.go
+++ b/modules/postgres/postgres.go
@@ -275,10 +275,11 @@ func (c *PostgresContainer) execCommandsSQL(ctx context.Context, cmds ...string)
 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,
+		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