feat!(cockroachdb): use recommended cluster settings

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
+}

modules/cockroachdb/cockroachdb.go

@@ -3,22 +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 (
@@ -39,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 {
@@ -48,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...)
@@ -91,6 +97,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)
+						},
+					},
 				},
 			},
 		},
@@ -172,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),
@@ -234,47 +228,33 @@ func addTLS(ctx context.Context, container testcontainers.Container, opts option
 	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(),
+// 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
 	}
 
-	return u.String()
-}
-
-func connTLS(opts options) (*tls.Config, error) {
-	if opts.TLS == nil {
-		return nil, ErrTLSNotEnabled
+	connStr, err := opts.containerConnString(ctx, container)
+	if err != nil {
+		return fmt.Errorf("connection string: %w", err)
 	}
 
-	keyPair, err := tls.X509KeyPair(opts.TLS.ClientCert, opts.TLS.ClientKey)
+	db, err := sql.Open("pgx/v5", connStr)
 	if err != nil {
-		return nil, err
+		return fmt.Errorf("sql.Open: %w", err)
 	}
+	defer func() {
+		cerr := db.Close()
+		if err == nil {
+			err = cerr
+		}
+	}()
 
-	certPool := x509.NewCertPool()
-	certPool.AddCert(opts.TLS.CACert)
+	for _, stmt := range opts.Statements {
+		if _, err = db.Exec(stmt); err != nil {
+			return fmt.Errorf("db.Exec: %w", err)
+		}
+	}
 
-	return &tls.Config{
-		RootCAs:      certPool,
-		Certificates: []tls.Certificate{keyPair},
-		ServerName:   "localhost",
-	}, nil
+	return nil
 }

modules/cockroachdb/cockroachdb_test.go

@@ -2,9 +2,6 @@ package cockroachdb_test
 
 import (
 	"context"
-	"errors"
-	"net/url"
-	"strings"
 	"testing"
 	"time"
 
@@ -18,26 +15,28 @@ 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
+			// lacking the needed MODIFYCLUSTERSETTING privilege to run them.
+			cockroachdb.WithStatements(),
 		},
 	})
 }
 
 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"),
+			// Do not run the default statements as the user used on this test is
+			// lacking the needed MODIFYCLUSTERSETTING privilege to run them.
+			cockroachdb.WithStatements(),
 		},
 	})
 }
@@ -47,16 +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),
 		},
 	})
 }
 
+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
 }
 
@@ -66,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() {
@@ -194,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
-}