galaxy-game/integration/internal/harness/postgres_container.go

package harness

import (
	"context"
	"fmt"
	"net"
	"net/url"
	"strings"
	"sync"
	"testing"
	"time"

	"galaxy/postgres"

	testcontainers "github.com/testcontainers/testcontainers-go"
	tcpostgres "github.com/testcontainers/testcontainers-go/modules/postgres"
	"github.com/testcontainers/testcontainers-go/wait"
)

const (
	defaultPostgresContainerImage = "postgres:16-alpine"
	defaultPostgresDatabase       = "galaxy_integration"
	defaultPostgresSuperuser      = "galaxy_integration"
	defaultPostgresSuperPassword  = "galaxy_integration"

	postgresAdminConnectTimeout = 5 * time.Second
	postgresStartupTimeout      = 60 * time.Second
)

// PostgresRuntime stores one started real PostgreSQL container together with
// the parsed connection coordinates and the per-test role credentials issued
// by EnsureRoleAndSchema.
//
// The struct is safe to call from concurrent tests because credential lookups
// guard the internal map with a mutex; each test should still keep its own
// PostgresRuntime to preserve container-level isolation.
type PostgresRuntime struct {
	Container *tcpostgres.PostgresContainer

	baseDSN  string
	host     string
	port     string
	database string

	mu    sync.Mutex
	creds map[string]string
}

// StartPostgresContainer starts one isolated PostgreSQL container and registers
// automatic cleanup for the suite. The container exposes a superuser created
// from the package-level constants; per-service roles are issued lazily by
// EnsureRoleAndSchema.
func StartPostgresContainer(t testing.TB) *PostgresRuntime {
	t.Helper()

	ctx := context.Background()

	container, err := tcpostgres.Run(ctx,
		defaultPostgresContainerImage,
		tcpostgres.WithDatabase(defaultPostgresDatabase),
		tcpostgres.WithUsername(defaultPostgresSuperuser),
		tcpostgres.WithPassword(defaultPostgresSuperPassword),
		// The default Postgres image emits the "ready to accept connections"
		// log line twice during startup: once during temporary bootstrap, once
		// after the real listener opens on the mapped port. Waiting for the
		// second occurrence avoids racing the temporary instance.
		testcontainers.WithWaitStrategy(
			wait.ForLog("database system is ready to accept connections").
				WithOccurrence(2).
				WithStartupTimeout(postgresStartupTimeout),
		),
	)
	if err != nil {
		t.Fatalf("start postgres container: %v", err)
	}

	t.Cleanup(func() {
		if err := testcontainers.TerminateContainer(container); err != nil {
			t.Errorf("terminate postgres container: %v", err)
		}
	})

	baseDSN, err := container.ConnectionString(ctx, "sslmode=disable")
	if err != nil {
		t.Fatalf("resolve postgres connection string: %v", err)
	}

	host, port, err := splitHostPort(baseDSN)
	if err != nil {
		t.Fatalf("parse postgres connection string: %v", err)
	}

	return &PostgresRuntime{
		Container: container,
		baseDSN:   baseDSN,
		host:      host,
		port:      port,
		database:  defaultPostgresDatabase,
		creds:     map[string]string{},
	}
}

// BaseDSN returns the superuser DSN exposed by the container, suitable for
// administrative tasks such as creating roles or schemas. Callers should
// prefer DSNForSchema for service-scoped access.
func (rt *PostgresRuntime) BaseDSN() string {
	return rt.baseDSN
}

// DSNForSchema returns a DSN that connects as role and pins search_path to
// schema. EnsureRoleAndSchema must have populated credentials for role first;
// otherwise the call panics, signalling a test setup bug.
func (rt *PostgresRuntime) DSNForSchema(schema, role string) string {
	rt.mu.Lock()
	password, ok := rt.creds[role]
	rt.mu.Unlock()
	if !ok {
		panic(fmt.Sprintf(
			"harness: DSNForSchema called for role %q with no credentials; call EnsureRoleAndSchema first",
			role,
		))
	}

	values := url.Values{}
	values.Set("search_path", schema)
	values.Set("sslmode", "disable")

	dsn := url.URL{
		Scheme:   "postgres",
		User:     url.UserPassword(role, password),
		Host:     net.JoinHostPort(rt.host, rt.port),
		Path:     "/" + rt.database,
		RawQuery: values.Encode(),
	}
	return dsn.String()
}

// EnsureRoleAndSchema creates role with the given password (idempotent) and a
// schema owned by that role (idempotent), then grants USAGE so the role can
// resolve table references inside it. The credentials are cached for later
// DSNForSchema lookups.
//
// The operation runs through a temporary administrative connection opened
// from BaseDSN; the connection is closed before the call returns.
func (rt *PostgresRuntime) EnsureRoleAndSchema(ctx context.Context, schema, role, password string) error {
	if strings.TrimSpace(schema) == "" {
		return fmt.Errorf("ensure role and schema: schema must not be empty")
	}
	if strings.TrimSpace(role) == "" {
		return fmt.Errorf("ensure role and schema: role must not be empty")
	}

	cfg := postgres.DefaultConfig()
	cfg.PrimaryDSN = rt.baseDSN
	cfg.OperationTimeout = postgresAdminConnectTimeout

	db, err := postgres.OpenPrimary(ctx, cfg)
	if err != nil {
		return fmt.Errorf("ensure role and schema: open admin connection: %w", err)
	}
	defer func() {
		_ = db.Close()
	}()

	createRole := fmt.Sprintf(`DO $$
BEGIN
	IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = %s) THEN
		CREATE ROLE %s LOGIN PASSWORD %s;
	END IF;
END $$;`,
		quoteSQLLiteral(role),
		quoteSQLIdentifier(role),
		quoteSQLLiteral(password),
	)
	if _, err := db.ExecContext(ctx, createRole); err != nil {
		return fmt.Errorf("ensure role and schema: create role %q: %w", role, err)
	}

	createSchema := fmt.Sprintf(`CREATE SCHEMA IF NOT EXISTS %s AUTHORIZATION %s;`,
		quoteSQLIdentifier(schema),
		quoteSQLIdentifier(role),
	)
	if _, err := db.ExecContext(ctx, createSchema); err != nil {
		return fmt.Errorf("ensure role and schema: create schema %q: %w", schema, err)
	}

	grantUsage := fmt.Sprintf(`GRANT USAGE ON SCHEMA %s TO %s;`,
		quoteSQLIdentifier(schema),
		quoteSQLIdentifier(role),
	)
	if _, err := db.ExecContext(ctx, grantUsage); err != nil {
		return fmt.Errorf("ensure role and schema: grant usage on %q to %q: %w", schema, role, err)
	}

	rt.mu.Lock()
	rt.creds[role] = password
	rt.mu.Unlock()

	return nil
}

// WithPostgres returns env entries pointing the service identified by
// envPrefix at schema/role inside rt. EnsureRoleAndSchema must have populated
// credentials for role first.
//
// The returned map carries only `<envPrefix>_POSTGRES_PRIMARY_DSN`; the other
// per-service Postgres knobs (operation timeout, pool sizes) keep the
// defaults provided by `pkg/postgres.DefaultConfig`.
func WithPostgres(rt *PostgresRuntime, envPrefix, schema, role string) map[string]string {
	return map[string]string{
		envPrefix + "_POSTGRES_PRIMARY_DSN": rt.DSNForSchema(schema, role),
	}
}

// quoteSQLIdentifier wraps name in double quotes and escapes any embedded
// double quote, producing a SQL identifier that survives reserved words such
// as `user`.
func quoteSQLIdentifier(name string) string {
	return `"` + strings.ReplaceAll(name, `"`, `""`) + `"`
}

// quoteSQLLiteral wraps value in single quotes and escapes any embedded single
// quote, producing a SQL literal usable in DDL statements where parameter
// binding is not available.
func quoteSQLLiteral(value string) string {
	return "'" + strings.ReplaceAll(value, "'", "''") + "'"
}

// splitHostPort extracts host and port from a postgres:// DSN.
func splitHostPort(dsn string) (string, string, error) {
	parsed, err := url.Parse(dsn)
	if err != nil {
		return "", "", fmt.Errorf("parse dsn: %w", err)
	}
	host := parsed.Hostname()
	port := parsed.Port()
	if host == "" || port == "" {
		return "", "", fmt.Errorf("dsn %q missing host or port", dsn)
	}
	return host, port, nil
}