diff --git a/backend/internal/setup/setup.go b/backend/internal/setup/setup.go index a1d528927e..c10193faba 100644 --- a/backend/internal/setup/setup.go +++ b/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 } diff --git a/backend/internal/setup/setup_test.go b/backend/internal/setup/setup_test.go index 52f95e38c0..65d14f57a8 100644 --- a/backend/internal/setup/setup_test.go +++ b/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) + } }