refactor(storage): replace batch dependent partial update

submitqueue/extension/storage/batch_dependent_store.go

@@ -37,8 +37,8 @@ type BatchDependentStore interface {
 	// Returns ErrAlreadyExists if the entry already exists for the given batch ID.
 	Create(ctx context.Context, batchDependent entity.BatchDependent) error
 
-	// UpdateDependents updates the dependents of a batch dependent and the version to newVersion
+	// Update replaces the non-key fields of a batch dependent and persists newVersion
 	// if the current persisted version matches oldVersion. If versions do not match, returns ErrVersionMismatch.
 	// Version arithmetic is owned by the caller; the store performs a pure conditional write.
-	UpdateDependents(ctx context.Context, batchID string, oldVersion, newVersion int32, dependents []string) error
+	Update(ctx context.Context, batchDependent entity.BatchDependent, oldVersion, newVersion int32) error
 }

submitqueue/extension/storage/mysql/batch_dependent_store.go

@@ -91,41 +91,41 @@ func (s *batchDependentStore) Create(ctx context.Context, batchDependent entity.
 	return nil
 }
 
-// UpdateDependents updates the dependents of a batch dependent and the version to newVersion
+// Update replaces the non-key fields of a batch dependent and persists newVersion
 // if the current persisted version matches oldVersion. If versions do not match, returns ErrVersionMismatch.
 // Version arithmetic is owned by the caller; this is a pure conditional write.
-func (s *batchDependentStore) UpdateDependents(ctx context.Context, batchID string, oldVersion, newVersion int32, dependents []string) (retErr error) {
+func (s *batchDependentStore) Update(ctx context.Context, batchDependent entity.BatchDependent, oldVersion, newVersion int32) (retErr error) {
 	op := metrics.Begin(s.scope, "update_dependents", metrics.StorageLatencyBuckets)
 	defer func() { op.Complete(retErr) }()
 
-	dependentsJSON, err := json.Marshal(dependents)
+	dependentsJSON, err := json.Marshal(batchDependent.Dependents)
 	if err != nil {
-		return fmt.Errorf("failed to marshal dependents batchID=%s for UpdateDependents batch dependent entity: %w", batchID, err)
+		return fmt.Errorf("failed to marshal dependents batchID=%s for Update batch dependent entity: %w", batchDependent.BatchID, err)
 	}
 
 	result, err := s.db.ExecContext(ctx,
 		"UPDATE batch_dependent SET dependents = ?, version = ? WHERE batch_id = ? AND version = ?",
-		dependentsJSON, newVersion, batchID, oldVersion,
+		dependentsJSON, newVersion, batchDependent.BatchID, oldVersion,
 	)
 	if err != nil {
 		return fmt.Errorf(
 			"failed to update batch dependent dependents for batchID=%q oldVersion=%d newVersion=%d: %w",
-			batchID, oldVersion, newVersion, err,
+			batchDependent.BatchID, oldVersion, newVersion, err,
 		)
 	}
 
 	rowsAffected, err := result.RowsAffected()
 	if err != nil {
 		return fmt.Errorf(
 			"failed to get rows affected from update for batchID=%q oldVersion=%d newVersion=%d: %w",
-			batchID, oldVersion, newVersion, err,
+			batchDependent.BatchID, oldVersion, newVersion, err,
 		)
 	}
 
 	if rowsAffected != 1 {
 		return fmt.Errorf(
 			"version mismatch for batch dependent update: batchID=%q expected_version=%d: %w",
-			batchID, oldVersion, storage.ErrVersionMismatch,
+			batchDependent.BatchID, oldVersion, storage.ErrVersionMismatch,
 		)
 	}
 

submitqueue/extension/storage/mysql/batch_dependent_store_test.go

@@ -70,6 +70,37 @@ func TestBatchDependentStore_Get(t *testing.T) {
 			},
 			want: want,
 		},
+		{
+			name:    "found with nil dependents",
+			batchID: "monorepo/batch/nil",
+			setup: func(mock sqlmock.Sqlmock) {
+				rows := sqlmock.NewRows([]string{"batch_id", "dependents", "version"}).
+					AddRow("monorepo/batch/nil", []byte("null"), int32(2))
+				mock.ExpectQuery("SELECT batch_id, dependents, version FROM batch_dependent").
+					WithArgs("monorepo/batch/nil").
+					WillReturnRows(rows)
+			},
+			want: entity.BatchDependent{
+				BatchID: "monorepo/batch/nil",
+				Version: 2,
+			},
+		},
+		{
+			name:    "found with empty dependents",
+			batchID: "monorepo/batch/empty",
+			setup: func(mock sqlmock.Sqlmock) {
+				rows := sqlmock.NewRows([]string{"batch_id", "dependents", "version"}).
+					AddRow("monorepo/batch/empty", []byte("[]"), int32(3))
+				mock.ExpectQuery("SELECT batch_id, dependents, version FROM batch_dependent").
+					WithArgs("monorepo/batch/empty").
+					WillReturnRows(rows)
+			},
+			want: entity.BatchDependent{
+				BatchID:    "monorepo/batch/empty",
+				Dependents: []string{},
+				Version:    3,
+			},
+		},
 		{
 			name:    "not found",
 			batchID: "missing",
@@ -178,49 +209,82 @@ func TestBatchDependentStore_Create(t *testing.T) {
 	}
 }
 
-func TestBatchDependentStore_UpdateDependents(t *testing.T) {
-	const batchID = "monorepo/batch/1"
+func TestBatchDependentStore_Update(t *testing.T) {
 	const oldVersion, newVersion = int32(1), int32(2)
-	dependents := []string{"monorepo/batch/2", "monorepo/batch/3"}
+	batchDependent := entity.BatchDependent{
+		BatchID:    "monorepo/batch/1",
+		Dependents: []string{"monorepo/batch/2", "monorepo/batch/3"},
+		Version:    oldVersion,
+	}
 
 	tests := []struct {
 		name      string
+		entity    entity.BatchDependent
 		setup     func(mock sqlmock.Sqlmock)
 		wantErr   bool
 		wantErrIs error
 	}{
 		{
-			name: "success",
+			name:   "success",
+			entity: batchDependent,
+			setup: func(mock sqlmock.Sqlmock) {
+				mock.ExpectExec("UPDATE batch_dependent").
+					WithArgs([]byte(`["monorepo/batch/2","monorepo/batch/3"]`), newVersion, batchDependent.BatchID, oldVersion).
+					WillReturnResult(sqlmock.NewResult(0, 1))
+			},
+		},
+		{
+			name: "success with nil dependents",
+			entity: entity.BatchDependent{
+				BatchID: "monorepo/batch/nil",
+				Version: oldVersion,
+			},
+			setup: func(mock sqlmock.Sqlmock) {
+				mock.ExpectExec("UPDATE batch_dependent").
+					WithArgs([]byte("null"), newVersion, "monorepo/batch/nil", oldVersion).
+					WillReturnResult(sqlmock.NewResult(0, 1))
+			},
+		},
+		{
+			name: "success with empty dependents",
+			entity: entity.BatchDependent{
+				BatchID:    "monorepo/batch/empty",
+				Dependents: []string{},
+				Version:    oldVersion,
+			},
 			setup: func(mock sqlmock.Sqlmock) {
 				mock.ExpectExec("UPDATE batch_dependent").
-					WithArgs(sqlmock.AnyArg(), newVersion, batchID, oldVersion).
+					WithArgs([]byte("[]"), newVersion, "monorepo/batch/empty", oldVersion).
 					WillReturnResult(sqlmock.NewResult(0, 1))
 			},
 		},
 		{
-			name: "version mismatch",
+			name:   "version mismatch",
+			entity: batchDependent,
 			setup: func(mock sqlmock.Sqlmock) {
 				mock.ExpectExec("UPDATE batch_dependent").
-					WithArgs(sqlmock.AnyArg(), newVersion, batchID, oldVersion).
+					WithArgs([]byte(`["monorepo/batch/2","monorepo/batch/3"]`), newVersion, batchDependent.BatchID, oldVersion).
 					WillReturnResult(sqlmock.NewResult(0, 0))
 			},
 			wantErr:   true,
 			wantErrIs: storage.ErrVersionMismatch,
 		},
 		{
-			name: "exec error",
+			name:   "exec error",
+			entity: batchDependent,
 			setup: func(mock sqlmock.Sqlmock) {
 				mock.ExpectExec("UPDATE batch_dependent").
-					WithArgs(sqlmock.AnyArg(), newVersion, batchID, oldVersion).
+					WithArgs([]byte(`["monorepo/batch/2","monorepo/batch/3"]`), newVersion, batchDependent.BatchID, oldVersion).
 					WillReturnError(fmt.Errorf("connection reset"))
 			},
 			wantErr: true,
 		},
 		{
-			name: "rows affected error",
+			name:   "rows affected error",
+			entity: batchDependent,
 			setup: func(mock sqlmock.Sqlmock) {
 				mock.ExpectExec("UPDATE batch_dependent").
-					WithArgs(sqlmock.AnyArg(), newVersion, batchID, oldVersion).
+					WithArgs([]byte(`["monorepo/batch/2","monorepo/batch/3"]`), newVersion, batchDependent.BatchID, oldVersion).
 					WillReturnResult(sqlmock.NewErrorResult(fmt.Errorf("driver error")))
 			},
 			wantErr: true,
@@ -234,7 +298,7 @@ func TestBatchDependentStore_UpdateDependents(t *testing.T) {
 
 			tt.setup(mock)
 
-			err := store.UpdateDependents(context.Background(), batchID, oldVersion, newVersion, dependents)
+			err := store.Update(context.Background(), tt.entity, oldVersion, newVersion)
 			if tt.wantErr {
 				require.Error(t, err)
 				if tt.wantErrIs != nil {

submitqueue/orchestrator/controller/batch/batch.go

@@ -325,9 +325,11 @@ func (c *Controller) populateBatch(ctx context.Context, batch entity.Batch) (ent
 			return entity.Batch{}, fmt.Errorf("failed to get batch dependent for batchID=%s: %w", dependencyID, err)
 		}
 
-		dependents := append(existing.Dependents, batch.ID)
+		updated := existing
+		updated.Dependents = append([]string(nil), existing.Dependents...)
+		updated.Dependents = append(updated.Dependents, batch.ID)
 		newVersion := existing.Version + 1
-		if err := c.store.GetBatchDependentStore().UpdateDependents(ctx, dependencyID, existing.Version, newVersion, dependents); err != nil {
+		if err := c.store.GetBatchDependentStore().Update(ctx, updated, existing.Version, newVersion); err != nil {
 			metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1)
 			return entity.Batch{}, fmt.Errorf("failed to update batch dependent index for existing batchID=%s and new batchID=%s: %w", dependencyID, batch.ID, err)
 		}

submitqueue/orchestrator/controller/batch/batch_test.go

@@ -357,14 +357,22 @@ func TestController_Process_WithDependencies(t *testing.T) {
 		BatchID: "test-queue/batch/1",
 		Version: 1,
 	}, nil)
-	mockBatchDependentStore.EXPECT().UpdateDependents(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), gomock.Any()).Return(nil)
+	mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{
+		BatchID:    "test-queue/batch/1",
+		Dependents: []string{"test-queue/batch/1"},
+		Version:    1,
+	}, int32(1), int32(2)).Return(nil)
 	// batch/2 already has an existing dependent.
 	mockBatchDependentStore.EXPECT().Get(gomock.Any(), "test-queue/batch/2").Return(entity.BatchDependent{
 		BatchID:    "test-queue/batch/2",
 		Dependents: []string{"test-queue/batch/99"},
 		Version:    2,
 	}, nil)
-	mockBatchDependentStore.EXPECT().UpdateDependents(gomock.Any(), "test-queue/batch/2", int32(2), int32(3), gomock.Any()).Return(nil)
+	mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{
+		BatchID:    "test-queue/batch/2",
+		Dependents: []string{"test-queue/batch/99", "test-queue/batch/1"},
+		Version:    2,
+	}, int32(2), int32(3)).Return(nil)
 	// Create empty reverse index for the new batch.
 	mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
 
@@ -418,7 +426,11 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) {
 		BatchID: "test-queue/batch/2",
 		Version: 5,
 	}, nil)
-	mockBatchDependentStore.EXPECT().UpdateDependents(gomock.Any(), "test-queue/batch/2", int32(5), int32(6), gomock.Any()).Return(nil)
+	mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{
+		BatchID:    "test-queue/batch/2",
+		Dependents: []string{"test-queue/batch/1"},
+		Version:    5,
+	}, int32(5), int32(6)).Return(nil)
 	mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
 
 	mockReqStore := storagemock.NewMockRequestStore(ctrl)
@@ -457,6 +469,56 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestController_Process_BatchDependentUpdateFailureDoesNotMutateFetchedDependents(t *testing.T) {
+	ctrl := gomock.NewController(t)
+
+	request := testRequest()
+	activeBatch := entity.Batch{
+		ID:      "test-queue/batch/99",
+		Queue:   "test-queue",
+		State:   entity.BatchStateCreated,
+		Version: 1,
+	}
+	dependents := make([]string, 1, 2)
+	dependents[0] = "test-queue/batch/98"
+	existing := entity.BatchDependent{
+		BatchID:    activeBatch.ID,
+		Dependents: dependents,
+		Version:    4,
+	}
+
+	mockBatchStore := storagemock.NewMockBatchStore(ctrl)
+	mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, gomock.Any()).Return([]entity.Batch{activeBatch}, nil)
+
+	mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl)
+	mockBatchDependentStore.EXPECT().Get(gomock.Any(), activeBatch.ID).Return(existing, nil)
+	mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{
+		BatchID:    activeBatch.ID,
+		Dependents: []string{"test-queue/batch/98", "test-queue/batch/1"},
+		Version:    existing.Version,
+	}, existing.Version, existing.Version+1).Return(errors.New("update failed"))
+
+	mockReqStore := storagemock.NewMockRequestStore(ctrl)
+	mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil)
+
+	mockStorage := storagemock.NewMockStorage(ctrl)
+	mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes()
+	mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes()
+	mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes()
+
+	controller := newTestController(t, ctrl, newSequentialCounter(ctrl), mockStorage, nil, nil)
+
+	msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil)
+	delivery := queuemock.NewMockDelivery(ctrl)
+	delivery.EXPECT().Message().Return(msg).AnyTimes()
+	delivery.EXPECT().Attempt().Return(1).AnyTimes()
+
+	err := controller.Process(context.Background(), delivery)
+	require.Error(t, err)
+	assert.Equal(t, "", dependents[:cap(dependents)][1])
+	assert.Equal(t, int32(4), existing.Version)
+}
+
 func TestController_Process_AnalyzerFailure(t *testing.T) {
 	ctrl := gomock.NewController(t)
 
@@ -894,9 +956,11 @@ func TestController_PopulateBatch_Errors(t *testing.T) {
 					BatchID: "test-queue/batch/0",
 					Version: 2,
 				}, nil)
-				dependentStore.EXPECT().UpdateDependents(
-					gomock.Any(), "test-queue/batch/0", int32(2), int32(3), []string{batch.ID},
-				).Return(storeErr)
+				dependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{
+					BatchID:    "test-queue/batch/0",
+					Dependents: []string{batch.ID},
+					Version:    2,
+				}, int32(2), int32(3)).Return(storeErr)
 			},
 			errMsg: "failed to update batch dependent index",
 		},
@@ -909,9 +973,11 @@ func TestController_PopulateBatch_Errors(t *testing.T) {
 					Dependents: []string{"test-queue/batch/old"},
 					Version:    2,
 				}, nil)
-				dependentStore.EXPECT().UpdateDependents(
-					gomock.Any(), "test-queue/batch/0", int32(2), int32(3), []string{"test-queue/batch/old", batch.ID},
-				).Return(nil)
+				dependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{
+					BatchID:    "test-queue/batch/0",
+					Dependents: []string{"test-queue/batch/old", batch.ID},
+					Version:    2,
+				}, int32(2), int32(3)).Return(nil)
 				batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCreated).Return(storeErr)
 			},
 			errMsg: "failed to mark batch",