From a2a7059cc587fe5b05b742678998c12a00dae5c7 Mon Sep 17 00:00:00 2001
From: Martin Asquino <martin.asquino@gmail.com>
Date: Tue, 29 Oct 2024 23:25:27 +0000
Subject: [PATCH 01/11] feat: use recommended cluster settings for cockroachdb

---
 modules/cockroachdb/cockroachdb.go | 46 ++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/modules/cockroachdb/cockroachdb.go b/modules/cockroachdb/cockroachdb.go
index 092efa4e2a..cb6a6610e0 100644
--- a/modules/cockroachdb/cockroachdb.go
+++ b/modules/cockroachdb/cockroachdb.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/tls"
 	"crypto/x509"
+	"database/sql"
 	"encoding/pem"
 	"errors"
 	"fmt"
@@ -91,6 +92,11 @@ func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustom
 							return addTLS(ctx, container, o)
 						},
 					},
+					PostReadies: []testcontainers.ContainerHook{
+						func(ctx context.Context, container testcontainers.Container) error {
+							return setRecommendedSettings(ctx, container, o)
+						},
+					},
 				},
 			},
 		},
@@ -234,6 +240,46 @@ func addTLS(ctx context.Context, container testcontainers.Container, opts option
 	return nil
 }
 
+// setRecommendedSettings applies the cluster settings recommended by cockroachlabs for testing clusters.
+// See https://www.cockroachlabs.com/docs/stable/local-testing for more information.
+func setRecommendedSettings(ctx context.Context, container testcontainers.Container, opts options) error {
+	port, err := container.MappedPort(ctx, defaultSQLPort)
+	if err != nil {
+		return err
+	}
+
+	host, err := container.Host(ctx)
+	if err != nil {
+		return err
+	}
+
+	db, err := sql.Open("pgx/v5", connString(opts, host, port))
+	if err != nil {
+		return err
+	}
+	defer db.Close()
+
+	stmts := []string{
+		"SET CLUSTER SETTING kv.range_merge.queue_interval = '50ms'",
+		"SET CLUSTER SETTING jobs.registry.interval.gc = '30s'",
+		"SET CLUSTER SETTING jobs.registry.interval.cancel = '180s'",
+		"SET CLUSTER SETTING jobs.retention_time = '15s'",
+		"SET CLUSTER SETTING sql.stats.automatic_collection.enabled = false",
+		"SET CLUSTER SETTING kv.range_split.by_load_merge_delay = '5s'",
+		`ALTER RANGE default CONFIGURE ZONE USING "gc.ttlseconds" = 600`,
+		`ALTER DATABASE system CONFIGURE ZONE USING "gc.ttlseconds" = 600`,
+	}
+
+	for _, stmt := range stmts {
+		_, err = db.Exec(stmt)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 func connString(opts options, host string, port nat.Port) string {
 	user := url.User(opts.User)
 	if opts.Password != "" {

From 1bcaa9d607e0c8a164ebb3a0e3a29e32382a0cbc Mon Sep 17 00:00:00 2001
From: Martin Asquino <martin.asquino@gmail.com>
Date: Wed, 30 Oct 2024 12:15:25 +0000
Subject: [PATCH 02/11] fix: wrap errors

---
 modules/cockroachdb/cockroachdb.go | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/modules/cockroachdb/cockroachdb.go b/modules/cockroachdb/cockroachdb.go
index cb6a6610e0..f17bc69c4a 100644
--- a/modules/cockroachdb/cockroachdb.go
+++ b/modules/cockroachdb/cockroachdb.go
@@ -245,17 +245,17 @@ func addTLS(ctx context.Context, container testcontainers.Container, opts option
 func setRecommendedSettings(ctx context.Context, container testcontainers.Container, opts options) error {
 	port, err := container.MappedPort(ctx, defaultSQLPort)
 	if err != nil {
-		return err
+		return fmt.Errorf("mapped port: %w", err)
 	}
 
 	host, err := container.Host(ctx)
 	if err != nil {
-		return err
+		return fmt.Errorf("host: %w", err)
 	}
 
 	db, err := sql.Open("pgx/v5", connString(opts, host, port))
 	if err != nil {
-		return err
+		return fmt.Errorf("sql.Open: %w", err)
 	}
 	defer db.Close()
 
@@ -273,7 +273,7 @@ func setRecommendedSettings(ctx context.Context, container testcontainers.Contai
 	for _, stmt := range stmts {
 		_, err = db.Exec(stmt)
 		if err != nil {
-			return err
+			return fmt.Errorf("db.Exec: %w", err)
 		}
 	}
 

From 4a7f3e3583594aa268b3d63520e7419d58465f6d Mon Sep 17 00:00:00 2001
From: Martin Asquino <martin.asquino@gmail.com>
Date: Wed, 30 Oct 2024 13:25:39 +0000
Subject: [PATCH 03/11] feat: add WithStatements option and expose recommended
 settings as a list of statements

---
 modules/cockroachdb/cockroachdb.go | 24 ++++++++---------------
 modules/cockroachdb/options.go     | 31 +++++++++++++++++++++++++-----
 2 files changed, 34 insertions(+), 21 deletions(-)

diff --git a/modules/cockroachdb/cockroachdb.go b/modules/cockroachdb/cockroachdb.go
index f17bc69c4a..53cb8701a9 100644
--- a/modules/cockroachdb/cockroachdb.go
+++ b/modules/cockroachdb/cockroachdb.go
@@ -94,7 +94,7 @@ func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustom
 					},
 					PostReadies: []testcontainers.ContainerHook{
 						func(ctx context.Context, container testcontainers.Container) error {
-							return setRecommendedSettings(ctx, container, o)
+							return runStatements(ctx, container, o)
 						},
 					},
 				},
@@ -240,9 +240,12 @@ func addTLS(ctx context.Context, container testcontainers.Container, opts option
 	return nil
 }
 
-// setRecommendedSettings applies the cluster settings recommended by cockroachlabs for testing clusters.
-// See https://www.cockroachlabs.com/docs/stable/local-testing for more information.
-func setRecommendedSettings(ctx context.Context, container testcontainers.Container, opts options) error {
+// runStatements runs the configured statements against the CockroachDB container.
+func runStatements(ctx context.Context, container testcontainers.Container, opts options) error {
+	if len(opts.Statements) == 0 {
+		return nil
+	}
+
 	port, err := container.MappedPort(ctx, defaultSQLPort)
 	if err != nil {
 		return fmt.Errorf("mapped port: %w", err)
@@ -259,18 +262,7 @@ func setRecommendedSettings(ctx context.Context, container testcontainers.Contai
 	}
 	defer db.Close()
 
-	stmts := []string{
-		"SET CLUSTER SETTING kv.range_merge.queue_interval = '50ms'",
-		"SET CLUSTER SETTING jobs.registry.interval.gc = '30s'",
-		"SET CLUSTER SETTING jobs.registry.interval.cancel = '180s'",
-		"SET CLUSTER SETTING jobs.retention_time = '15s'",
-		"SET CLUSTER SETTING sql.stats.automatic_collection.enabled = false",
-		"SET CLUSTER SETTING kv.range_split.by_load_merge_delay = '5s'",
-		`ALTER RANGE default CONFIGURE ZONE USING "gc.ttlseconds" = 600`,
-		`ALTER DATABASE system CONFIGURE ZONE USING "gc.ttlseconds" = 600`,
-	}
-
-	for _, stmt := range stmts {
+	for _, stmt := range opts.Statements {
 		_, err = db.Exec(stmt)
 		if err != nil {
 			return fmt.Errorf("db.Exec: %w", err)
diff --git a/modules/cockroachdb/options.go b/modules/cockroachdb/options.go
index a2211d77e7..2acde60ee6 100644
--- a/modules/cockroachdb/options.go
+++ b/modules/cockroachdb/options.go
@@ -3,11 +3,12 @@ package cockroachdb
 import "github.com/testcontainers/testcontainers-go"
 
 type options struct {
-	Database  string
-	User      string
-	Password  string
-	StoreSize string
-	TLS       *TLSConfig
+	Database   string
+	User       string
+	Password   string
+	StoreSize  string
+	TLS        *TLSConfig
+	Statements []string
 }
 
 func defaultOptions() options {
@@ -67,3 +68,23 @@ func WithTLS(cfg *TLSConfig) Option {
 		o.TLS = cfg
 	}
 }
+
+// ClusterDefaults are the settings recommended by Cockroach Labs for testing clusters.
+// See https://www.cockroachlabs.com/docs/stable/local-testing for more information.
+var ClusterDefaults []string = []string{
+	"SET CLUSTER SETTING kv.range_merge.queue_interval = '50ms'",
+	"SET CLUSTER SETTING jobs.registry.interval.gc = '30s'",
+	"SET CLUSTER SETTING jobs.registry.interval.cancel = '180s'",
+	"SET CLUSTER SETTING jobs.retention_time = '15s'",
+	"SET CLUSTER SETTING sql.stats.automatic_collection.enabled = false",
+	"SET CLUSTER SETTING kv.range_split.by_load_merge_delay = '5s'",
+	`ALTER RANGE default CONFIGURE ZONE USING "gc.ttlseconds" = 600`,
+	`ALTER DATABASE system CONFIGURE ZONE USING "gc.ttlseconds" = 600`,
+}
+
+// WithStatements sets the statements to run on the CockroachDB cluster once the container is ready.
+func WithStatements(statements ...string) Option {
+	return func(o *options) {
+		o.Statements = statements
+	}
+}

From 404652797878e35ed5d326bcbaea03eb9ae16453 Mon Sep 17 00:00:00 2001
From: martskins <martskins@users.noreply.github.com>
Date: Wed, 30 Oct 2024 14:16:52 +0000
Subject: [PATCH 04/11] Update modules/cockroachdb/options.go

Co-authored-by: Steven Hartland <stevenmhartland@gmail.com>
---
 modules/cockroachdb/options.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/modules/cockroachdb/options.go b/modules/cockroachdb/options.go
index 2acde60ee6..fee22f3cc2 100644
--- a/modules/cockroachdb/options.go
+++ b/modules/cockroachdb/options.go
@@ -71,7 +71,7 @@ func WithTLS(cfg *TLSConfig) Option {
 
 // ClusterDefaults are the settings recommended by Cockroach Labs for testing clusters.
 // See https://www.cockroachlabs.com/docs/stable/local-testing for more information.
-var ClusterDefaults []string = []string{
+var ClusterDefaults = []string{
 	"SET CLUSTER SETTING kv.range_merge.queue_interval = '50ms'",
 	"SET CLUSTER SETTING jobs.registry.interval.gc = '30s'",
 	"SET CLUSTER SETTING jobs.registry.interval.cancel = '180s'",

From d08af654749912737f7255668695af2e2b2c04bb Mon Sep 17 00:00:00 2001
From: martskins <martskins@users.noreply.github.com>
Date: Wed, 30 Oct 2024 14:16:57 +0000
Subject: [PATCH 05/11] Update modules/cockroachdb/cockroachdb.go

Co-authored-by: Steven Hartland <stevenmhartland@gmail.com>
---
 modules/cockroachdb/cockroachdb.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/modules/cockroachdb/cockroachdb.go b/modules/cockroachdb/cockroachdb.go
index 53cb8701a9..fcd83c90f6 100644
--- a/modules/cockroachdb/cockroachdb.go
+++ b/modules/cockroachdb/cockroachdb.go
@@ -263,8 +263,7 @@ func runStatements(ctx context.Context, container testcontainers.Container, opts
 	defer db.Close()
 
 	for _, stmt := range opts.Statements {
-		_, err = db.Exec(stmt)
-		if err != nil {
+		if _, err = db.Exec(stmt); err != nil {
 			return fmt.Errorf("db.Exec: %w", err)
 		}
 	}

From a03ed97a10638ff13e167b46f4c42c207064e725 Mon Sep 17 00:00:00 2001
From: Martin Asquino <martin.asquino@gmail.com>
Date: Wed, 30 Oct 2024 14:29:51 +0000
Subject: [PATCH 06/11] fix: return close error on defer

---
 modules/cockroachdb/cockroachdb.go | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/modules/cockroachdb/cockroachdb.go b/modules/cockroachdb/cockroachdb.go
index fcd83c90f6..884d8f076f 100644
--- a/modules/cockroachdb/cockroachdb.go
+++ b/modules/cockroachdb/cockroachdb.go
@@ -241,7 +241,7 @@ func addTLS(ctx context.Context, container testcontainers.Container, opts option
 }
 
 // runStatements runs the configured statements against the CockroachDB container.
-func runStatements(ctx context.Context, container testcontainers.Container, opts options) error {
+func runStatements(ctx context.Context, container testcontainers.Container, opts options) (err error) {
 	if len(opts.Statements) == 0 {
 		return nil
 	}
@@ -260,7 +260,12 @@ func runStatements(ctx context.Context, container testcontainers.Container, opts
 	if err != nil {
 		return fmt.Errorf("sql.Open: %w", err)
 	}
-	defer db.Close()
+	defer func() {
+		cerr := db.Close()
+		if err == nil {
+			err = cerr
+		}
+	}()
 
 	for _, stmt := range opts.Statements {
 		if _, err = db.Exec(stmt); err != nil {

From 94a81ba485e24bc0af4306d76171597a59c19c76 Mon Sep 17 00:00:00 2001
From: Martin Asquino <martin.asquino@gmail.com>
Date: Wed, 30 Oct 2024 15:29:02 +0000
Subject: [PATCH 07/11] chore: improvement documentation and set default value
 for statements in defaultOptions

---
 modules/cockroachdb/examples_test.go | 59 ++++++++++++++++++++++++++++
 modules/cockroachdb/options.go       | 11 ++++--
 2 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/modules/cockroachdb/examples_test.go b/modules/cockroachdb/examples_test.go
index c06c97596b..2f6c5ea1a5 100644
--- a/modules/cockroachdb/examples_test.go
+++ b/modules/cockroachdb/examples_test.go
@@ -2,6 +2,7 @@ package cockroachdb_test
 
 import (
 	"context"
+	"database/sql"
 	"fmt"
 	"log"
 	"net/url"
@@ -50,3 +51,61 @@ func ExampleRun() {
 	// true
 	// postgres://root@localhost:xxx/defaultdb?sslmode=disable
 }
+
+func ExampleRun_withRecommendedSettings() {
+	ctx := context.Background()
+
+	cockroachdbContainer, err := cockroachdb.Run(ctx, "cockroachdb/cockroach:latest-v23.1", cockroachdb.WithStatements(cockroachdb.ClusterDefaults...))
+	defer func() {
+		if err := testcontainers.TerminateContainer(cockroachdbContainer); err != nil {
+			log.Printf("failed to terminate container: %s", err)
+		}
+	}()
+	if err != nil {
+		log.Printf("failed to start container: %s", err)
+		return
+	}
+
+	state, err := cockroachdbContainer.State(ctx)
+	if err != nil {
+		log.Printf("failed to get container state: %s", err)
+		return
+	}
+	fmt.Println(state.Running)
+
+	addr, err := cockroachdbContainer.ConnectionString(ctx)
+	if err != nil {
+		log.Printf("failed to get connection string: %s", err)
+		return
+	}
+
+	db, err := sql.Open("pgx/v5", addr)
+	if err != nil {
+		log.Printf("failed to open connection: %s", err)
+		return
+	}
+	defer func() {
+		if err := db.Close(); err != nil {
+			log.Printf("failed to close connection: %s", err)
+		}
+	}()
+
+	var queueInterval string
+	if err := db.QueryRow("SHOW CLUSTER SETTING kv.range_merge.queue_interval").Scan(&queueInterval); err != nil {
+		log.Printf("failed to scan row: %s", err)
+		return
+	}
+	fmt.Println(queueInterval)
+
+	var statsCollectionEnabled bool
+	if err := db.QueryRow("SHOW CLUSTER SETTING sql.stats.automatic_collection.enabled").Scan(&statsCollectionEnabled); err != nil {
+		log.Printf("failed to scan row: %s", err)
+		return
+	}
+	fmt.Println(statsCollectionEnabled)
+
+	// Output:
+	// true
+	// 00:00:00.05
+	// false
+}
diff --git a/modules/cockroachdb/options.go b/modules/cockroachdb/options.go
index fee22f3cc2..192693825e 100644
--- a/modules/cockroachdb/options.go
+++ b/modules/cockroachdb/options.go
@@ -13,10 +13,11 @@ type options struct {
 
 func defaultOptions() options {
 	return options{
-		User:      defaultUser,
-		Password:  defaultPassword,
-		Database:  defaultDatabase,
-		StoreSize: defaultStoreSize,
+		User:       defaultUser,
+		Password:   defaultPassword,
+		Database:   defaultDatabase,
+		StoreSize:  defaultStoreSize,
+		Statements: ClusterDefaults,
 	}
 }
 
@@ -83,6 +84,8 @@ var ClusterDefaults = []string{
 }
 
 // WithStatements sets the statements to run on the CockroachDB cluster once the container is ready.
+// This, in combination with ClusterDefaults, can be used to configure the cluster with the settings
+// recommended by Cockroach Labs.
 func WithStatements(statements ...string) Option {
 	return func(o *options) {
 		o.Statements = statements

From 457039780998c0d32c805529a81c38b1b70c8665 Mon Sep 17 00:00:00 2001
From: martskins <martskins@users.noreply.github.com>
Date: Wed, 30 Oct 2024 15:31:57 +0000
Subject: [PATCH 08/11] Update modules/cockroachdb/cockroachdb.go
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Manuel de la Peña <social.mdelapenya@gmail.com>
---
 modules/cockroachdb/cockroachdb.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/modules/cockroachdb/cockroachdb.go b/modules/cockroachdb/cockroachdb.go
index 884d8f076f..34bba7c9ab 100644
--- a/modules/cockroachdb/cockroachdb.go
+++ b/modules/cockroachdb/cockroachdb.go
@@ -241,7 +241,7 @@ func addTLS(ctx context.Context, container testcontainers.Container, opts option
 }
 
 // runStatements runs the configured statements against the CockroachDB container.
-func runStatements(ctx context.Context, container testcontainers.Container, opts options) (err error) {
+func runStatements(ctx context.Context, container testcontainers.Container, opts options) (err error) { //nolint:nonamedreturns // Needed for error checking.
 	if len(opts.Statements) == 0 {
 		return nil
 	}

From cf8625d597c8f47ea3514099b06d405495d32ef8 Mon Sep 17 00:00:00 2001
From: Martin Asquino <martin.asquino@gmail.com>
Date: Wed, 30 Oct 2024 15:59:27 +0000
Subject: [PATCH 09/11] chore: remove unnecessary nolint directive and rename
 ClusterDefaults to DefaultStatements

---
 modules/cockroachdb/cockroachdb.go   | 2 +-
 modules/cockroachdb/examples_test.go | 2 +-
 modules/cockroachdb/options.go       | 8 ++++----
 3 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/modules/cockroachdb/cockroachdb.go b/modules/cockroachdb/cockroachdb.go
index 34bba7c9ab..884d8f076f 100644
--- a/modules/cockroachdb/cockroachdb.go
+++ b/modules/cockroachdb/cockroachdb.go
@@ -241,7 +241,7 @@ func addTLS(ctx context.Context, container testcontainers.Container, opts option
 }
 
 // runStatements runs the configured statements against the CockroachDB container.
-func runStatements(ctx context.Context, container testcontainers.Container, opts options) (err error) { //nolint:nonamedreturns // Needed for error checking.
+func runStatements(ctx context.Context, container testcontainers.Container, opts options) (err error) {
 	if len(opts.Statements) == 0 {
 		return nil
 	}
diff --git a/modules/cockroachdb/examples_test.go b/modules/cockroachdb/examples_test.go
index 2f6c5ea1a5..9a8fb12881 100644
--- a/modules/cockroachdb/examples_test.go
+++ b/modules/cockroachdb/examples_test.go
@@ -55,7 +55,7 @@ func ExampleRun() {
 func ExampleRun_withRecommendedSettings() {
 	ctx := context.Background()
 
-	cockroachdbContainer, err := cockroachdb.Run(ctx, "cockroachdb/cockroach:latest-v23.1", cockroachdb.WithStatements(cockroachdb.ClusterDefaults...))
+	cockroachdbContainer, err := cockroachdb.Run(ctx, "cockroachdb/cockroach:latest-v23.1", cockroachdb.WithStatements(cockroachdb.DefaultStatements...))
 	defer func() {
 		if err := testcontainers.TerminateContainer(cockroachdbContainer); err != nil {
 			log.Printf("failed to terminate container: %s", err)
diff --git a/modules/cockroachdb/options.go b/modules/cockroachdb/options.go
index 192693825e..90e67a1c2d 100644
--- a/modules/cockroachdb/options.go
+++ b/modules/cockroachdb/options.go
@@ -17,7 +17,7 @@ func defaultOptions() options {
 		Password:   defaultPassword,
 		Database:   defaultDatabase,
 		StoreSize:  defaultStoreSize,
-		Statements: ClusterDefaults,
+		Statements: DefaultStatements,
 	}
 }
 
@@ -70,9 +70,9 @@ func WithTLS(cfg *TLSConfig) Option {
 	}
 }
 
-// ClusterDefaults are the settings recommended by Cockroach Labs for testing clusters.
+// DefaultStatements are the settings recommended by Cockroach Labs for testing clusters.
 // See https://www.cockroachlabs.com/docs/stable/local-testing for more information.
-var ClusterDefaults = []string{
+var DefaultStatements = []string{
 	"SET CLUSTER SETTING kv.range_merge.queue_interval = '50ms'",
 	"SET CLUSTER SETTING jobs.registry.interval.gc = '30s'",
 	"SET CLUSTER SETTING jobs.registry.interval.cancel = '180s'",
@@ -84,7 +84,7 @@ var ClusterDefaults = []string{
 }
 
 // WithStatements sets the statements to run on the CockroachDB cluster once the container is ready.
-// This, in combination with ClusterDefaults, can be used to configure the cluster with the settings
+// This, in combination with DefaultStatements, can be used to configure the cluster with the settings
 // recommended by Cockroach Labs.
 func WithStatements(statements ...string) Option {
 	return func(o *options) {

From 50878ca1cc7374e9f907d6eb940e237805ef4a41 Mon Sep 17 00:00:00 2001
From: Martin Asquino <martin.asquino@gmail.com>
Date: Wed, 30 Oct 2024 22:06:44 +0000
Subject: [PATCH 10/11] fix: only use recommended settings in root user tests

---
 modules/cockroachdb/cockroachdb_test.go | 9 +++++++++
 modules/cockroachdb/options.go          | 1 +
 2 files changed, 10 insertions(+)

diff --git a/modules/cockroachdb/cockroachdb_test.go b/modules/cockroachdb/cockroachdb_test.go
index cc355e9168..45df7909bb 100644
--- a/modules/cockroachdb/cockroachdb_test.go
+++ b/modules/cockroachdb/cockroachdb_test.go
@@ -28,6 +28,9 @@ func TestCockroach_NotRoot(t *testing.T) {
 		url: "postgres://test@localhost:xxxxx/defaultdb?sslmode=disable",
 		opts: []testcontainers.ContainerCustomizer{
 			cockroachdb.WithUser("test"),
+			// Do not run the default statements as the user used on this test is
+			// lacking the needed MODIFYCLUSTERSETTING privilege to run them.
+			cockroachdb.WithStatements(),
 		},
 	})
 }
@@ -38,6 +41,9 @@ func TestCockroach_Password(t *testing.T) {
 		opts: []testcontainers.ContainerCustomizer{
 			cockroachdb.WithUser("foo"),
 			cockroachdb.WithPassword("bar"),
+			// Do not run the default statements as the user used on this test is
+			// lacking the needed MODIFYCLUSTERSETTING privilege to run them.
+			cockroachdb.WithStatements(),
 		},
 	})
 }
@@ -50,6 +56,9 @@ func TestCockroach_TLS(t *testing.T) {
 		url: "postgres://root@localhost:xxxxx/defaultdb?sslmode=verify-full",
 		opts: []testcontainers.ContainerCustomizer{
 			cockroachdb.WithTLS(tlsCfg),
+			// Do not run the default statements as the user used on this test is
+			// lacking the needed MODIFYCLUSTERSETTING privilege to run them.
+			cockroachdb.WithStatements(),
 		},
 	})
 }
diff --git a/modules/cockroachdb/options.go b/modules/cockroachdb/options.go
index 90e67a1c2d..eba101834e 100644
--- a/modules/cockroachdb/options.go
+++ b/modules/cockroachdb/options.go
@@ -71,6 +71,7 @@ func WithTLS(cfg *TLSConfig) Option {
 }
 
 // DefaultStatements are the settings recommended by Cockroach Labs for testing clusters.
+// Note that to use these defaults the user needs to have MODIFYCLUSTERSETTING privilege.
 // See https://www.cockroachlabs.com/docs/stable/local-testing for more information.
 var DefaultStatements = []string{
 	"SET CLUSTER SETTING kv.range_merge.queue_interval = '50ms'",

From c58631bf667c03061ec5a10745fd9048655f9e7a Mon Sep 17 00:00:00 2001
From: Steven Hartland <stevenmhartland@gmail.com>
Date: Fri, 1 Nov 2024 14:37:34 +0000
Subject: [PATCH 11/11] feat!(cockroachdb): simplify connection handling

Simplify the connection handling in cockroachdb so that ConnectionString
can be used without the user doing extra work to handle TLS if enabled.

Deprecate TLSConfig which is no longer needed separately.

BREAKING_CHANGE: This now returns a registered connection string so is
no longer compatible with pgx.ParseConfig, use ConnectionConfig for this
use case instead.
---
 modules/cockroachdb/certs.go            |  19 ++++
 modules/cockroachdb/cockroachdb.go      | 118 ++++++------------------
 modules/cockroachdb/cockroachdb_test.go |  52 +++--------
 modules/cockroachdb/examples_test.go    |  26 ++++--
 modules/cockroachdb/options.go          |  92 +++++++++++++++++-
 5 files changed, 168 insertions(+), 139 deletions(-)

diff --git a/modules/cockroachdb/certs.go b/modules/cockroachdb/certs.go
index afa12fcd1a..61280b4db9 100644
--- a/modules/cockroachdb/certs.go
+++ b/modules/cockroachdb/certs.go
@@ -1,8 +1,10 @@
 package cockroachdb
 
 import (
+	"crypto/tls"
 	"crypto/x509"
 	"errors"
+	"fmt"
 	"net"
 	"time"
 
@@ -65,3 +67,20 @@ func NewTLSConfig() (*TLSConfig, error) {
 		ClientKey:  clientCert.KeyBytes,
 	}, nil
 }
+
+// tlsConfig returns a [tls.Config] for options.
+func (c *TLSConfig) tlsConfig() (*tls.Config, error) {
+	keyPair, err := tls.X509KeyPair(c.ClientCert, c.ClientKey)
+	if err != nil {
+		return nil, fmt.Errorf("x509 key pair: %w", err)
+	}
+
+	certPool := x509.NewCertPool()
+	certPool.AddCert(c.CACert)
+
+	return &tls.Config{
+		RootCAs:      certPool,
+		Certificates: []tls.Certificate{keyPair},
+		ServerName:   "localhost",
+	}, nil
+}
diff --git a/modules/cockroachdb/cockroachdb.go b/modules/cockroachdb/cockroachdb.go
index 884d8f076f..98c5ebf149 100644
--- a/modules/cockroachdb/cockroachdb.go
+++ b/modules/cockroachdb/cockroachdb.go
@@ -3,23 +3,20 @@ package cockroachdb
 import (
 	"context"
 	"crypto/tls"
-	"crypto/x509"
 	"database/sql"
 	"encoding/pem"
 	"errors"
 	"fmt"
-	"net"
-	"net/url"
 	"path/filepath"
 
 	"github.com/docker/go-connections/nat"
 	"github.com/jackc/pgx/v5"
-	"github.com/jackc/pgx/v5/stdlib"
 
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/wait"
 )
 
+// ErrTLSNotEnabled is returned when trying to get a TLS config from a container that does not have TLS enabled.
 var ErrTLSNotEnabled = errors.New("tls not enabled")
 
 const (
@@ -40,7 +37,9 @@ type CockroachDBContainer struct {
 	opts options
 }
 
-// MustConnectionString panics if the address cannot be determined.
+// MustConnectionString returns a connection string to open a new connection to CockroachDB
+// as described by [CockroachDBContainer.ConnectionString].
+// It panics if an error occurs.
 func (c *CockroachDBContainer) MustConnectionString(ctx context.Context) string {
 	addr, err := c.ConnectionString(ctx)
 	if err != nil {
@@ -49,27 +48,33 @@ func (c *CockroachDBContainer) MustConnectionString(ctx context.Context) string
 	return addr
 }
 
-// ConnectionString returns the dial address to open a new connection to CockroachDB.
+// ConnectionString returns a connection string to open a new connection to CockroachDB.
+// The returned string is suitable for use by [sql.Open] but is not be compatible with
+// [pgx.ParseConfig], so if you want to call [pgx.ConnectConfig] use the
+// [CockroachDBContainer.ConnectionConfig] method instead.
 func (c *CockroachDBContainer) ConnectionString(ctx context.Context) (string, error) {
-	port, err := c.MappedPort(ctx, defaultSQLPort)
-	if err != nil {
-		return "", err
-	}
-
-	host, err := c.Host(ctx)
-	if err != nil {
-		return "", err
-	}
+	return c.opts.containerConnString(ctx, c.Container)
+}
 
-	return connString(c.opts, host, port), nil
+// ConnectionConfig returns a [pgx.ConnConfig] for the CockroachDB container.
+// This can be passed to [pgx.ConnectConfig] to open a new connection.
+func (c *CockroachDBContainer) ConnectionConfig(ctx context.Context) (*pgx.ConnConfig, error) {
+	return c.opts.containerConnConfig(ctx, c.Container)
 }
 
 // TLSConfig returns config necessary to connect to CockroachDB over TLS.
+//
+// Deprecated: use [CockroachDBContainer.ConnectionConfig] or
+// [CockroachDBContainer.ConnectionConfig] instead.
 func (c *CockroachDBContainer) TLSConfig() (*tls.Config, error) {
-	return connTLS(c.opts)
+	if c.opts.TLS == nil {
+		return nil, ErrTLSNotEnabled
+	}
+
+	return c.opts.TLS.tlsConfig()
 }
 
-// Deprecated: use Run instead
+// Deprecated: use Run instead.
 // RunContainer creates an instance of the CockroachDB container type
 func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomizer) (*CockroachDBContainer, error) {
 	return Run(ctx, "cockroachdb/cockroach:latest-v23.1", opts...)
@@ -178,29 +183,12 @@ func addEnvs(req *testcontainers.GenericContainerRequest, opts options) error {
 }
 
 func addWaitingFor(req *testcontainers.GenericContainerRequest, opts options) error {
-	var tlsConfig *tls.Config
-	if opts.TLS != nil {
-		cfg, err := connTLS(opts)
-		if err != nil {
-			return err
-		}
-		tlsConfig = cfg
-	}
-
 	sqlWait := wait.ForSQL(defaultSQLPort, "pgx/v5", func(host string, port nat.Port) string {
-		connStr := connString(opts, host, port)
-		if tlsConfig == nil {
-			return connStr
-		}
-
-		// register TLS config with pgx driver
-		connCfg, err := pgx.ParseConfig(connStr)
+		connStr, err := opts.connString(host, port)
 		if err != nil {
 			panic(err)
 		}
-		connCfg.TLSConfig = tlsConfig
-
-		return stdlib.RegisterConnConfig(connCfg)
+		return connStr
 	})
 	defaultStrategy := wait.ForAll(
 		wait.ForHTTP("/health").WithPort(defaultAdminPort),
@@ -246,17 +234,12 @@ func runStatements(ctx context.Context, container testcontainers.Container, opts
 		return nil
 	}
 
-	port, err := container.MappedPort(ctx, defaultSQLPort)
-	if err != nil {
-		return fmt.Errorf("mapped port: %w", err)
-	}
-
-	host, err := container.Host(ctx)
+	connStr, err := opts.containerConnString(ctx, container)
 	if err != nil {
-		return fmt.Errorf("host: %w", err)
+		return fmt.Errorf("connection string: %w", err)
 	}
 
-	db, err := sql.Open("pgx/v5", connString(opts, host, port))
+	db, err := sql.Open("pgx/v5", connStr)
 	if err != nil {
 		return fmt.Errorf("sql.Open: %w", err)
 	}
@@ -275,48 +258,3 @@ func runStatements(ctx context.Context, container testcontainers.Container, opts
 
 	return nil
 }
-
-func connString(opts options, host string, port nat.Port) string {
-	user := url.User(opts.User)
-	if opts.Password != "" {
-		user = url.UserPassword(opts.User, opts.Password)
-	}
-
-	sslMode := "disable"
-	if opts.TLS != nil {
-		sslMode = "verify-full"
-	}
-	params := url.Values{
-		"sslmode": []string{sslMode},
-	}
-
-	u := url.URL{
-		Scheme:   "postgres",
-		User:     user,
-		Host:     net.JoinHostPort(host, port.Port()),
-		Path:     opts.Database,
-		RawQuery: params.Encode(),
-	}
-
-	return u.String()
-}
-
-func connTLS(opts options) (*tls.Config, error) {
-	if opts.TLS == nil {
-		return nil, ErrTLSNotEnabled
-	}
-
-	keyPair, err := tls.X509KeyPair(opts.TLS.ClientCert, opts.TLS.ClientKey)
-	if err != nil {
-		return nil, err
-	}
-
-	certPool := x509.NewCertPool()
-	certPool.AddCert(opts.TLS.CACert)
-
-	return &tls.Config{
-		RootCAs:      certPool,
-		Certificates: []tls.Certificate{keyPair},
-		ServerName:   "localhost",
-	}, nil
-}
diff --git a/modules/cockroachdb/cockroachdb_test.go b/modules/cockroachdb/cockroachdb_test.go
index 45df7909bb..d3d05d1d86 100644
--- a/modules/cockroachdb/cockroachdb_test.go
+++ b/modules/cockroachdb/cockroachdb_test.go
@@ -2,9 +2,6 @@ package cockroachdb_test
 
 import (
 	"context"
-	"errors"
-	"net/url"
-	"strings"
 	"testing"
 	"time"
 
@@ -18,14 +15,11 @@ import (
 )
 
 func TestCockroach_Insecure(t *testing.T) {
-	suite.Run(t, &AuthNSuite{
-		url: "postgres://root@localhost:xxxxx/defaultdb?sslmode=disable",
-	})
+	suite.Run(t, &AuthNSuite{})
 }
 
 func TestCockroach_NotRoot(t *testing.T) {
 	suite.Run(t, &AuthNSuite{
-		url: "postgres://test@localhost:xxxxx/defaultdb?sslmode=disable",
 		opts: []testcontainers.ContainerCustomizer{
 			cockroachdb.WithUser("test"),
 			// Do not run the default statements as the user used on this test is
@@ -37,7 +31,6 @@ func TestCockroach_NotRoot(t *testing.T) {
 
 func TestCockroach_Password(t *testing.T) {
 	suite.Run(t, &AuthNSuite{
-		url: "postgres://foo:bar@localhost:xxxxx/defaultdb?sslmode=disable",
 		opts: []testcontainers.ContainerCustomizer{
 			cockroachdb.WithUser("foo"),
 			cockroachdb.WithPassword("bar"),
@@ -53,19 +46,26 @@ func TestCockroach_TLS(t *testing.T) {
 	require.NoError(t, err)
 
 	suite.Run(t, &AuthNSuite{
-		url: "postgres://root@localhost:xxxxx/defaultdb?sslmode=verify-full",
 		opts: []testcontainers.ContainerCustomizer{
 			cockroachdb.WithTLS(tlsCfg),
-			// Do not run the default statements as the user used on this test is
-			// lacking the needed MODIFYCLUSTERSETTING privilege to run them.
-			cockroachdb.WithStatements(),
 		},
 	})
 }
 
+func TestTLS(t *testing.T) {
+	tlsCfg, err := cockroachdb.NewTLSConfig()
+	require.NoError(t, err)
+
+	ctx := context.Background()
+
+	ctr, err := cockroachdb.Run(ctx, "cockroachdb/cockroach:latest-v23.1", cockroachdb.WithTLS(tlsCfg))
+	testcontainers.CleanupContainer(t, ctr)
+	require.NoError(t, err)
+	require.NotNil(t, ctr)
+}
+
 type AuthNSuite struct {
 	suite.Suite
-	url  string
 	opts []testcontainers.ContainerCustomizer
 }
 
@@ -75,11 +75,6 @@ func (suite *AuthNSuite) TestConnectionString() {
 	ctr, err := cockroachdb.Run(ctx, "cockroachdb/cockroach:latest-v23.1", suite.opts...)
 	testcontainers.CleanupContainer(suite.T(), ctr)
 	suite.Require().NoError(err)
-
-	connStr, err := removePort(ctr.MustConnectionString(ctx))
-	suite.Require().NoError(err)
-
-	suite.Equal(suite.url, connStr)
 }
 
 func (suite *AuthNSuite) TestPing() {
@@ -203,29 +198,10 @@ func (suite *AuthNSuite) TestWithWaitStrategyAndDeadline() {
 }
 
 func conn(ctx context.Context, container *cockroachdb.CockroachDBContainer) (*pgx.Conn, error) {
-	cfg, err := pgx.ParseConfig(container.MustConnectionString(ctx))
+	cfg, err := container.ConnectionConfig(ctx)
 	if err != nil {
 		return nil, err
 	}
 
-	tlsCfg, err := container.TLSConfig()
-	switch {
-	case err != nil:
-		if !errors.Is(err, cockroachdb.ErrTLSNotEnabled) {
-			return nil, err
-		}
-	default:
-		// apply TLS config
-		cfg.TLSConfig = tlsCfg
-	}
-
 	return pgx.ConnectConfig(ctx, cfg)
 }
-
-func removePort(s string) (string, error) {
-	u, err := url.Parse(s)
-	if err != nil {
-		return "", err
-	}
-	return strings.Replace(s, ":"+u.Port(), ":xxxxx", 1), nil
-}
diff --git a/modules/cockroachdb/examples_test.go b/modules/cockroachdb/examples_test.go
index 9a8fb12881..4cc14e8b7b 100644
--- a/modules/cockroachdb/examples_test.go
+++ b/modules/cockroachdb/examples_test.go
@@ -5,7 +5,8 @@ import (
 	"database/sql"
 	"fmt"
 	"log"
-	"net/url"
+
+	"github.com/jackc/pgx/v5"
 
 	"github.com/testcontainers/testcontainers-go"
 	"github.com/testcontainers/testcontainers-go/modules/cockroachdb"
@@ -34,25 +35,34 @@ func ExampleRun() {
 	}
 	fmt.Println(state.Running)
 
-	addr, err := cockroachdbContainer.ConnectionString(ctx)
+	cfg, err := cockroachdbContainer.ConnectionConfig(ctx)
 	if err != nil {
 		log.Printf("failed to get connection string: %s", err)
 		return
 	}
-	u, err := url.Parse(addr)
+
+	conn, err := pgx.ConnectConfig(ctx, cfg)
 	if err != nil {
-		log.Printf("failed to parse connection string: %s", err)
+		log.Printf("failed to connect: %s", err)
+		return
+	}
+
+	defer func() {
+		if err := conn.Close(ctx); err != nil {
+			log.Printf("failed to close connection: %s", err)
+		}
+	}()
+
+	if err = conn.Ping(ctx); err != nil {
+		log.Printf("failed to ping: %s", err)
 		return
 	}
-	u.Host = fmt.Sprintf("%s:%s", u.Hostname(), "xxx")
-	fmt.Println(u.String())
 
 	// Output:
 	// true
-	// postgres://root@localhost:xxx/defaultdb?sslmode=disable
 }
 
-func ExampleRun_withRecommendedSettings() {
+func ExampleRun_withCustomStatements() {
 	ctx := context.Background()
 
 	cockroachdbContainer, err := cockroachdb.Run(ctx, "cockroachdb/cockroach:latest-v23.1", cockroachdb.WithStatements(cockroachdb.DefaultStatements...))
diff --git a/modules/cockroachdb/options.go b/modules/cockroachdb/options.go
index eba101834e..81a1843f71 100644
--- a/modules/cockroachdb/options.go
+++ b/modules/cockroachdb/options.go
@@ -1,6 +1,17 @@
 package cockroachdb
 
-import "github.com/testcontainers/testcontainers-go"
+import (
+	"context"
+	"fmt"
+	"net"
+	"net/url"
+
+	"github.com/docker/go-connections/nat"
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/stdlib"
+
+	"github.com/testcontainers/testcontainers-go"
+)
 
 type options struct {
 	Database   string
@@ -11,6 +22,81 @@ type options struct {
 	Statements []string
 }
 
+// containerConnConfig returns the [pgx.ConnConfig] for the given container and options.
+func (opts options) containerConnConfig(ctx context.Context, container testcontainers.Container) (*pgx.ConnConfig, error) {
+	port, err := container.MappedPort(ctx, defaultSQLPort)
+	if err != nil {
+		return nil, fmt.Errorf("mapped port: %w", err)
+	}
+
+	host, err := container.Host(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("host: %w", err)
+	}
+
+	return opts.connConfig(host, port)
+}
+
+// containerConnString returns the connection string for the given container and options.
+func (opts options) containerConnString(ctx context.Context, container testcontainers.Container) (string, error) {
+	cfg, err := opts.containerConnConfig(ctx, container)
+	if err != nil {
+		return "", fmt.Errorf("container connection config: %w", err)
+	}
+
+	return stdlib.RegisterConnConfig(cfg), nil
+}
+
+// connString returns a connection string for the given host, port and options.
+func (opts options) connString(host string, port nat.Port) (string, error) {
+	cfg, err := opts.connConfig(host, port)
+	if err != nil {
+		return "", fmt.Errorf("connection config: %w", err)
+	}
+
+	return stdlib.RegisterConnConfig(cfg), nil
+}
+
+// connConfig returns a [pgx.ConnConfig] for the given host, port and options.
+func (opts options) connConfig(host string, port nat.Port) (*pgx.ConnConfig, error) {
+	user := url.User(opts.User)
+	if opts.Password != "" {
+		user = url.UserPassword(opts.User, opts.Password)
+	}
+
+	sslMode := "disable"
+	if opts.TLS != nil {
+		sslMode = "require" // We can't use "verify-full" as it might be a self signed cert.
+	}
+	params := url.Values{
+		"sslmode": []string{sslMode},
+	}
+
+	u := url.URL{
+		Scheme:   "postgres",
+		User:     user,
+		Host:     net.JoinHostPort(host, port.Port()),
+		Path:     opts.Database,
+		RawQuery: params.Encode(),
+	}
+
+	cfg, err := pgx.ParseConfig(u.String())
+	if err != nil {
+		return nil, fmt.Errorf("parse config: %w", err)
+	}
+
+	if opts.TLS != nil {
+		tlsCfg, err := opts.TLS.tlsConfig()
+		if err != nil {
+			return nil, fmt.Errorf("tls config: %w", err)
+		}
+
+		cfg.TLSConfig = tlsCfg
+	}
+
+	return cfg, nil
+}
+
 func defaultOptions() options {
 	return options{
 		User:       defaultUser,
@@ -85,8 +171,8 @@ var DefaultStatements = []string{
 }
 
 // WithStatements sets the statements to run on the CockroachDB cluster once the container is ready.
-// This, in combination with DefaultStatements, can be used to configure the cluster with the settings
-// recommended by Cockroach Labs.
+// By default, the container will run the statements in [DefaultStatements] as recommended by
+// Cockroach Labs however that is not always possible due to the user not having the required privileges.
 func WithStatements(statements ...string) Option {
 	return func(o *options) {
 		o.Statements = statements