From 81290cb08304a3e8192e056568c12d95e5c66340 Mon Sep 17 00:00:00 2001 From: Steven Hartland Date: Wed, 27 Nov 2024 14:05:03 +0000 Subject: [PATCH] chore: various clean ups --- modules/cockroachdb/cockroachdb.go | 177 +++++++++++++++-------------- modules/cockroachdb/options.go | 9 +- 2 files changed, 93 insertions(+), 93 deletions(-) diff --git a/modules/cockroachdb/cockroachdb.go b/modules/cockroachdb/cockroachdb.go index a1130f2fc1..966214f5a0 100644 --- a/modules/cockroachdb/cockroachdb.go +++ b/modules/cockroachdb/cockroachdb.go @@ -126,7 +126,7 @@ func (c *CockroachDBContainer) ConnectionConfig(ctx context.Context) (*pgx.ConnC // TLSConfig returns config necessary to connect to CockroachDB over TLS. // Returns [ErrTLSNotEnabled] if TLS is not enabled. // -// Deprecated: use [CockroachDBContainer.ConnectionConfig] or +// Deprecated: use [CockroachDBContainer.ConnectionString] or // [CockroachDBContainer.ConnectionConfig] instead. func (c *CockroachDBContainer) TLSConfig() (*tls.Config, error) { if cfg := c.tlsStrategy.TLSConfig(); cfg != nil { @@ -136,89 +136,6 @@ func (c *CockroachDBContainer) TLSConfig() (*tls.Config, error) { return nil, ErrTLSNotEnabled } -// connString returns a connection string for the given host, port and options. -func (c *CockroachDBContainer) connString(host string, port nat.Port) (string, error) { - cfg, err := c.connConfig(host, port) - if err != nil { - return "", fmt.Errorf("connection config: %w", err) - } - - return stdlib.RegisterConnConfig(cfg), nil -} - -// connConfig returns a [pgx.ConnConfig] for the given host, port and options. -func (c *CockroachDBContainer) connConfig(host string, port nat.Port) (*pgx.ConnConfig, error) { - var user *url.Userinfo - if c.password != "" { - user = url.UserPassword(c.user, c.password) - } else { - user = url.User(c.user) - } - - sslMode := "disable" - tlsConfig := c.tlsStrategy.TLSConfig() - if tlsConfig != nil { - sslMode = "verify-full" - } - params := url.Values{ - "sslmode": []string{sslMode}, - } - - u := url.URL{ - Scheme: "postgres", - User: user, - Host: net.JoinHostPort(host, port.Port()), - Path: c.database, - RawQuery: params.Encode(), - } - - cfg, err := pgx.ParseConfig(u.String()) - if err != nil { - return nil, fmt.Errorf("parse config: %w", err) - } - - cfg.TLSConfig = tlsConfig - - return cfg, nil -} - -// setOptions sets the CockroachDBContainer options from a request. -func (c *CockroachDBContainer) setOptions(req *testcontainers.GenericContainerRequest) error { - c.database = req.Env[envDatabase] - c.user = req.Env[envUser] - c.password = req.Env[envPassword] - - var insecure bool - for _, arg := range req.Cmd { - if arg == insecureFlag { - insecure = true - break - } - } - - if err := wait.Walk(&req.WaitingFor, func(strategy wait.Strategy) error { - if cert, ok := strategy.(*wait.TLSStrategy); ok { - if insecure { - // If insecure mode is enabled, the certificate strategy is removed. - return errors.Join(wait.VisitRemove, wait.VisitStop) - } - - // Update the client certificate files to match the user which may have changed. - cert.WithCert(certsDir+"/client."+c.user+".crt", certsDir+"/client."+c.user+".key") - - c.tlsStrategy = cert - - // Stop the walk as the certificate strategy has been found. - return wait.VisitStop - } - return nil - }); err != nil { - return fmt.Errorf("walk strategies: %w", err) - } - - return nil -} - // Deprecated: use Run instead. // RunContainer creates an instance of the CockroachDB container type func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomizer) (*CockroachDBContainer, error) { @@ -227,7 +144,7 @@ func RunContainer(ctx context.Context, opts ...testcontainers.ContainerCustomize // Run start an instance of the CockroachDB container type using the given image and options. // -// By default, the container will: +// By default, the container will configured with: // - Cluster: Single node // - Storage: 100% in-memory // - User: root @@ -292,8 +209,7 @@ func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustom } } - // Extract the options from the request so they can used by wait strategies and connection methods. - if err := ctr.setOptions(&req); err != nil { + if err := ctr.configure(&req); err != nil { return nil, fmt.Errorf("set options: %w", err) } @@ -305,3 +221,90 @@ func Run(ctx context.Context, img string, opts ...testcontainers.ContainerCustom return ctr, nil } + +// connString returns a connection string for the given host, port and options. +func (c *CockroachDBContainer) connString(host string, port nat.Port) (string, error) { + cfg, err := c.connConfig(host, port) + if err != nil { + return "", fmt.Errorf("connection config: %w", err) + } + + return stdlib.RegisterConnConfig(cfg), nil +} + +// connConfig returns a [pgx.ConnConfig] for the given host, port and options. +func (c *CockroachDBContainer) connConfig(host string, port nat.Port) (*pgx.ConnConfig, error) { + var user *url.Userinfo + if c.password != "" { + user = url.UserPassword(c.user, c.password) + } else { + user = url.User(c.user) + } + + sslMode := "disable" + tlsConfig := c.tlsStrategy.TLSConfig() + if tlsConfig != nil { + sslMode = "verify-full" + } + params := url.Values{ + "sslmode": []string{sslMode}, + } + + u := url.URL{ + Scheme: "postgres", + User: user, + Host: net.JoinHostPort(host, port.Port()), + Path: c.database, + RawQuery: params.Encode(), + } + + cfg, err := pgx.ParseConfig(u.String()) + if err != nil { + return nil, fmt.Errorf("parse config: %w", err) + } + + cfg.TLSConfig = tlsConfig + + return cfg, nil +} + +// configure sets the CockroachDBContainer options from the given request and updates the request +// wait strategies to match the options. +func (c *CockroachDBContainer) configure(req *testcontainers.GenericContainerRequest) error { + c.database = req.Env[envDatabase] + c.user = req.Env[envUser] + c.password = req.Env[envPassword] + + var insecure bool + for _, arg := range req.Cmd { + if arg == insecureFlag { + insecure = true + break + } + } + + // Walk the wait strategies to find the TLS strategy and either remove it or + // update the client certificate files to match the user and configure the + // container to use the TLS strategy. + if err := wait.Walk(&req.WaitingFor, func(strategy wait.Strategy) error { + if cert, ok := strategy.(*wait.TLSStrategy); ok { + if insecure { + // If insecure mode is enabled, the certificate strategy is removed. + return errors.Join(wait.VisitRemove, wait.VisitStop) + } + + // Update the client certificate files to match the user which may have changed. + cert.WithCert(certsDir+"/client."+c.user+".crt", certsDir+"/client."+c.user+".key") + + c.tlsStrategy = cert + + // Stop the walk as the certificate strategy has been found. + return wait.VisitStop + } + return nil + }); err != nil { + return fmt.Errorf("walk strategies: %w", err) + } + + return nil +} diff --git a/modules/cockroachdb/options.go b/modules/cockroachdb/options.go index 738036d188..9efac532c6 100644 --- a/modules/cockroachdb/options.go +++ b/modules/cockroachdb/options.go @@ -16,9 +16,7 @@ var errInsecureWithPassword = errors.New("insecure mode cannot be used with a pa // The database creation will be skipped if data exists in the `/cockroach/cockroach-data` directory within the container. func WithDatabase(database string) testcontainers.CustomizeRequestOption { return func(req *testcontainers.GenericContainerRequest) error { - // CockroachDB forces the database to be lowercase. - database = strings.ToLower(database) - req.Env[envDatabase] = database + req.Env[envDatabase] = strings.ToLower(database) return nil } } @@ -28,15 +26,14 @@ func WithDatabase(database string) testcontainers.CustomizeRequestOption { // The user creation will be skipped if data exists in the `/cockroach/cockroach-data` directory within the container. func WithUser(user string) testcontainers.CustomizeRequestOption { return func(req *testcontainers.GenericContainerRequest) error { - // CockroachDB forces the user to be lowercase. - user = strings.ToLower(user) - req.Env[envUser] = user + req.Env[envUser] = strings.ToLower(user) return nil } } // WithPassword sets the password of the user to create and connect as. // The user creation will be skipped if data exists in the `/cockroach/cockroach-data` directory within the container. +// This will error if insecure mode is enabled. func WithPassword(password string) testcontainers.CustomizeRequestOption { return func(req *testcontainers.GenericContainerRequest) error { for _, arg := range req.Cmd {