feat: use recommended cluster settings for cockroachdb

modules/cockroachdb/cockroachdb.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/tls"
 	"crypto/x509"
+	"database/sql"
 	"encoding/pem"
 	"fmt"
 	"net"
@@ -90,6 +91,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 runStatements(ctx, container, o)
+						},
+					},
 				},
 			},
 		},
@@ -233,6 +239,42 @@ func addTLS(ctx context.Context, container testcontainers.Container, opts option
 	return nil
 }
 
+// runStatements runs the configured statements against the CockroachDB container.
+func runStatements(ctx context.Context, container testcontainers.Container, opts options) (err error) {
+	if len(opts.Statements) == 0 {
+		return nil
+	}
+
+	port, err := container.MappedPort(ctx, defaultSQLPort)
+	if err != nil {
+		return fmt.Errorf("mapped port: %w", err)
+	}
+
+	host, err := container.Host(ctx)
+	if err != nil {
+		return fmt.Errorf("host: %w", err)
+	}
+
+	db, err := sql.Open("pgx/v5", connString(opts, host, port))
+	if err != nil {
+		return fmt.Errorf("sql.Open: %w", err)
+	}
+	defer func() {
+		cerr := db.Close()
+		if err == nil {
+			err = cerr
+		}
+	}()
+
+	for _, stmt := range opts.Statements {
+		if _, err = db.Exec(stmt); err != nil {
+			return fmt.Errorf("db.Exec: %w", err)
+		}
+	}
+
+	return nil
+}
+
 func connString(opts options, host string, port nat.Port) string {
 	user := url.User(opts.User)
 	if opts.Password != "" {

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(),
 		},
 	})
 }

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.DefaultStatements...))
+	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
+}

modules/cockroachdb/options.go

@@ -3,19 +3,21 @@ 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 {
 	return options{
-		User:      defaultUser,
-		Password:  defaultPassword,
-		Database:  defaultDatabase,
-		StoreSize: defaultStoreSize,
+		User:       defaultUser,
+		Password:   defaultPassword,
+		Database:   defaultDatabase,
+		StoreSize:  defaultStoreSize,
+		Statements: DefaultStatements,
 	}
 }
 
@@ -67,3 +69,26 @@ func WithTLS(cfg *TLSConfig) Option {
 		o.TLS = cfg
 	}
 }
+
+// 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'",
+	"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.
+// 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) {
+		o.Statements = statements
+	}
+}