TsDB2/DB/core/updater.go

package core

import (
	"TsDB2/DB/core/page"
	"TsDB2/DB/core/types"
	"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) {
	// 扫描匹配行
	q := u.table.Query().Where(u.where)
	results, err := q.Find()
	if err != nil {
		return 0, err
	}
	// 找主键索引名
	var pkIdxName string
	for i := range u.table.indexMeta.metas {
		m := &u.table.indexMeta.metas[i]
		if m.Flag&0x80 != 0 {
			continue
		}
		if m.IndexType == page.IndexTypePrimary || m.IndexType == page.IndexTypeTime {
			pkIdxName = util.BytesToString(m.IndexName[:])
			break
		}
	}
	if pkIdxName == "" {
		return 0, nil
	}
	var count int64
	for _, row := range results {
		// 找主键字段 (第一个 uint64 字段作为 pkID)
		for _, fv := range row {
			if v, ok := fv.Value.(uint64); ok {
				key := make([]byte, 8)
				binary.BigEndian.PutUint64(key, v)
				// 构建新行值
				var values []types.FieldValue
				for field, val := range u.updates {
					values = append(values, types.FieldValue{
						FieldName: []byte(field),
						Value:     val,
					})
				}
				if err := u.table.UpdateRow(pkIdxName, key, values); err == nil {
					count++
				}
				break
			}
		}
	}
	return count, nil
}