fix(planetscale): persist recovered migration context

Amp-Thread-ID: https://ampcode.com/threads/T-019e7457-78b6-742d-8b89-33144cc92e27
Co-authored-by: Amp <amp@ampcode.com>

Author: aparajon

pkg/tern/local_client.go

@@ -734,6 +734,7 @@ func (c *LocalClient) Progress(ctx context.Context, req *ternv1.ProgressRequest)
 	// can poll the deploy request and query SHOW VITESS_MIGRATIONS.
 	var engineResult *engine.ProgressResult
 	var vitessApplyIsInstant bool
+	var vitessApplyData *storage.VitessApplyData
 	// Query engine for live progress. For Vitess, also query during pending state
 	// to surface PlanetScale states (preparing branch, deploy request, etc.).
 	queryDuringPending := c.config.Type == storage.DatabaseTypeVitess
@@ -756,6 +757,7 @@ func (c *LocalClient) Progress(ctx context.Context, req *ternv1.ProgressRequest)
 			case vad == nil:
 				c.logger.Warn("VitessApplyData not found for progress — apply may still be initializing", "apply_id", activeTask.ApplyID)
 			default:
+				vitessApplyData = vad
 				vitessApplyIsInstant = vad.IsInstant
 				resumeState, resumeErr := planetscale.BuildResumeState(planetscaleResumeData(vad))
 				if resumeErr != nil {
@@ -768,6 +770,12 @@ func (c *LocalClient) Progress(ctx context.Context, req *ternv1.ProgressRequest)
 		result, err := eng.Progress(ctx, progressReq)
 		if err == nil {
 			engineResult = result
+			if c.config.Type == storage.DatabaseTypeVitess && vitessApplyData != nil && result.ResumeState != nil && result.ResumeState.MigrationContext != "" && vitessApplyData.MigrationContext != result.ResumeState.MigrationContext {
+				vitessApplyData.MigrationContext = result.ResumeState.MigrationContext
+				if saveErr := c.storage.VitessApplyData().Save(ctx, vitessApplyData); saveErr != nil {
+					c.logger.Warn("failed to persist recovered Vitess migration context", "apply_id", activeTask.ApplyID, "migration_context", result.ResumeState.MigrationContext, "error", saveErr)
+				}
+			}
 			c.logger.Info("Progress: engine returned", "engine_state", result.State, "message", result.Message, "task_id", activeTask.TaskIdentifier, "storage_state", activeTask.State)
 			engineTaskState := taskStateFromProgressResult(result)
 			taskState := taskStateWithNoBackwardProgress(activeTask.State, engineTaskState)

pkg/tern/local_client_test.go

@@ -71,10 +71,12 @@ func (s *exactProgressTaskStore) Update(context.Context, *storage.Task) error {
 
 type fakeControlEngine struct {
 	engine.Engine
-	stopCount     int
-	cutoverCount  int
-	cutoverResult *engine.ControlResult
-	cutoverErr    error
+	stopCount      int
+	cutoverCount   int
+	progressResult *engine.ProgressResult
+	progressErr    error
+	cutoverResult  *engine.ControlResult
+	cutoverErr     error
 }
 
 func (e *fakeControlEngine) Name() string { return "fake" }
@@ -88,6 +90,9 @@ func (e *fakeControlEngine) Apply(context.Context, *engine.ApplyRequest) (*engin
 }
 
 func (e *fakeControlEngine) Progress(context.Context, *engine.ProgressRequest) (*engine.ProgressResult, error) {
+	if e.progressResult != nil || e.progressErr != nil {
+		return e.progressResult, e.progressErr
+	}
 	return &engine.ProgressResult{State: engine.StateRunning}, nil
 }
 
@@ -126,6 +131,7 @@ type exactProgressStorage struct {
 	tasks           storage.TaskStore
 	logs            storage.ApplyLogStore
 	controlRequests storage.ControlRequestStore
+	vitessApplyData storage.VitessApplyDataStore
 }
 
 func (s *exactProgressStorage) Applies() storage.ApplyStore { return s.applies }
@@ -139,6 +145,30 @@ func (s *exactProgressStorage) ApplyLogs() storage.ApplyLogStore {
 func (s *exactProgressStorage) ControlRequests() storage.ControlRequestStore {
 	return s.controlRequests
 }
+func (s *exactProgressStorage) VitessApplyData() storage.VitessApplyDataStore {
+	return s.vitessApplyData
+}
+
+type exactProgressVitessApplyDataStore struct {
+	storage.VitessApplyDataStore
+	data     *storage.VitessApplyData
+	saveData *storage.VitessApplyData
+	err      error
+}
+
+func (s *exactProgressVitessApplyDataStore) GetByApplyID(context.Context, int64) (*storage.VitessApplyData, error) {
+	return s.data, s.err
+}
+
+func (s *exactProgressVitessApplyDataStore) Save(_ context.Context, data *storage.VitessApplyData) error {
+	if s.err != nil {
+		return s.err
+	}
+	clone := *data
+	s.saveData = &clone
+	s.data = data
+	return nil
+}
 
 func TestApplyCancelHandleDoesNotCancelNewerOwner(t *testing.T) {
 	client := &LocalClient{}
@@ -291,6 +321,67 @@ func TestLocalClient_ProgressByApplyIDReturnsNotFoundForMissingApplyData(t *test
 	}
 }
 
+func TestLocalClient_ProgressPersistsRecoveredVitessMigrationContext(t *testing.T) {
+	apply := &storage.Apply{
+		ID:              42,
+		ApplyIdentifier: "apply-vitess-recovered-context",
+		Database:        "testdb",
+		DatabaseType:    storage.DatabaseTypeVitess,
+		Environment:     "staging",
+		State:           state.Apply.Running,
+		Engine:          storage.EnginePlanetScale,
+	}
+	task := &storage.Task{
+		ID:             43,
+		ApplyID:        apply.ID,
+		TaskIdentifier: "task-vitess-recovered-context",
+		Database:       "testdb",
+		DatabaseType:   storage.DatabaseTypeVitess,
+		Namespace:      "testdb",
+		TableName:      "users",
+		State:          state.Task.Running,
+	}
+	vitessData := &storage.VitessApplyData{
+		ApplyID:         apply.ID,
+		BranchName:      "schemabot-apply",
+		DeployRequestID: 123,
+		ExistingMigrationCtxs: map[string]storage.VitessMigrationContextTimestamps{
+			"preexisting": {RequestedTimestamp: "2026-01-01 00:00:00"},
+		},
+	}
+	vitessStore := &exactProgressVitessApplyDataStore{data: vitessData}
+	client := &LocalClient{
+		config: LocalConfig{
+			Database: "testdb",
+			Type:     storage.DatabaseTypeVitess,
+		},
+		storage: &exactProgressStorage{
+			applies:         &exactProgressApplyStore{apply: apply},
+			tasks:           &exactProgressTaskStore{tasks: []*storage.Task{task}},
+			vitessApplyData: vitessStore,
+		},
+		planetscaleEngine: &fakeControlEngine{progressResult: &engine.ProgressResult{
+			State: engine.StateRunning,
+			ResumeState: &engine.ResumeState{
+				MigrationContext: "recovered-context",
+			},
+		}},
+		logger: slog.Default(),
+	}
+
+	_, err := client.Progress(t.Context(), &ternv1.ProgressRequest{
+		ApplyId:     apply.ApplyIdentifier,
+		Environment: apply.Environment,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, vitessStore.saveData)
+	assert.Equal(t, "recovered-context", vitessStore.saveData.MigrationContext)
+	assert.Equal(t, apply.ID, vitessStore.saveData.ApplyID)
+	assert.Equal(t, "schemabot-apply", vitessStore.saveData.BranchName)
+	assert.Equal(t, uint64(123), vitessStore.saveData.DeployRequestID)
+	assert.Contains(t, vitessStore.saveData.ExistingMigrationCtxs, "preexisting")
+}
+
 func TestLocalClient_ProcessPendingStopControlRequest(t *testing.T) {
 	apply := &storage.Apply{
 		ID:              123,