TsDB2/DB/core/updater.go

package core

import (
	"git.ycl.pub/yl32940/TsDB2/DB/core/types"
	"git.ycl.pub/yl32940/TsDB2/DB/core/util"
	"encoding/binary"
)

// Updater 更新操作对象
type Updater struct {
	table   *Table
	updates map[string]interface{}
	where   Condition
}

// Update 返回更新操作对象
func (table *Table) Update() *Updater {
	return &Updater{
		table:   table,
		updates: make(map[string]interface{}),
	}
}

// Set 设置更新字段
func (u *Updater) Set(field string, value interface{}) *Updater {
	u.updates[field] = value
	return u
}

// Where 设置更新条件
func (u *Updater) Where(cond Condition) *Updater {
	u.where = cond
	return u
}

// Exec 执行更新
func (u *Updater) Exec() (int64, error) {
	defer u.table.db.lockOp()() // 批量更新全程串行; 内部 scanWithPKID/UpdateRow 同 goroutine 重入
	// 找承载行数据的一级索引名 (实现 RowStoreHandler 者, 本层不认识索引类型)
	pkIdxName := ""
	if pkMeta, _, ok := u.table.indexMeta.rowStoreMeta(); ok {
		pkIdxName = util.BytesToString(pkMeta.IndexName[:])
	}
	if pkIdxName == "" {
		return 0, nil
	}

	// 先收集匹配行的真实 pkID, 再逐行改。
	// 两处都不能沿用旧实现:
	//   - 用 Find() 收集会把每一行都反序列化成 []FieldValue, 万行级的更新白占内存,
	//     这里只物化 pkID;
	//   - pkID 必须取行前缀里的真实值。行数据里没有 pkID 字段, 旧实现只能"取第一个
	//     uint64 字段值当 pkID", 主键列不是行内第一个 8 字节列时会改到另一行上。
	var pkIDs []uint64
	if err := u.table.Query().Where(u.where).scanWithPKID(func(pkID uint64, _ []types.FieldValue) bool {
		pkIDs = append(pkIDs, pkID)
		return true
	}); err != nil {
		return 0, err
	}

	var count int64
	for _, pkID := range pkIDs {
		key := make([]byte, 8)
		binary.BigEndian.PutUint64(key, pkID)
		// 只把本次要改的字段交给 UpdateRow, 其余列由 serializeRow 沿用旧值
		var values []types.FieldValue
		for field, val := range u.updates {
			values = append(values, types.FieldValue{
				FieldName: []byte(field),
				Value:     val,
			})
		}
		// 过去这里写成 `if err == nil { count++ }`, 把失败行悄悄跳过:
		// 调用方会拿到成功的行数, 却不知道有行没更新。必须向上抛出。
		if err := u.table.UpdateRow(pkIdxName, key, values); err != nil {
			return count, err
		}
		count++
	}
	return count, nil
}