diff --git a/cmd/cleanup-lambda/main.go b/cmd/cleanup-lambda/main.go index 327a165c7..ed3cc8309 100644 --- a/cmd/cleanup-lambda/main.go +++ b/cmd/cleanup-lambda/main.go @@ -7,7 +7,6 @@ import ( "time" "github.com/LeanerCloud/CUDly/internal/database" - "github.com/LeanerCloud/CUDly/internal/secrets" "github.com/aws/aws-lambda-go/lambda" ) @@ -30,7 +29,7 @@ func cleanupExpiredRecords(ctx context.Context, event CleanupEvent) (*CleanupRes // A new DB connection is opened per invocation (no connection reuse across warm starts). // This is intentional: the cleanup Lambda runs infrequently and the simpler, stateless // design is preferred over the shared-connection pattern used in cmd/lambda/main.go. - db, err := initDB(ctx) + db, err := database.OpenFromEnv(ctx) if err != nil { return nil, err } @@ -56,31 +55,6 @@ func cleanupExpiredRecords(ctx context.Context, event CleanupEvent) (*CleanupRes return result, nil } -// initDB creates and returns a database connection using env config and optional secret resolution. -func initDB(ctx context.Context) (*database.Connection, error) { - dbConfig, err := database.LoadFromEnv() - if err != nil { - return nil, fmt.Errorf("failed to load database config: %w", err) - } - - var secretResolver database.SecretResolver - if dbConfig.PasswordSecret != "" { - secretConfig := secrets.LoadConfigFromEnv() - resolver, err := secrets.NewResolver(ctx, secretConfig) - if err != nil { - return nil, fmt.Errorf("failed to create secret resolver: %w", err) - } - defer resolver.Close() - secretResolver = resolver - } - - db, err := database.NewConnection(ctx, dbConfig, secretResolver) - if err != nil { - return nil, fmt.Errorf("failed to connect to database: %w", err) - } - return db, nil -} - // dryRunCount counts records that would be deleted without actually deleting them. func dryRunCount(ctx context.Context, db *database.Connection, now time.Time, result *CleanupResult) error { if err := db.QueryRow(ctx, "SELECT COUNT(*) FROM sessions WHERE expires_at < $1", now).Scan(&result.SessionsDeleted); err != nil { diff --git a/cmd/lambda/clear-rate-limit/main.go b/cmd/lambda/clear-rate-limit/main.go index 39f385ea1..2f654a571 100644 --- a/cmd/lambda/clear-rate-limit/main.go +++ b/cmd/lambda/clear-rate-limit/main.go @@ -7,7 +7,6 @@ import ( "os" "github.com/LeanerCloud/CUDly/internal/database" - "github.com/LeanerCloud/CUDly/internal/secrets" "github.com/aws/aws-lambda-go/lambda" ) @@ -27,27 +26,9 @@ type Response struct { } func clearRateLimit(ctx context.Context) (Response, error) { - // Initialize database connection with secret resolution - dbConfig, err := database.LoadFromEnv() + db, err := database.OpenFromEnv(ctx) if err != nil { - return Response{}, fmt.Errorf("failed to load database config: %w", err) - } - - // Create secret resolver if password secret is specified - var secretResolver database.SecretResolver - if dbConfig.PasswordSecret != "" { - secretConfig := secrets.LoadConfigFromEnv() - resolver, err := secrets.NewResolver(ctx, secretConfig) - if err != nil { - return Response{}, fmt.Errorf("failed to create secret resolver: %w", err) - } - defer resolver.Close() - secretResolver = resolver - } - - db, err := database.NewConnection(ctx, dbConfig, secretResolver) - if err != nil { - return Response{}, fmt.Errorf("failed to connect to database: %w", err) + return Response{}, err } defer db.Close() diff --git a/cmd/rekey/main.go b/cmd/rekey/main.go index 46e3a9974..ed61eefea 100644 --- a/cmd/rekey/main.go +++ b/cmd/rekey/main.go @@ -36,7 +36,7 @@ func main() { flag.Parse() if os.Getenv(safetyEnv) != "1" { - log.Fatalf("rekey: refusing to run without %s=1 (see docs/runbooks/rekey-from-zero-key.md)", safetyEnv) + log.Fatalf("rekey: refusing to run without %s=1 (see cmd/rekey/README.md)", safetyEnv) } ctx, cancel := context.WithTimeout(context.Background(), *timeout) diff --git a/internal/database/open_from_env.go b/internal/database/open_from_env.go new file mode 100644 index 000000000..aa78ffc33 --- /dev/null +++ b/internal/database/open_from_env.go @@ -0,0 +1,51 @@ +package database + +import ( + "context" + "fmt" + + "github.com/LeanerCloud/CUDly/internal/secrets" +) + +// OpenFromEnv creates a database connection using environment-variable configuration. +// +// It loads the DB config from the environment, optionally builds a secrets.Resolver +// when DB_PASSWORD_SECRET is set (resolving the password synchronously before +// returning), then returns an open connection pool. The resolver is closed as +// soon as the password has been resolved; callers do not need to manage it. +// +// This is the canonical bootstrap helper for cmd/* entrypoints that only need +// a DB connection and have no other use for the secrets resolver. If you also +// need the resolver for a different purpose (e.g. cmd/rekey loads a key via +// credentials.LoadKey before opening the DB), wire the resolver yourself and +// call NewConnection directly. +func OpenFromEnv(ctx context.Context) (*Connection, error) { + dbConfig, err := LoadFromEnv() + if err != nil { + return nil, fmt.Errorf("database: load config from env: %w", err) + } + + var sr SecretResolver + if dbConfig.PasswordSecret != "" { + secretConfig := secrets.LoadConfigFromEnv() + resolver, err := secrets.NewResolver(ctx, secretConfig) + if err != nil { + return nil, fmt.Errorf("database: create secret resolver: %w", err) + } + // NewConnection resolves the password synchronously before returning, so + // the resolver is safe to close immediately after the connection is open. + defer func() { + if cerr := resolver.Close(); cerr != nil { + // Non-fatal: the connection is already established. + _ = cerr + } + }() + sr = resolver + } + + conn, err := NewConnection(ctx, dbConfig, sr) + if err != nil { + return nil, fmt.Errorf("database: connect: %w", err) + } + return conn, nil +} diff --git a/internal/database/open_from_env_test.go b/internal/database/open_from_env_test.go new file mode 100644 index 000000000..fad9064a3 --- /dev/null +++ b/internal/database/open_from_env_test.go @@ -0,0 +1,106 @@ +package database + +import ( + "context" + "os" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestOpenFromEnv_MissingConfig verifies that OpenFromEnv surfaces a config +// validation error (e.g. missing DB_PASSWORD) rather than panicking or +// returning a nil connection. +// +// A live Postgres is not available in unit tests; we cover the config-loading +// and error-propagation paths here. The connection-establishment path is +// exercised by integration tests against testcontainers Postgres in CI. +func TestOpenFromEnv_MissingConfig(t *testing.T) { + // Remove all DB env vars so LoadFromEnv returns an error. + keys := []string{ + "DB_HOST", "DB_PORT", "DB_NAME", "DB_USER", + "DB_PASSWORD", "DB_PASSWORD_SECRET", "DB_SSL_MODE", + } + originals := make(map[string]string, len(keys)) + for _, k := range keys { + originals[k] = os.Getenv(k) + os.Unsetenv(k) + } + t.Cleanup(func() { + for k, v := range originals { + if v != "" { + os.Setenv(k, v) + } else { + os.Unsetenv(k) + } + } + }) + + conn, err := OpenFromEnv(context.Background()) + require.Error(t, err, "OpenFromEnv must fail when DB config is missing") + assert.Nil(t, conn) + // Error must be wrapped with the database: prefix so callers can tell + // it came from the bootstrap layer. + assert.Contains(t, err.Error(), "database:") +} + +// TestOpenFromEnv_NoPasswordSecret verifies that OpenFromEnv does not attempt +// to build a secrets.Resolver when DB_PASSWORD_SECRET is absent, even if +// DB_PASSWORD is provided (the resolver path is unreachable). +// +// This test cannot make a real connection, so it expects a connection error, +// not a config error -- the config is valid, but there is no Postgres to +// connect to. +func TestOpenFromEnv_NoPasswordSecret(t *testing.T) { + // Provide a structurally valid config that will fail at the TCP dial. + env := map[string]string{ + "DB_HOST": "127.0.0.1", + "DB_PORT": "19999", // nothing listening here + "DB_NAME": "testdb", + "DB_USER": "testuser", + "DB_PASSWORD": "testpass", + "DB_SSL_MODE": "disable", + "DB_MAX_CONN_IDLE_TIME": "1s", + "DB_MAX_CONN_LIFETIME": "1s", + "DB_HEALTH_CHECK_PERIOD": "1s", + } + originals := make(map[string]string, len(env)) + for k, v := range env { + originals[k] = os.Getenv(k) + os.Setenv(k, v) + } + os.Unsetenv("DB_PASSWORD_SECRET") + t.Cleanup(func() { + for k, v := range originals { + if v != "" { + os.Setenv(k, v) + } else { + os.Unsetenv(k) + } + } + os.Unsetenv("DB_PASSWORD_SECRET") + }) + + // Shrink retry knobs so the test finishes quickly. + origMaxRetries := maxConnectRetries + origBase := connectRetryBaseDelay + origMax := connectRetryMaxDelay + origPerAttempt := perAttemptConnectTimeout + maxConnectRetries = 1 + connectRetryBaseDelay = 0 + connectRetryMaxDelay = 0 + perAttemptConnectTimeout = 50e6 // 50ms + t.Cleanup(func() { + maxConnectRetries = origMaxRetries + connectRetryBaseDelay = origBase + connectRetryMaxDelay = origMax + perAttemptConnectTimeout = origPerAttempt + }) + + conn, err := OpenFromEnv(context.Background()) + // We expect a connection error (no Postgres at 19999), not a config error. + require.Error(t, err, "OpenFromEnv must fail when Postgres is unreachable") + assert.Nil(t, conn) + assert.Contains(t, err.Error(), "database:") +}