TsDB2/DB/query.go

package DB

import (
	"TsDB2/DB/core"
	"TsDB2/DB/core/types"
	"fmt"
)

// Query 泛型查询构建器
type Query[T any] struct {
	table  *Table[T]
	wheres []whereClause
	orders []orderClause
	limit  int
	offset int
}

type whereClause struct {
	field string
	op    string
	value interface{}
}

type orderClause struct {
	field string
	desc  bool
}

func (q *Query[T]) addWhere(field, op string, value interface{}) *Query[T] {
	q.wheres = append(q.wheres, whereClause{field: field, op: op, value: value})
	return q
}

// Where 添加条件
func (q *Query[T]) Where(field string, op string, value interface{}) *Query[T] {
	q.wheres = append(q.wheres, whereClause{field: field, op: op, value: value})
	return q
}

// OrderBy 排序
func (q *Query[T]) OrderBy(field string, desc ...bool) *Query[T] {
	d := false
	if len(desc) > 0 && desc[0] {
		d = true
	}
	q.orders = append(q.orders, orderClause{field: field, desc: d})
	return q
}

// Limit 限制行数
func (q *Query[T]) Limit(n int) *Query[T] {
	q.limit = n
	return q
}

// Offset 跳过行数
func (q *Query[T]) Offset(n int) *Query[T] {
	q.offset = n
	return q
}

// Find 执行查询返回所有结果
func (q *Query[T]) Find() ([]*T, error) {
	cq := q.buildCoreQuery()
	rows, err := cq.Find()
	if err != nil {
		return nil, err
	}
	result := make([]*T, 0, len(rows))
	for _, row := range rows {
		obj := new(T)
		if err := fieldValuesToStruct(row, q.table.model, obj); err != nil {
			continue
		}
		result = append(result, obj)
	}
	return result, nil
}

// First 返回第一条结果
func (q *Query[T]) First() (*T, error) {
	q.limit = 1
	results, err := q.Find()
	if err != nil {
		return nil, err
	}
	if len(results) == 0 {
		return nil, fmt.Errorf("tsdb: record not found")
	}
	return results[0], nil
}

// Count 返回匹配行数
func (q *Query[T]) Count() (int64, error) {
	cq := q.buildCoreQuery()
	return cq.Count()
}

// Scan 流式遍历
func (q *Query[T]) Scan(callback func(*T) bool) error {
	cq := q.buildCoreQuery()
	return cq.Scan(func(row []types.FieldValue) bool {
		obj := new(T)
		if err := fieldValuesToStruct(row, q.table.model, obj); err != nil {
			return true
		}
		return callback(obj)
	})
}

// buildCoreQuery 构建 core 层的 Query
func (q *Query[T]) buildCoreQuery() *core.Query {
	cq := q.table.coreTable.Query()

	// 构建 WHERE 条件
	if len(q.wheres) > 0 {
		cond := q.buildCondition(q.wheres[0])
		for i := 1; i < len(q.wheres); i++ {
			cond = core.And(cond, q.buildCondition(q.wheres[i]))
		}
		cq = cq.Where(cond)
	}

	// ORDER BY
	for _, o := range q.orders {
		cq = cq.OrderBy(o.field, o.desc)
	}

	// LIMIT / OFFSET
	if q.limit > 0 {
		cq = cq.Limit(q.limit)
	}
	if q.offset > 0 {
		cq = cq.Offset(q.offset)
	}

	return cq
}

// buildCondition 将 whereClause 转换为 core.Condition
func (q *Query[T]) buildCondition(w whereClause) core.Condition {
	ft := q.findFieldType(w.field)
	val := normalizeCondValue(w.value, ft)
	switch w.op {
	case "=", "==", "eq":
		return core.Eq(w.field, val, ft)
	case "!=", "<>":
		return core.Not(core.Eq(w.field, val, ft))
	case ">", "gt":
		return core.Gt(w.field, val, ft)
	case ">=", "ge", "gte":
		return core.Ge(w.field, val, ft)
	case "<", "lt":
		return core.Lt(w.field, val, ft)
	case "<=", "le", "lte":
		return core.Le(w.field, val, ft)
	case "like":
		if s, ok := w.value.(string); ok {
			return core.Like(w.field, s)
		}
		return core.Eq(w.field, val, ft)
	default:
		return core.Eq(w.field, val, ft)
	}
}

// normalizeCondValue 将条件值转换为 core 期望的类型
func normalizeCondValue(v interface{}, ft types.FieldType) interface{} {
	switch ft {
	case types.FieldType_8byte:
		return toUint64(v)
	case types.FieldType_4byte:
		return toInt32(v)
	case types.FieldType_2byte:
		return int16(toUint64(v))
	case types.FieldType_1byte:
		return int8(toUint64(v))
	default:
		// 定长字符串类型: 转换为 []byte (core 的 encodeIndexKey 需要)
		if s, ok := v.(string); ok {
			return toFixedBytes(s, int(types.GetSize(ft)))
		}
		return v
	}
}

// findFieldType 查找字段的 FieldType
// 优先从 DB schema 查找，回退到 struct model
func (q *Query[T]) findFieldType(name string) types.FieldType {
	// 优先从 DB schema 查找
	if ft, ok := q.table.coreTable.GetActiveFieldType(name); ok {
		return ft
	}
	// 回退到 struct model
	for _, fd := range q.table.model.AllFields {
		if fd.ColName == name || fd.Name == name {
			return fd.FieldType
		}
	}
	return types.FieldType_8byte // 默认
}