diff --git a/modules/cockroachdb/cockroachdb.go b/modules/cockroachdb/cockroachdb.go index 4d412f04d3..a3b4952548 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" "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 != "" { 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/examples_test.go b/modules/cockroachdb/examples_test.go index c06c97596b..9a8fb12881 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.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 +} diff --git a/modules/cockroachdb/options.go b/modules/cockroachdb/options.go index a2211d77e7..eba101834e 100644 --- a/modules/cockroachdb/options.go +++ b/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 + } +}