Add UpdateItem to database service for partial attribute updates

database/database.go

@@ -135,6 +135,22 @@ func (db *Database) GetItem(ctx context.Context, table string, key map[string]an
 	return out.(map[string]any), nil
 }
 
+func (db *Database) UpdateItem(ctx context.Context, input driver.UpdateItemInput) (map[string]any, error) {
+	out, err := db.do(ctx, "UpdateItem", map[string]any{"table": input.Table}, func() (any, error) {
+		return db.driver.UpdateItem(ctx, input)
+	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	if out == nil {
+		return nil, nil
+	}
+
+	return out.(map[string]any), nil
+}
+
 func (db *Database) DeleteItem(ctx context.Context, table string, key map[string]any) error {
 	_, err := db.do(ctx, "DeleteItem", map[string]any{"table": table}, func() (any, error) {
 		return nil, db.driver.DeleteItem(ctx, table, key)

database/driver/driver.go

@@ -52,6 +52,20 @@ type GSIConfig struct {
 	SortKey      string
 }
 
+// UpdateAction represents a single field-level update action.
+type UpdateAction struct {
+	Action string // "SET" or "REMOVE"
+	Field  string
+	Value  any // ignored for REMOVE
+}
+
+// UpdateItemInput describes an update operation on an existing item.
+type UpdateItemInput struct {
+	Table   string
+	Key     map[string]any
+	Actions []UpdateAction
+}
+
 // KeyCondition defines a key condition for queries.
 type KeyCondition struct {
 	PartitionKey string
@@ -102,6 +116,7 @@ type Database interface {
 
 	PutItem(ctx context.Context, table string, item map[string]any) error
 	GetItem(ctx context.Context, table string, key map[string]any) (map[string]any, error)
+	UpdateItem(ctx context.Context, input UpdateItemInput) (map[string]any, error)
 	DeleteItem(ctx context.Context, table string, key map[string]any) error
 	Query(ctx context.Context, input QueryInput) (*QueryResult, error)
 	Scan(ctx context.Context, input ScanInput) (*QueryResult, error)

providers/aws/dynamodb/dynamodb.go

@@ -190,6 +190,50 @@ func (m *Mock) GetItem(_ context.Context, table string, key map[string]any) (map
 	return item, nil
 }
 
+// UpdateItem applies partial updates to an existing item.
+func (m *Mock) UpdateItem(_ context.Context, input driver.UpdateItemInput) (map[string]any, error) {
+	m.mu.Lock()
+
+	td, exists := m.tables[input.Table]
+	if !exists {
+		m.mu.Unlock()
+		return nil, cerrors.Newf(cerrors.NotFound, "table %s not found", input.Table)
+	}
+
+	k := itemKey(td.config, input.Key)
+	item, ok := td.items.Get(k)
+
+	if !ok {
+		m.mu.Unlock()
+		return nil, cerrors.New(cerrors.NotFound, "item not found")
+	}
+
+	oldItem := copyItem(item)
+	updated := copyItem(item)
+
+	for _, action := range input.Actions {
+		switch action.Action {
+		case "SET":
+			updated[action.Field] = action.Value
+		case "REMOVE":
+			delete(updated, action.Field)
+		default:
+			m.mu.Unlock()
+			return nil, cerrors.Newf(cerrors.InvalidArgument, "unsupported action: %s", action.Action)
+		}
+	}
+
+	td.items.Set(k, updated)
+	m.recordStreamEvent(td, oldItem, updated, true)
+	m.mu.Unlock()
+
+	dims := map[string]string{"TableName": input.Table}
+	m.emitMetric("ConsumedWriteCapacityUnits", 1, dims)
+	m.emitMetric("SuccessfulRequestCount", 1, dims)
+
+	return updated, nil
+}
+
 func (m *Mock) DeleteItem(_ context.Context, table string, key map[string]any) error {
 	m.mu.Lock()
 

providers/aws/dynamodb/dynamodb_test.go

@@ -714,3 +714,193 @@ func assertNotEmpty(t *testing.T, s string) {
 		t.Error("expected non-empty string")
 	}
 }
+
+func TestUpdateItemSetFields(t *testing.T) {
+	m := newTestMock()
+	ctx := context.Background()
+	createTestTable(m, "tbl")
+
+	_ = m.PutItem(ctx, "tbl", map[string]any{"pk": "u1", "sk": "info", "name": "Alice", "age": 30})
+
+	updated, err := m.UpdateItem(ctx, driver.UpdateItemInput{
+		Table: "tbl",
+		Key:   map[string]any{"pk": "u1", "sk": "info"},
+		Actions: []driver.UpdateAction{
+			{Action: "SET", Field: "name", Value: "Alice Smith"},
+			{Action: "SET", Field: "email", Value: "alice@test.com"},
+		},
+	})
+	requireNoError(t, err)
+	assertEqual(t, "Alice Smith", updated["name"])
+	assertEqual(t, "alice@test.com", updated["email"])
+	assertEqual(t, 30, updated["age"])
+}
+
+func TestUpdateItemRemoveFields(t *testing.T) {
+	m := newTestMock()
+	ctx := context.Background()
+	createTestTable(m, "tbl")
+
+	_ = m.PutItem(ctx, "tbl", map[string]any{"pk": "u1", "sk": "info", "name": "Alice", "city": "NYC"})
+
+	updated, err := m.UpdateItem(ctx, driver.UpdateItemInput{
+		Table: "tbl",
+		Key:   map[string]any{"pk": "u1", "sk": "info"},
+		Actions: []driver.UpdateAction{
+			{Action: "REMOVE", Field: "city"},
+		},
+	})
+	requireNoError(t, err)
+
+	if _, has := updated["city"]; has {
+		t.Error("expected city to be removed")
+	}
+
+	assertEqual(t, "Alice", updated["name"])
+}
+
+func TestUpdateItemSetAndRemoveCombined(t *testing.T) {
+	m := newTestMock()
+	ctx := context.Background()
+	createTestTable(m, "tbl")
+
+	_ = m.PutItem(ctx, "tbl", map[string]any{"pk": "u1", "sk": "info", "name": "Alice", "old_field": "remove_me"})
+
+	updated, err := m.UpdateItem(ctx, driver.UpdateItemInput{
+		Table: "tbl",
+		Key:   map[string]any{"pk": "u1", "sk": "info"},
+		Actions: []driver.UpdateAction{
+			{Action: "SET", Field: "name", Value: "Bob"},
+			{Action: "REMOVE", Field: "old_field"},
+		},
+	})
+	requireNoError(t, err)
+	assertEqual(t, "Bob", updated["name"])
+
+	if _, has := updated["old_field"]; has {
+		t.Error("expected old_field to be removed")
+	}
+}
+
+func TestUpdateItemPersistsChanges(t *testing.T) {
+	m := newTestMock()
+	ctx := context.Background()
+	createTestTable(m, "tbl")
+
+	_ = m.PutItem(ctx, "tbl", map[string]any{"pk": "u1", "sk": "info", "v": "old"})
+
+	_, err := m.UpdateItem(ctx, driver.UpdateItemInput{
+		Table: "tbl",
+		Key:   map[string]any{"pk": "u1", "sk": "info"},
+		Actions: []driver.UpdateAction{
+			{Action: "SET", Field: "v", Value: "new"},
+		},
+	})
+	requireNoError(t, err)
+
+	got, err := m.GetItem(ctx, "tbl", map[string]any{"pk": "u1", "sk": "info"})
+	requireNoError(t, err)
+	assertEqual(t, "new", got["v"])
+}
+
+func TestUpdateItemTableNotFound(t *testing.T) {
+	m := newTestMock()
+	ctx := context.Background()
+
+	_, err := m.UpdateItem(ctx, driver.UpdateItemInput{
+		Table:   "nonexistent",
+		Key:     map[string]any{"pk": "x"},
+		Actions: []driver.UpdateAction{{Action: "SET", Field: "v", Value: 1}},
+	})
+	assertError(t, err, true)
+}
+
+func TestUpdateItemItemNotFound(t *testing.T) {
+	m := newTestMock()
+	ctx := context.Background()
+	createTestTable(m, "tbl")
+
+	_, err := m.UpdateItem(ctx, driver.UpdateItemInput{
+		Table:   "tbl",
+		Key:     map[string]any{"pk": "missing", "sk": "missing"},
+		Actions: []driver.UpdateAction{{Action: "SET", Field: "v", Value: 1}},
+	})
+	assertError(t, err, true)
+}
+
+func TestUpdateItemInvalidAction(t *testing.T) {
+	m := newTestMock()
+	ctx := context.Background()
+	createTestTable(m, "tbl")
+
+	_ = m.PutItem(ctx, "tbl", map[string]any{"pk": "u1", "sk": "info", "v": 1})
+
+	_, err := m.UpdateItem(ctx, driver.UpdateItemInput{
+		Table:   "tbl",
+		Key:     map[string]any{"pk": "u1", "sk": "info"},
+		Actions: []driver.UpdateAction{{Action: "ADD", Field: "v", Value: 1}},
+	})
+	assertError(t, err, true)
+}
+
+func TestUpdateItemEmitsStreamRecord(t *testing.T) {
+	m := newTestMock()
+	ctx := context.Background()
+	createTestTable(m, "tbl")
+
+	_ = m.UpdateStreamConfig(ctx, "tbl", driver.StreamConfig{
+		Enabled: true, ViewType: "NEW_AND_OLD_IMAGES",
+	})
+
+	_ = m.PutItem(ctx, "tbl", map[string]any{"pk": "u1", "sk": "info", "val": "old"})
+
+	_, err := m.UpdateItem(ctx, driver.UpdateItemInput{
+		Table: "tbl",
+		Key:   map[string]any{"pk": "u1", "sk": "info"},
+		Actions: []driver.UpdateAction{
+			{Action: "SET", Field: "val", Value: "new"},
+		},
+	})
+	requireNoError(t, err)
+
+	iter, err := m.GetStreamRecords(ctx, "tbl", 10, "")
+	requireNoError(t, err)
+	assertEqual(t, 2, len(iter.Records))
+	assertEqual(t, "MODIFY", iter.Records[1].EventType)
+	assertEqual(t, "old", iter.Records[1].OldImage["val"])
+	assertEqual(t, "new", iter.Records[1].NewImage["val"])
+}
+
+func TestUpdateItemEmitsMetrics(t *testing.T) {
+	fc := config.NewFakeClock(time.Date(2025, 1, 1, 0, 0, 0, 0, time.UTC))
+	opts := config.NewOptions(config.WithClock(fc))
+	m := New(opts)
+	ctx := context.Background()
+
+	cw := cloudwatch.New(opts)
+	m.SetMonitoring(cw)
+	createTestTable(m, "tbl")
+
+	_ = m.PutItem(ctx, "tbl", map[string]any{"pk": "u1", "sk": "info"})
+
+	_, err := m.UpdateItem(ctx, driver.UpdateItemInput{
+		Table: "tbl",
+		Key:   map[string]any{"pk": "u1", "sk": "info"},
+		Actions: []driver.UpdateAction{
+			{Action: "SET", Field: "v", Value: "x"},
+		},
+	})
+	requireNoError(t, err)
+
+	result, err := cw.GetMetricData(ctx, mondriver.GetMetricInput{
+		Namespace:  "AWS/DynamoDB",
+		MetricName: "ConsumedWriteCapacityUnits",
+		Dimensions: map[string]string{"TableName": "tbl"},
+		StartTime:  fc.Now().Add(-1 * time.Hour),
+		EndTime:    fc.Now().Add(1 * time.Hour),
+		Period:     60,
+		Stat:       "Sum",
+	})
+	requireNoError(t, err)
+	assertEqual(t, true, len(result.Values) > 0)
+}

providers/azure/cosmosdb/cosmosdb.go

@@ -206,6 +206,48 @@ func (m *Mock) GetItem(_ context.Context, table string, key map[string]any) (map
 	return item, nil
 }
 
+// UpdateItem applies partial updates to an existing document in a container.
+func (m *Mock) UpdateItem(_ context.Context, input driver.UpdateItemInput) (map[string]any, error) {
+	m.mu.Lock()
+
+	td, exists := m.tables[input.Table]
+	if !exists {
+		m.mu.Unlock()
+		return nil, cerrors.Newf(cerrors.NotFound, "container %s not found", input.Table)
+	}
+
+	k := itemKey(td.config, input.Key)
+	item, ok := td.items.Get(k)
+
+	if !ok {
+		m.mu.Unlock()
+		return nil, cerrors.New(cerrors.NotFound, "item not found")
+	}
+
+	oldItem := copyItem(item)
+	updated := copyItem(item)
+
+	for _, action := range input.Actions {
+		switch action.Action {
+		case "SET":
+			updated[action.Field] = action.Value
+		case "REMOVE":
+			delete(updated, action.Field)
+		default:
+			m.mu.Unlock()
+			return nil, cerrors.Newf(cerrors.InvalidArgument, "unsupported action: %s", action.Action)
+		}
+	}
+
+	td.items.Set(k, updated)
+	m.recordStreamEvent(td, oldItem, updated, true)
+	m.mu.Unlock()
+
+	m.emitMetric(input.Table, map[string]float64{"TotalRequests": 1, "TotalRequestUnits": 1})
+
+	return updated, nil
+}
+
 // DeleteItem deletes an item from a container by key.
 func (m *Mock) DeleteItem(_ context.Context, table string, key map[string]any) error {
 	m.mu.Lock()