diff --git a/modules/cockroachdb/cockroachdb.go b/modules/cockroachdb/cockroachdb.go index d9cfb9cb36..38f0eae839 100644 --- a/modules/cockroachdb/cockroachdb.go +++ b/modules/cockroachdb/cockroachdb.go @@ -93,7 +93,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) }, }, }, @@ -239,9 +239,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) @@ -258,18 +261,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 + } +}