fix(rollback): plan from requested apply state

Rollback now plans from the exact completed apply being rolled back instead of whichever completed apply happens to be latest for the PR/database. It stores whether original schema capture succeeded separately from the schema map so create-table rollbacks can target an empty original schema, while older plans without capture still fail closed.

The confirmation path also keeps rollback locks distinct from normal apply locks, routes observers through the stored rollback deployment, makes rollback-confirm database-explicit, and rejects unsupported planning-time defer-cutover flags.

Vitess VSchema rollback remains fail-closed until original VSchema capture is supported.

Co-authored-by: Amp <amp@ampcode.com>
Amp-Thread-ID: https://ampcode.com/threads/T-019eb380-ca81-749e-94a2-cb4982ed90dd

Author: aparajon

pkg/api/control_handlers.go

@@ -1685,7 +1685,7 @@ func (s *Service) handleRollbackPlan(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	resp, err := s.ExecuteRollbackPlan(r.Context(), apply.Database, apply.Environment, apply.Deployment)
+	resp, err := s.ExecuteRollbackPlanForApply(r.Context(), apply)
 	if err != nil {
 		metrics.RecordControlOperation(r.Context(), "rollback_plan", apply.Database, apply.Deployment, apply.Environment, "error")
 		s.writeControlError(w, "rollback plan", apply, err)

pkg/api/handlers_test.go

@@ -596,9 +596,6 @@ func (m *mockTernClient) SkipRevert(ctx context.Context, req *ternv1.SkipRevertR
 	}
 	return nil, m.skipRevertErr
 }
-func (m *mockTernClient) RollbackPlan(ctx context.Context, database, environment string) (*ternv1.PlanResponse, error) {
-	return nil, nil
-}
 func (m *mockTernClient) ResumeApply(ctx context.Context, apply *storage.Apply) error {
 	m.resumeMu.Lock()
 	m.resumeApply = apply
@@ -3656,3 +3653,47 @@ func TestDeriveErrorCode(t *testing.T) {
 		})
 	}
 }
+
+func TestRollbackSchemaFilesAllowsCapturedEmptyOriginalSchema(t *testing.T) {
+	plan := &storage.Plan{
+		Namespaces: map[string]*storage.NamespacePlanData{
+			"shop": {
+				OriginalSchema:         map[string]string{},
+				OriginalSchemaCaptured: true,
+			},
+		},
+	}
+
+	schemaFiles, err := rollbackSchemaFiles(plan)
+	require.NoError(t, err)
+	require.Contains(t, schemaFiles, "shop")
+	assert.Empty(t, schemaFiles["shop"].Files)
+}
+
+func TestRollbackSchemaFilesRejectsLegacyPlanWithoutCapturedOriginalSchema(t *testing.T) {
+	plan := &storage.Plan{
+		Namespaces: map[string]*storage.NamespacePlanData{
+			"shop": {
+				OriginalSchema: map[string]string{},
+			},
+		},
+	}
+
+	schemaFiles, err := rollbackSchemaFiles(plan)
+	require.Error(t, err)
+	assert.Nil(t, schemaFiles)
+	assert.Contains(t, err.Error(), `no original schema available for rollback namespace "shop"`)
+}
+
+func TestPlanContainsVSchemaChange(t *testing.T) {
+	assert.True(t, planContainsVSchemaChange(&storage.Plan{
+		Namespaces: map[string]*storage.NamespacePlanData{
+			"commerce": {VSchema: []byte(`{"tables":{"users":{}}}`)},
+		},
+	}))
+	assert.False(t, planContainsVSchemaChange(&storage.Plan{
+		Namespaces: map[string]*storage.NamespacePlanData{
+			"commerce": {OriginalSchemaCaptured: true},
+		},
+	}))
+}

pkg/api/plan_handlers.go

@@ -466,27 +466,56 @@ func (s *Service) ExecutePlan(ctx context.Context, req PlanRequest) (*apitypes.P
 		}
 	}
 
-	// Store plan in SchemaBot's storage (idempotent — duplicate is ignored)
+	route := storedPlanRoute{
+		DatabaseType: resolvedTarget.DatabaseType,
+		Deployment:   deployment,
+		Target:       resolvedTarget.Target,
+	}
+	if err := s.storePlanResponse(ctx, req, resp, route); err != nil {
+		return nil, err
+	}
+
+	return planResponseFromProto(resp), nil
+}
+
+type storedPlanRoute struct {
+	DatabaseType string
+	Deployment   string
+	Target       string
+}
+
+func (s *Service) storePlanResponse(ctx context.Context, req PlanRequest, resp *ternv1.PlanResponse, route storedPlanRoute) error {
+	prInt := 0
+	if req.PullRequest != nil {
+		prInt = int(*req.PullRequest)
+	}
+	trustedSchemaPath := ""
+	if req.SourceTrusted {
+		trustedSchemaPath = req.SchemaPath
+	}
+	headSHA := ""
+	if req.HeadSHA != nil {
+		headSHA = *req.HeadSHA
+	}
 	storedPlan := &storage.Plan{
 		PlanIdentifier: resp.PlanId,
 		Database:       req.Database,
-		DatabaseType:   resolvedTarget.DatabaseType,
-		Deployment:     deployment,
-		Target:         resolvedTarget.Target,
+		DatabaseType:   route.DatabaseType,
+		Deployment:     route.Deployment,
+		Target:         route.Target,
 		Repository:     req.Repository,
 		PullRequest:    prInt,
 		SchemaPath:     trustedSchemaPath,
 		Environment:    req.Environment,
 		SchemaFiles:    protoToSchemaFiles(req.SchemaFiles),
-		Namespaces:     protoChangesToNamespaces(resp.Changes),
-		HeadSHA:        ternReq.HeadSha,
+		Namespaces:     protoChangesToNamespaces(resp.Changes, req.SchemaFiles),
+		HeadSHA:        headSHA,
 		CreatedAt:      time.Now(),
 	}
 	if _, err := s.storage.Plans().Create(ctx, storedPlan); err != nil && !errors.Is(err, storage.ErrPlanIDExists) {
-		return nil, fmt.Errorf("store plan: %w", err)
+		return fmt.Errorf("store plan: %w", err)
 	}
-
-	return planResponseFromProto(resp), nil
+	return nil
 }
 
 // handleApply handles POST /api/apply requests.
@@ -899,29 +928,114 @@ func applyTaskChanges(plan *storage.Plan) []storage.TableChange {
 	return changes
 }
 
-// ExecuteRollbackPlan generates a rollback plan via the Tern client.
-// The plan is automatically stored by the Tern client's RollbackPlan method
-// (which calls Plan internally). This is the shared implementation used by
-// both the HTTP handler and the webhook handler.
-func (s *Service) ExecuteRollbackPlan(ctx context.Context, database, environment, deployment string) (*apitypes.PlanResponse, error) {
-	deployment, err := s.deploymentForDatabaseEnvironment(database, deployment, environment)
+// ExecuteRollbackPlanForApply generates a rollback plan for a specific apply.
+func (s *Service) ExecuteRollbackPlanForApply(ctx context.Context, apply *storage.Apply) (*apitypes.PlanResponse, error) {
+	if apply == nil {
+		return nil, fmt.Errorf("apply is required")
+	}
+	if !state.IsState(apply.State, state.Apply.Completed) {
+		return nil, fmt.Errorf("apply %s is in state %q; only completed applies can be rolled back", apply.ApplyIdentifier, apply.State)
+	}
+
+	plan, err := s.storage.Plans().GetByID(ctx, apply.PlanID)
+	if err != nil {
+		return nil, fmt.Errorf("get rollback source plan: %w", err)
+	}
+	if plan == nil {
+		return nil, fmt.Errorf("plan not found for apply %s", apply.ApplyIdentifier)
+	}
+	if !rollbackSourcePlanMatchesApply(plan, apply) {
+		return nil, fmt.Errorf("source plan %s belongs to %s/%s/%s, not apply %s for %s/%s/%s",
+			plan.PlanIdentifier, plan.Database, plan.DatabaseType, plan.Environment,
+			apply.ApplyIdentifier, apply.Database, apply.DatabaseType, apply.Environment)
+	}
+	if apply.DatabaseType == storage.DatabaseTypeVitess && planContainsVSchemaChange(plan) {
+		return nil, fmt.Errorf("rollback of Vitess VSchema changes is not supported because the original VSchema was not captured; reconcile manually")
+	}
+
+	schemaFiles, err := rollbackSchemaFiles(plan)
 	if err != nil {
-		return nil, fmt.Errorf("resolve deployment for rollback plan: %w", err)
+		return nil, err
 	}
 
-	client, err := s.TernClient(deployment, environment)
+	deployment, err := storedDeploymentForApply(apply)
+	if err != nil {
+		return nil, err
+	}
+	if plan.Target == "" {
+		return nil, fmt.Errorf("plan %s is missing server-side routing metadata field %q; create a new plan and retry rollback", plan.PlanIdentifier, "target")
+	}
+	client, err := s.TernClient(deployment, apply.Environment)
 	if err != nil {
-		return nil, fmt.Errorf("database %q (%s): %w", database, environment, err)
+		return nil, fmt.Errorf("database %q (%s): %w", apply.Database, apply.Environment, err)
 	}
 
-	resp, err := client.RollbackPlan(ctx, database, environment)
+	prNumber := int32(apply.PullRequest)
+	req := PlanRequest{
+		Database:    apply.Database,
+		Environment: apply.Environment,
+		Type:        apply.DatabaseType,
+		SchemaFiles: schemaFiles,
+		Repository:  apply.Repository,
+		PullRequest: &prNumber,
+	}
+	resp, err := client.Plan(ctx, &ternv1.PlanRequest{
+		Database:    req.Database,
+		Type:        req.Type,
+		SchemaFiles: req.SchemaFiles,
+		Repository:  req.Repository,
+		PullRequest: prNumber,
+		Environment: req.Environment,
+		Target:      plan.Target,
+	})
 	if err != nil {
 		return nil, err
 	}
+	route := storedPlanRoute{
+		DatabaseType: apply.DatabaseType,
+		Deployment:   deployment,
+		Target:       plan.Target,
+	}
+	if err := s.storePlanResponse(ctx, req, resp, route); err != nil {
+		return nil, err
+	}
 
 	return planResponseFromProto(resp), nil
 }
 
+func rollbackSourcePlanMatchesApply(plan *storage.Plan, apply *storage.Apply) bool {
+	return plan.Database == apply.Database &&
+		plan.DatabaseType == apply.DatabaseType &&
+		plan.Environment == apply.Environment
+}
+
+func planContainsVSchemaChange(plan *storage.Plan) bool {
+	for _, nsData := range plan.Namespaces {
+		if nsData != nil && len(nsData.VSchema) > 0 {
+			return true
+		}
+	}
+	return false
+}
+
+func rollbackSchemaFiles(plan *storage.Plan) (map[string]*ternv1.SchemaFiles, error) {
+	schemaFiles := make(map[string]*ternv1.SchemaFiles)
+	for ns, nsData := range plan.Namespaces {
+		if nsData == nil || !nsData.OriginalSchemaCaptured {
+			return nil, fmt.Errorf("no original schema available for rollback namespace %q (plan may predate rollback feature)", ns)
+		}
+		sqlFiles := make(map[string]string, len(nsData.OriginalSchema))
+		for tableName, createSQL := range nsData.OriginalSchema {
+			sqlFiles[tableName+".sql"] = createSQL
+		}
+		schemaFiles[ns] = &ternv1.SchemaFiles{Files: sqlFiles}
+	}
+	if len(schemaFiles) == 0 {
+		return nil, fmt.Errorf("no namespaces available for rollback")
+	}
+	return schemaFiles, nil
+}
+
 // validateSchemaFiles checks that schema_files has at least one namespace and
 // that every namespace carries a non-null value. An empty Files map within a
 // namespace is valid (signals "drop all tables"), so we only reject when

pkg/api/proto_helpers.go

@@ -84,26 +84,40 @@ func planResponseFromProto(resp *ternv1.PlanResponse) *apitypes.PlanResponse {
 }
 
 // protoChangesToNamespaces converts proto SchemaChanges to storage namespace plan data.
-func protoChangesToNamespaces(changes []*ternv1.SchemaChange) map[string]*storage.NamespacePlanData {
+func protoChangesToNamespaces(changes []*ternv1.SchemaChange, schemaFiles map[string]*ternv1.SchemaFiles) map[string]*storage.NamespacePlanData {
 	result := make(map[string]*storage.NamespacePlanData)
 	for _, sc := range changes {
 		ns := sc.Namespace
 		if ns == "" {
 			ns = "default"
 		}
-		nsData := &storage.NamespacePlanData{}
+		nsData := &storage.NamespacePlanData{
+			OriginalSchema:         sc.OriginalSchema,
+			OriginalSchemaCaptured: sc.OriginalSchemaCaptured,
+		}
 		for _, t := range sc.TableChanges {
 			nsData.Tables = append(nsData.Tables, storage.TableChange{
 				Table:     t.TableName,
 				DDL:       t.Ddl,
 				Operation: protoChangeTypeToOperation(t.ChangeType),
 			})
 		}
+		if schemaChangeHasVSchema(sc) {
+			if sf := schemaFiles[ns]; sf != nil {
+				if raw := sf.Files["vschema.json"]; raw != "" {
+					nsData.VSchema = []byte(raw)
+				}
+			}
+		}
 		result[ns] = nsData
 	}
 	return result
 }
 
+func schemaChangeHasVSchema(sc *ternv1.SchemaChange) bool {
+	return sc.Metadata["vschema"] != "" || sc.Metadata["vschema_changed"] == "true"
+}
+
 // protoChangeTypeToOperation converts a proto ChangeType enum to a storage operation string.
 func protoChangeTypeToOperation(ct ternv1.ChangeType) string {
 	switch ct {

pkg/api/proto_helpers_test.go

@@ -115,3 +115,19 @@ func TestValidateSchemaFiles_MissingRejected(t *testing.T) {
 	assert.Empty(t, warning)
 	assert.Contains(t, err.Error(), "schema_files is required")
 }
+
+func TestProtoChangesToNamespacesPreservesVSchemaFile(t *testing.T) {
+	changes := []*ternv1.SchemaChange{
+		{
+			Namespace: "commerce",
+			Metadata:  map[string]string{"vschema_changed": "true"},
+		},
+	}
+	schemaFiles := map[string]*ternv1.SchemaFiles{
+		"commerce": {Files: map[string]string{"vschema.json": `{"tables":{"users":{}}}`}},
+	}
+
+	namespaces := protoChangesToNamespaces(changes, schemaFiles)
+
+	assert.JSONEq(t, `{"tables":{"users":{}}}`, string(namespaces["commerce"].VSchema))
+}