package database import ( "context" "fmt" "os" "strings" "sync" "github.com/Tencent/WeKnora/internal/logger" "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" ) var ( currentMigrationVersion uint currentMigrationDirty bool migrationVersionOnce sync.Once migrationVersionSet bool ) // CachedMigrationVersion returns the migration version captured at startup. // Returns (version, dirty, ok). ok is false if the version was never captured. func CachedMigrationVersion() (uint, bool, bool) { return currentMigrationVersion, currentMigrationDirty, migrationVersionSet } func setMigrationVersion(version uint, dirty bool) { migrationVersionOnce.Do(func() { currentMigrationVersion = version currentMigrationDirty = dirty migrationVersionSet = true }) } // RunMigrations executes all pending database migrations // This should be called during application startup func RunMigrations(dsn string) error { return RunMigrationsWithOptions(dsn, MigrationOptions{AutoRecoverDirty: false}) } // MigrationOptions configures migration behavior type MigrationOptions struct { // AutoRecoverDirty when true, automatically attempts to recover from dirty state // by forcing to the previous version and retrying the migration AutoRecoverDirty bool } // RunMigrationsWithOptions executes all pending database migrations with custom options func RunMigrationsWithOptions(dsn string, opts MigrationOptions) error { ctx := context.Background() logger.Infof(ctx, "Starting database migration...") migrationsPath := "file://migrations/versioned" if strings.HasPrefix(dsn, "sqlite3://") { migrationsPath = "file://migrations/sqlite" } m, err := migrate.New(migrationsPath, dsn) if err != nil { logger.Errorf(ctx, "Failed to create migrate instance: %v", err) return fmt.Errorf("failed to create migrate instance: %w", err) } defer m.Close() // Check current version and dirty state before migration oldVersion, oldDirty, versionErr := m.Version() if versionErr != nil && versionErr != migrate.ErrNilVersion { logger.Errorf(ctx, "Failed to get migration version: %v", versionErr) return fmt.Errorf("failed to get migration version: %w", versionErr) } if versionErr == migrate.ErrNilVersion { logger.Infof(ctx, "Database has no migration history, will start from version 0") } else { logger.Infof(ctx, "Current migration version: %d, dirty: %v", oldVersion, oldDirty) } // If database is in dirty state, try to recover or return error if oldDirty { logger.Warnf(ctx, "Database is in dirty state at version %d", oldVersion) if opts.AutoRecoverDirty { logger.Infof(ctx, "AutoRecoverDirty is enabled, attempting recovery...") if err := recoverFromDirtyState(ctx, m, oldVersion); err != nil { return err } // Update oldVersion after recovery oldVersion, _, _ = m.Version() } else { // Calculate the version to force to (usually the previous version) forceVersion := int(oldVersion) - 1 if oldVersion == 0 || forceVersion < 0 { forceVersion = 0 } return fmt.Errorf( "database is in dirty state at version %d. This usually means a migration failed partway through. "+ "To fix this:\n"+ "1. Check if the migration partially applied changes and manually fix if needed\n"+ "2. Use the force command to set the version to the last successful migration (usually %d):\n"+ " ./scripts/migrate.sh force %d\n"+ " Or if using make: make migrate-force version=%d\n"+ "3. After fixing, restart the application to retry the migration\n"+ "Or enable AutoRecoverDirty option to automatically retry", oldVersion, forceVersion, forceVersion, forceVersion, ) } } // Run all pending migrations logger.Infof(ctx, "Running pending migrations...") if err := m.Up(); err != nil && err != migrate.ErrNoChange { logger.Errorf(ctx, "Migration failed: %v", err) // Check if error is due to dirty state (in case it became dirty during migration) currentVersion, currentDirty, versionCheckErr := m.Version() if versionCheckErr == nil && currentDirty { logger.Warnf(ctx, "Migration caused dirty state at version %d", currentVersion) if opts.AutoRecoverDirty { logger.Infof(ctx, "Attempting to recover from dirty state...") // Try to recover and retry if recoverErr := recoverFromDirtyState(ctx, m, currentVersion); recoverErr != nil { return recoverErr } // Retry migration after recovery logger.Infof(ctx, "Retrying migration after recovery...") if retryErr := m.Up(); retryErr != nil && retryErr != migrate.ErrNoChange { logger.Errorf(ctx, "Migration failed after recovery attempt: %v", retryErr) return fmt.Errorf("migration failed after recovery attempt: %w", retryErr) } } else { // Calculate the version to force to (usually the previous version) forceVersion := currentVersion - 1 if currentVersion == 0 { forceVersion = 0 } return fmt.Errorf( "migration failed and database is now in dirty state at version %d. "+ "To fix this:\n"+ "1. Check if the migration partially applied changes and manually fix if needed\n"+ "2. Use the force command to set the version to the last successful migration (usually %d):\n"+ " ./scripts/migrate.sh force %d\n"+ " Or if using make: make migrate-force version=%d\n"+ "3. After fixing, restart the application to retry the migration\n"+ "Or enable AutoRecoverDirty option to automatically retry", currentVersion, forceVersion, forceVersion, forceVersion, ) } } else { return fmt.Errorf("failed to run migrations: %w", err) } } // Get current version after migration version, dirty, err := m.Version() if err != nil && err != migrate.ErrNilVersion { return fmt.Errorf("failed to get migration version: %w", err) } setMigrationVersion(version, dirty) if oldVersion != version { logger.Infof(ctx, "Database migrated from version %d to %d", oldVersion, version) } else { logger.Infof(ctx, "Database is up to date (version: %d)", version) } if dirty { logger.Warnf(ctx, "Database is in dirty state! Manual intervention may be required.") } return nil } // recoverFromDirtyState attempts to recover from a dirty migration state // by forcing to the previous version and allowing the migration to be retried func recoverFromDirtyState(ctx context.Context, m *migrate.Migrate, dirtyVersion uint) error { // Special case: if dirty at version 0 (init migration), we cannot go back further // The only option is to force to version 0 and retry, but this requires the migration to be idempotent if dirtyVersion == 0 { logger.Warnf(ctx, "Database is in dirty state at version 0 (init migration). "+ "This is the initial migration, cannot rollback further. "+ "Will attempt to clear dirty flag and retry. "+ "Note: This only works if the init migration uses IF NOT EXISTS clauses.") // Force to version -1 (no version) to allow re-running version 0 // This effectively tells migrate that no migrations have been applied if err := m.Force(-1); err != nil { return fmt.Errorf( "failed to recover from dirty state at version 0. "+ "Manual intervention required:\n"+ "1. Check what was partially created in the database\n"+ "2. Either drop all created objects and retry, or\n"+ "3. Manually complete the migration and run: ./scripts/migrate.sh force 0\n"+ "Error: %w", err) } logger.Infof(ctx, "Cleared migration state, will retry from version 0") return nil } forceVersion := int(dirtyVersion) - 1 logger.Warnf(ctx, "Database is in dirty state at version %d, attempting auto-recovery by forcing to version %d", dirtyVersion, forceVersion) // Force to previous version to clear dirty state if err := m.Force(forceVersion); err != nil { return fmt.Errorf("failed to force migration version during recovery: %w", err) } logger.Infof(ctx, "Successfully forced migration to version %d, migration will be retried", forceVersion) return nil } // GetMigrationVersion returns the current migration version func GetMigrationVersion() (uint, bool, error) { dbURL := fmt.Sprintf( "postgres://%s:%s@%s:%s/%s?sslmode=disable", os.Getenv("DB_USER"), os.Getenv("DB_PASSWORD"), os.Getenv("DB_HOST"), os.Getenv("DB_PORT"), os.Getenv("DB_NAME"), ) migrationsPath := "file://migrations/versioned" m, err := migrate.New(migrationsPath, dbURL) if err != nil { return 0, false, fmt.Errorf("failed to create migrate instance: %w", err) } defer m.Close() version, dirty, err := m.Version() if err != nil { return 0, false, err } return version, dirty, nil }