feat: add WithStatements option and expose recommended settings as a …

…list of statements
testcontainers · Oct 30, 2024 · fe3d3ee · fe3d3ee
1 parent 1fe0c72
commit fe3d3ee
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 21 deletions.
diff --git 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
@@ -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
+	}
+}