gochat/backend/internal/database/migrate.go

package database

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"strings"

	"github.com/golang-migrate/migrate/v4"
	_ "github.com/golang-migrate/migrate/v4/database/postgres"
	_ "github.com/golang-migrate/migrate/v4/database/sqlite3"
	_ "github.com/golang-migrate/migrate/v4/source/file"
	"gorm.io/gorm"
)

var unsupportedPQEnvironmentKeys = []string{"PGSERVICE", "PGSERVICEFILE", "PGREALM"}

// SanitizePostgresEnvironment removes libpq service-file settings that are
// incompatible with GoChat's explicit database configuration and can make
// both pgx and lib/pq ignore or reject the configured host/user/database.
func SanitizePostgresEnvironment() {
	for _, key := range unsupportedPQEnvironmentKeys {
		_ = os.Unsetenv(key)
	}
}

// RunMigrations applies all pending database migrations from the given path.
// dbURL should be a database connection string (PostgreSQL or SQLite).
// migrationsPath should be a file path to the migrations directory (e.g. "migrations").
//
// Returns nil if all migrations were applied successfully, or an error if:
// - The migration source cannot be opened
// - The database connection fails
// - A migration fails to apply
// - The database is in a dirty state (a previous migration partially failed)
func RunMigrations(dbURL string, migrationsPath string) error {
	return withSanitizedPQEnvironment(func() error {
		return runMigrations(dbURL, migrationsPath)
	})
}

func runMigrations(dbURL string, migrationsPath string) error {
	m, err := newMigrate(dbURL, migrationsPath)
	if err != nil {
		return err
	}
	defer m.Close()

	if err := m.Up(); err != nil {
		if errors.Is(err, migrate.ErrNoChange) {
			// No pending migrations — this is not an error
			return nil
		}
		// Check for dirty state
		version, dirty, dirtyErr := m.Version()
		if dirtyErr == nil && dirty {
			return fmt.Errorf("database is in dirty state at version %d; run 'migrate force %d' to fix before retrying: %w", version, version, err)
		}
		return fmt.Errorf("migration failed: %w", err)
	}

	return nil
}

// withSanitizedPQEnvironment prevents lib/pq from panicking when PostgreSQL
// service-file variables are inherited from the user's shell. GoChat passes a
// complete connection URL to golang-migrate, and lib/pq does not support these
// libpq service variables. Restore them after the migration operation so this
// workaround remains scoped to the legacy migration driver.
func withSanitizedPQEnvironment(fn func() error) error {
	type savedValue struct {
		value string
		set   bool
	}
	saved := make(map[string]savedValue, len(unsupportedPQEnvironmentKeys))
	for _, key := range unsupportedPQEnvironmentKeys {
		value, set := os.LookupEnv(key)
		saved[key] = savedValue{value: value, set: set}
		_ = os.Unsetenv(key)
	}
	defer func() {
		for _, key := range unsupportedPQEnvironmentKeys {
			previous := saved[key]
			if previous.set {
				_ = os.Setenv(key, previous.value)
			} else {
				_ = os.Unsetenv(key)
			}
		}
	}()

	return fn()
}

// MigrateSteps applies N migration steps (positive = up, negative = down).
func MigrateSteps(dbURL string, migrationsPath string, steps int) error {
	m, err := newMigrate(dbURL, migrationsPath)
	if err != nil {
		return err
	}
	defer m.Close()

	if err := m.Steps(steps); err != nil {
		if errors.Is(err, migrate.ErrNoChange) {
			return nil
		}
		return fmt.Errorf("migration steps failed: %w", err)
	}

	return nil
}

// RollbackMigrations rolls back all migrations (drops all tables).
func RollbackMigrations(dbURL string, migrationsPath string) error {
	m, err := newMigrate(dbURL, migrationsPath)
	if err != nil {
		return err
	}
	defer m.Close()

	if err := m.Down(); err != nil {
		if errors.Is(err, migrate.ErrNoChange) {
			return nil
		}
		return fmt.Errorf("rollback failed: %w", err)
	}

	return nil
}

// ForceVersion sets the migration version to a specific number, clearing dirty state.
// Use this to recover from a partially-applied migration.
func ForceVersion(dbURL string, migrationsPath string, version int) error {
	m, err := newMigrate(dbURL, migrationsPath)
	if err != nil {
		return err
	}
	defer m.Close()

	if err := m.Force(version); err != nil {
		return fmt.Errorf("force version failed: %w", err)
	}

	return nil
}

// CurrentVersion returns the current migration version and whether the DB is in a dirty state.
// When no migrations have been applied (version 0), returns (0, false, nil) instead of an error.
func CurrentVersion(dbURL string, migrationsPath string) (uint, bool, error) {
	m, err := newMigrate(dbURL, migrationsPath)
	if err != nil {
		return 0, false, err
	}
	defer m.Close()

	version, dirty, err := m.Version()
	if err != nil {
		if errors.Is(err, migrate.ErrNilVersion) {
			// No migrations applied yet — version 0, clean state
			return 0, false, nil
		}
		return 0, false, fmt.Errorf("failed to get version: %w", err)
	}

	return version, dirty, nil
}

// LatestVersion returns the highest numbered up migration on disk.
func LatestVersion(migrationsPath string) (uint, error) {
	entries, err := os.ReadDir(migrationsPath)
	if err != nil {
		return 0, fmt.Errorf("read migrations: %w", err)
	}
	var latest uint64
	for _, entry := range entries {
		name := entry.Name()
		if entry.IsDir() || !strings.HasSuffix(name, ".up.sql") {
			continue
		}
		prefix := strings.SplitN(filepath.Base(name), "_", 2)[0]
		version, parseErr := strconv.ParseUint(prefix, 10, 64)
		if parseErr != nil {
			continue
		}
		if version > latest {
			latest = version
		}
	}
	if latest == 0 {
		return 0, errors.New("no numbered up migrations found")
	}
	return uint(latest), nil
}

// CheckVersion verifies that golang-migrate reached the expected clean version.
func CheckVersion(ctx context.Context, db *gorm.DB, expected uint) error {
	if db == nil {
		return errors.New("database is not configured")
	}
	var state struct {
		Version uint
		Dirty   bool
	}
	if err := db.WithContext(ctx).Table("schema_migrations").Select("version, dirty").Take(&state).Error; err != nil {
		return fmt.Errorf("read schema migration state: %w", err)
	}
	if state.Dirty {
		return fmt.Errorf("schema migration %d is dirty", state.Version)
	}
	if state.Version != expected {
		return fmt.Errorf("schema migration version %d, expected %d", state.Version, expected)
	}
	return nil
}

func newMigrate(dbURL, migrationsPath string) (*migrate.Migrate, error) {
	if strings.TrimSpace(migrationsPath) == "" {
		return nil, errors.New("migration path must not be empty")
	}

	m, err := migrate.New(fmt.Sprintf("file://%s", migrationsPath), dbURL)
	if err != nil {
		return nil, fmt.Errorf("failed to create migrate instance: %w", err)
	}
	return m, nil
}