fix(setup): avoid hardcoded postgres connection

backend/internal/setup/setup.go

@@ -6,6 +6,7 @@ import (
 	"crypto/tls"
 	"database/sql"
 	"encoding/hex"
+	"errors"
 	"fmt"
 	"os"
 	"strconv"
@@ -17,7 +18,7 @@ import (
 	"github.com/Wei-Shaw/sub2api/internal/repository"
 	"github.com/Wei-Shaw/sub2api/internal/service"
 
-	_ "github.com/lib/pq"
+	"github.com/lib/pq"
 	"github.com/redis/go-redis/v9"
 	"go.uber.org/zap"
 	"gopkg.in/yaml.v3"
@@ -30,6 +31,8 @@ const (
 	defaultUserConcurrency     = 5
 	simpleModeAdminConcurrency = 30
 	defaultMigrationTimeout    = 60 * time.Second
+	postgresBootstrapDatabase  = "postgres"
+	databasePingTimeout        = 5 * time.Second
 )
 
 func setupDefaultAdminConcurrency() int {
@@ -185,39 +188,63 @@ func buildPostgresDSN(cfg *DatabaseConfig, dbName string) string {
 	)
 }
 
-func buildDatabaseConnectionDSNs(cfg *DatabaseConfig) (bootstrapDSN, targetDSN string) {
-	return buildPostgresDSN(cfg, "postgres"), buildPostgresDSN(cfg, cfg.DBName)
+func isDatabaseNotFoundError(err error) bool {
+	var pqErr *pq.Error
+	return errors.As(err, &pqErr) && pqErr.Code == "3D000"
 }
 
+func openAndPingPostgresDatabase(cfg *DatabaseConfig, dbName string) (*sql.DB, error) {
+	db, err := sql.Open("postgres", buildPostgresDSN(cfg, dbName))
+	if err != nil {
+		return nil, err
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), databasePingTimeout)
+	err = db.PingContext(ctx)
+	cancel()
+	if err != nil {
+		_ = db.Close()
+		return nil, err
+	}
+
+	return db, nil
+}
+
+type postgresDatabaseOpener func(*DatabaseConfig, string) (*sql.DB, error)
+
 // TestDatabaseConnection tests the database connection and creates database if not exists
 func TestDatabaseConnection(cfg *DatabaseConfig) error {
-	// First, connect to the default 'postgres' database to check/create target database.
-	// Connecting to cfg.DBName here fails when the target database has not been
-	// created yet, so the bootstrap connection must use PostgreSQL's maintenance DB.
-	defaultDSN, targetDSN := buildDatabaseConnectionDSNs(cfg)
+	return testDatabaseConnection(cfg, openAndPingPostgresDatabase)
+}
 
-	db, err := sql.Open("postgres", defaultDSN)
-	if err != nil {
-		return fmt.Errorf("failed to connect to PostgreSQL: %w", err)
+func testDatabaseConnection(cfg *DatabaseConfig, openDatabase postgresDatabaseOpener) error {
+	// Prefer the configured database so existing installations remain
+	// independent of any server-specific maintenance database.
+	targetDB, err := openDatabase(cfg, cfg.DBName)
+	if err == nil {
+		if closeErr := targetDB.Close(); closeErr != nil {
+			logger.LegacyPrintf("setup", "failed to close postgres connection: %v", closeErr)
+		}
+		return nil
+	}
+	if !isDatabaseNotFoundError(err) {
+		return fmt.Errorf("ping target database failed: %w", err)
 	}
 
+	// Preserve the existing postgres bootstrap path only when the target is missing.
+	db, err := openDatabase(cfg, postgresBootstrapDatabase)
+	if err != nil {
+		return fmt.Errorf("target database '%s' does not exist; failed to connect to bootstrap database '%s': %w", cfg.DBName, postgresBootstrapDatabase, err)
+	}
 	defer func() {
-		if db == nil {
-			return
-		}
 		if err := db.Close(); err != nil {
-			logger.LegacyPrintf("setup", "failed to close postgres connection: %v", err)
+			logger.LegacyPrintf("setup", "failed to close %s connection: %v", postgresBootstrapDatabase, err)
 		}
 	}()
 
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), databasePingTimeout)
 	defer cancel()
 
-	if err := db.PingContext(ctx); err != nil {
-		return fmt.Errorf("ping failed: %w", err)
-	}
-
-	// Check if target database exists
 	var exists bool
 	row := db.QueryRowContext(ctx, "SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname = $1)", cfg.DBName)
 	if err := row.Scan(&exists); err != nil {
@@ -236,30 +263,17 @@ func TestDatabaseConnection(cfg *DatabaseConfig) error {
 		logger.LegacyPrintf("setup", "Database '%s' created successfully", cfg.DBName)
 	}
 
-	// Now connect to the target database to verify
-	if err := db.Close(); err != nil {
-		logger.LegacyPrintf("setup", "failed to close postgres connection: %v", err)
-	}
-	db = nil
-
-	targetDB, err := sql.Open("postgres", targetDSN)
+	// Now connect to the target database to verify.
+	targetDB, err = openDatabase(cfg, cfg.DBName)
 	if err != nil {
 		return fmt.Errorf("failed to connect to database '%s': %w", cfg.DBName, err)
 	}
-
 	defer func() {
 		if err := targetDB.Close(); err != nil {
 			logger.LegacyPrintf("setup", "failed to close postgres connection: %v", err)
 		}
 	}()
 
-	ctx2, cancel2 := context.WithTimeout(context.Background(), 5*time.Second)
-	defer cancel2()
-
-	if err := targetDB.PingContext(ctx2); err != nil {
-		return fmt.Errorf("ping target database failed: %w", err)
-	}
-
 	return nil
 }
 

backend/internal/setup/setup_test.go

@@ -1,11 +1,16 @@
 package setup
 
 import (
+	"database/sql"
+	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
 	"testing"
 	"time"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/lib/pq"
 )
 
 func TestDecideAdminBootstrap(t *testing.T) {
@@ -226,7 +231,7 @@ func TestWriteConfigFileIncludesRedisUsername(t *testing.T) {
 	}
 }
 
-func TestBuildDatabaseConnectionDSNsUsesPostgresForBootstrap(t *testing.T) {
+func TestDatabaseConnectionDSNsUseConfiguredTargetAndLegacyBootstrapDatabase(t *testing.T) {
 	cfg := &DatabaseConfig{
 		Host:     "db",
 		Port:     5432,
@@ -236,15 +241,117 @@ func TestBuildDatabaseConnectionDSNsUsesPostgresForBootstrap(t *testing.T) {
 		SSLMode:  "disable",
 	}
 
-	bootstrapDSN, targetDSN := buildDatabaseConnectionDSNs(cfg)
+	targetDSN := buildPostgresDSN(cfg, cfg.DBName)
+	bootstrapDSN := buildPostgresDSN(cfg, postgresBootstrapDatabase)
 
-	if !strings.Contains(bootstrapDSN, "dbname=postgres") {
-		t.Fatalf("bootstrap DSN = %q, want default postgres database", bootstrapDSN)
-	}
-	if strings.Contains(bootstrapDSN, "dbname=sub2api") {
-		t.Fatalf("bootstrap DSN = %q, should not connect to target database before checking/creating it", bootstrapDSN)
-	}
 	if !strings.Contains(targetDSN, "dbname=sub2api") {
 		t.Fatalf("target DSN = %q, want configured database", targetDSN)
 	}
+	if !strings.Contains(bootstrapDSN, "dbname=postgres") {
+		t.Fatalf("bootstrap DSN = %q, want legacy postgres bootstrap database", bootstrapDSN)
+	}
+}
+
+func TestIsDatabaseNotFoundError(t *testing.T) {
+	if !isDatabaseNotFoundError(&pq.Error{Code: "3D000"}) {
+		t.Fatal("isDatabaseNotFoundError() = false, want true for PostgreSQL invalid_catalog_name")
+	}
+	if isDatabaseNotFoundError(&pq.Error{Code: "28P01"}) {
+		t.Fatal("isDatabaseNotFoundError() = true, want false for invalid_password")
+	}
+	if !isDatabaseNotFoundError(fmt.Errorf("wrapped: %w", &pq.Error{Code: "3D000"})) {
+		t.Fatal("isDatabaseNotFoundError() = false, want true for wrapped PostgreSQL error")
+	}
+}
+
+func TestDatabaseConnectionUsesConfiguredTargetBeforeBootstrapDatabase(t *testing.T) {
+	cfg := &DatabaseConfig{DBName: "customdb"}
+	targetDB, _, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("sqlmock.New() error = %v", err)
+	}
+	defer targetDB.Close()
+
+	var opened []string
+	openDatabase := func(_ *DatabaseConfig, dbName string) (*sql.DB, error) {
+		opened = append(opened, dbName)
+		return targetDB, nil
+	}
+
+	if err := testDatabaseConnection(cfg, openDatabase); err != nil {
+		t.Fatalf("testDatabaseConnection() error = %v", err)
+	}
+	if len(opened) != 1 || opened[0] != cfg.DBName {
+		t.Fatalf("opened databases = %v, want only configured target %q", opened, cfg.DBName)
+	}
+}
+
+func TestDatabaseConnectionUsesLegacyBootstrapOnlyForMissingTarget(t *testing.T) {
+	cfg := &DatabaseConfig{DBName: "customdb"}
+	bootstrapDB, bootstrapMock, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("sqlmock.New() bootstrap error = %v", err)
+	}
+	defer bootstrapDB.Close()
+	targetDB, _, err := sqlmock.New()
+	if err != nil {
+		t.Fatalf("sqlmock.New() target error = %v", err)
+	}
+	defer targetDB.Close()
+
+	bootstrapMock.ExpectQuery(`SELECT EXISTS\(SELECT 1 FROM pg_database WHERE datname = \$1\)`).
+		WithArgs(cfg.DBName).
+		WillReturnRows(sqlmock.NewRows([]string{"exists"}).AddRow(false))
+	bootstrapMock.ExpectExec(`CREATE DATABASE customdb`).
+		WillReturnResult(sqlmock.NewResult(0, 1))
+
+	var opened []string
+	targetAttempts := 0
+	openDatabase := func(_ *DatabaseConfig, dbName string) (*sql.DB, error) {
+		opened = append(opened, dbName)
+		switch dbName {
+		case cfg.DBName:
+			targetAttempts++
+			if targetAttempts == 1 {
+				return nil, &pq.Error{Code: "3D000"}
+			}
+			return targetDB, nil
+		case postgresBootstrapDatabase:
+			return bootstrapDB, nil
+		default:
+			return nil, fmt.Errorf("unexpected database %q", dbName)
+		}
+	}
+
+	if err := testDatabaseConnection(cfg, openDatabase); err != nil {
+		t.Fatalf("testDatabaseConnection() error = %v", err)
+	}
+	if err := bootstrapMock.ExpectationsWereMet(); err != nil {
+		t.Fatalf("bootstrap database expectations: %v", err)
+	}
+	wantOpened := []string{cfg.DBName, postgresBootstrapDatabase, cfg.DBName}
+	if len(opened) != len(wantOpened) {
+		t.Fatalf("opened databases = %v, want %v", opened, wantOpened)
+	}
+	for i := range wantOpened {
+		if opened[i] != wantOpened[i] {
+			t.Fatalf("opened databases = %v, want %v", opened, wantOpened)
+		}
+	}
+}
+
+func TestDatabaseConnectionDoesNotFallbackForTargetAuthenticationError(t *testing.T) {
+	cfg := &DatabaseConfig{DBName: "customdb"}
+	var opened []string
+	openDatabase := func(_ *DatabaseConfig, dbName string) (*sql.DB, error) {
+		opened = append(opened, dbName)
+		return nil, &pq.Error{Code: "28P01"}
+	}
+
+	if err := testDatabaseConnection(cfg, openDatabase); err == nil {
+		t.Fatal("testDatabaseConnection() error = nil, want authentication error")
+	}
+	if len(opened) != 1 || opened[0] != cfg.DBName {
+		t.Fatalf("opened databases = %v, want only configured target %q", opened, cfg.DBName)
+	}
 }