tdengine-mapper-go/mapping.go

package tdmap

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"reflect"
	"strings"
)

func NewMapper() *Mapper {
	return &Mapper{}
}

// TableNamer 定义了一个获取表名的方法。
type TableNamer interface {
	TableName() string
}

type Mapper struct {
	structMateMap SyncMap[string, *StructMeta] // 结构体 元信息缓存 key: struct 的唯一类型名称  value: 元信息
}

func (b *Mapper) scanStruct(data ...any) error {
	for _, datum := range data {
		if reflect.TypeOf(datum).Kind() != reflect.Ptr {
			return fmt.Errorf("need a pointer type: %v", reflect.TypeOf(data))
		}

		// 分析结构体并缓存结果
		if _, err := b.analyzeStructWithCache(datum); err != nil {
			return err
		}
	}

	return nil
}

// analyzeStructWithCache 带缓存的结构体分析
func (b *Mapper) analyzeStructWithCache(data any) (*StructMeta, error) {
	// 获取结构体的类型
	t, v := extractReflectInfo(data)

	// 确保是结构体类型
	if t.Kind() != reflect.Struct {
		return nil, fmt.Errorf("input data is not a struct")
	}

	// 获取类型的唯一标识符
	uniqueTypeName := buildTypeKey(t)

	// 先检查缓存
	if cached, ok := b.structMateMap.Load(uniqueTypeName); ok {
		return cached, nil
	}

	// 初始化结果结构体
	sr := StructMeta{
		UniqueTypeName:    uniqueTypeName,
		DBName2IndexCache: make(map[string][]int, t.NumField()),
		Field2IndexCache:  make(map[string][]int, t.NumField()),
		Field2DBNameCache: make(map[string]string, t.NumField()),
		DBAnnotatedNames:  make([]string, 0, t.NumField()),
		TaggedFieldNames:  make([]string, 0, t.NumField()),
		SuperTableName:    "",
	}

	// 遍历结构体的字段
	for i := 0; i < t.NumField(); i++ {
		field := t.Field(i)
		fieldValue := v.Field(i)

		if field.Anonymous {
			// 处理匿名结构体
			if field.Type.Kind() == reflect.Ptr {
				if fieldValue.IsNil() {
					// 如果指针是 nil，创建一个该类型的零值实例
					zeroValue := reflect.Zero(field.Type.Elem())
					fieldValue = zeroValue
				} else {
					fieldValue = fieldValue.Elem()
				}
			}

			if !fieldValue.CanInterface() {
				continue
			}

			// 递归分析匿名结构体，使用相同的缓存
			subResult, err := b.analyzeStructWithCache(fieldValue.Interface())
			if err != nil {
				return nil, err
			}

			// 合并分析结果
			for k, v := range subResult.DBName2IndexCache {
				sr.DBName2IndexCache[k] = append(sr.DBName2IndexCache[k], i)
				sr.DBName2IndexCache[k] = append(sr.DBName2IndexCache[k], v...)
			}

			for k, v := range subResult.Field2IndexCache {
				sr.Field2IndexCache[k] = append(sr.Field2IndexCache[k], i)
				sr.Field2IndexCache[k] = append(sr.Field2IndexCache[k], v...)
			}
			for k, v := range subResult.Field2DBNameCache {
				sr.Field2DBNameCache[k] = v
			}

			sr.DBAnnotatedNames = append(sr.DBAnnotatedNames, subResult.DBAnnotatedNames...)
			sr.TaggedFieldNames = append(sr.TaggedFieldNames, subResult.TaggedFieldNames...)

			continue
		}

		// 处理普通字段
		columnName := field.Tag.Get("db")
		if columnName == "-" || columnName == "" {
			continue
		}

		if len(sr.Field2IndexCache[field.Name]) > 0 {
			return nil, fmt.Errorf("duplicate field [%s %s `db:%s`]", field.Name, field.Type.Name(), columnName)
		}

		sr.Field2IndexCache[field.Name] = append([]int{}, i)
		sr.DBName2IndexCache[columnName] = append(sr.DBName2IndexCache[columnName], i)
		sr.Field2DBNameCache[field.Name] = columnName

		// 检查字段是否有taos注解
		if field.Tag.Get("taos") == "tag" {
			sr.TaggedFieldNames = append(sr.TaggedFieldNames, field.Name)
		} else {
			sr.DBAnnotatedNames = append(sr.DBAnnotatedNames, field.Name)
		}
	}

	if sTableName, ok := tryGetSuperTableName(data); ok {
		sr.SuperTableName = sTableName
	}

	// 将结果存入缓存
	b.structMateMap.Store(uniqueTypeName, &sr)

	return &sr, nil
}

// 提取 struct 内的信息
// @param data 包含 struct 类型的数据
// @return 返回一个 map[表名][]*TableRowMaterial， 如果超级表名为空，则表示普通表
func (b *Mapper) extractTableMaterial(data ...any) (map[string][]*TableRowMaterial, error) {
	result := make(map[string][]*TableRowMaterial)

	for _, item := range data {
		tf, vf := extractReflectInfo(item)
		uniqueTypeName := buildTypeKey(tf)

		// 获取表名
		var tableName string
		{
			if vf.CanAddr() {
				if a, ok := vf.Addr().Interface().(TableNamer); ok {
					tableName = a.TableName()
				}
			}

			if tableName == "" {
				if a, ok := item.(TableNamer); ok {
					tableName = a.TableName()
				}
			}

			if tableName == "" {
				return nil, fmt.Errorf("not func TableName() string func for struct type: %s", uniqueTypeName)
			}
		}

		mate, ok := b.structMateMap.Load(uniqueTypeName)
		if !ok {
			return nil, fmt.Errorf("not found struct type: %s", uniqueTypeName)
		}

		trMaterial := &TableRowMaterial{
			SuperTableName: mate.SuperTableName,
			TableName:      tableName,
			TagColumns:     make([]string, 0, len(mate.TaggedFieldNames)),
			TagValues:      make([]any, 0, len(mate.TaggedFieldNames)),
			Columns:        make([]string, 0, len(mate.DBAnnotatedNames)),
			Values:         make([]any, 0, len(mate.DBAnnotatedNames)),
		}

		for _, name := range mate.DBAnnotatedNames {
			// db tag 名称 -- 数据库列名
			dbColumn := mate.Field2DBNameCache[name]

			field := vf.FieldByIndex(mate.Field2IndexCache[name])
			for field.Kind() == reflect.Ptr {
				if field.IsNil() {
					break
				}
				field = field.Elem()
			}
			// 字段值
			dbValue := field.Interface()

			trMaterial.Columns = append(trMaterial.Columns, dbColumn)
			trMaterial.Values = append(trMaterial.Values, dbValue)
		}

		for _, name := range mate.TaggedFieldNames {
			// db tag 名称 -- 数据库列名
			tagColumn := mate.Field2DBNameCache[name]
			field := vf.FieldByIndex(mate.Field2IndexCache[name])
			for field.Kind() == reflect.Ptr {
				if field.IsNil() {
					break
				}
				field = field.Elem()
			}
			// 字段值
			tagValue := field.Interface()

			trMaterial.TagColumns = append(trMaterial.TagColumns, tagColumn)
			trMaterial.TagValues = append(trMaterial.TagValues, tagValue)
		}

		// 添加到结果
		key := trMaterial.TableName
		result[key] = append(result[key], trMaterial)
	}
	return result, nil
}

// flattenDataToPtrs 对输入数据中不同类型的数据统一转换为指针形式，并展平切片/数组类型的数据
// 支持输入为任意数量的参数，参数可以是单个结构体指针或结构体，也可以是结构体切片或指针切片
// 处理逻辑包括:
//   1. 将切片类型的参数拆开放入一个统一的一维切片中（展平）
//   2. 跳过其中的 nil 元素，确保数据有效性
//   3. 将所有非指针类型的元素转换为对应的指针类型，方便后续统一处理
func (b *Mapper) flattenDataToPtrs(data ...any) ([]any, error) {
	var flattenedData = make([]any, 0, len(data))

	for _, item := range data {
		if isNil(item) { // 跳过nil值，避免后续反射调用panic
			continue
		}
		itemType := reflect.TypeOf(item)

		if itemType.Kind() == reflect.Ptr || itemType.Kind() == reflect.Struct {
			// 非切片类型直接追加 指针
			flattenedData = append(flattenedData, ensurePtr(item))
			continue
		}

		if itemType.Kind() == reflect.Slice || itemType.Kind() == reflect.Array {
			// 处理切片类型参数，展平成单个元素
			sliceValue := reflect.ValueOf(item)
			for i := 0; i < sliceValue.Len(); i++ {
				elem := sliceValue.Index(i).Interface()
				if isNil(elem) { // 跳过nil元素
					continue
				}
				// 非切片类型直接追加 指针
				flattenedData = append(flattenedData, ensurePtr(elem))
			}
			continue
		}

		return nil, fmt.Errorf("unsupported data type: %s, expected struct, pointer, slice or array", itemType.String())

	}

	return flattenedData, nil
}

func (b *Mapper) ToInsertSQL(data ...any) (string, error) {
	// 展平数据为指针切片
	flattenedData, err := b.flattenDataToPtrs(data...)
	if err != nil {
		return "", fmt.Errorf("failed to flatten data: %w", err)
	}

	if len(flattenedData) == 0 {
		return "", fmt.Errorf("data is empty")
	}

	// 扫描 struct 类型数据
	if err := b.scanStruct(flattenedData...); err != nil {
		return "", fmt.Errorf("failed to analyzeStruct struct: %w", err)
	}

	// 提取 struct 内的信息
	tableMap, err := b.extractTableMaterial(flattenedData...)
	if err != nil {
		return "", fmt.Errorf("failed to extract struct data: %w", err)
	}

	if len(tableMap) == 0 {
		return "", fmt.Errorf("no data to insert")
	}

	// 构建 insert 语句
	/*
		INSERT INTO
		表名1 USING 超级表名1 (tag列名) TAGS(tag值) (列名) VALUES (),()
		表名2 USING 超级表名2 (tag列名)TAGS(tag值) (列名) VALUES (),()
	*/
	var buf strings.Builder
	buf.WriteString("INSERT INTO \n")
	for _, materials := range tableMap {
		if len(materials) == 0 {
			continue
		}

		var (
			rowSql string
			err    error
		)
		if materials[0].SuperTableName == "" {
			// 表名1 (列名) VALUES (),()
			rowSql, err = buildInsertStatementForNormalTable(materials...)
		} else {
			// 表名1 USING 超级表名1 (tag列名) TAGS(tag值) (列名) VALUES (),()
			rowSql, err = buildInsertStatementForSuperTable(materials...)
		}
		if err != nil {
			return "", fmt.Errorf("failed to build insert statement: %w", err)
		}

		buf.WriteString(rowSql)
		buf.WriteString(" \n")
	}

	return buf.String(), nil
}

// ScanRows 扫描多行数据
// A Superior Option: Use ScanRowsWithContext instead.
func (b *Mapper) ScanRows(target any, rows *sql.Rows) error {
	return b.ScanRowsWithContext(context.Background(), target, rows)
}

// ScanRowsWithContext 扫描多行数据
// @param ctx 上下文
// @param target 结构体的指针 或 slice 指针 结构体需要使用 db 标签注解
// @param rows 数据库返回的行数据
// @return 返回错误
func (b *Mapper) ScanRowsWithContext(ctx context.Context, target any, rows *sql.Rows) error {
	// 确保rows不为nil
	if isNil(rows) {
		return errors.New("rows cannot be nil")
	}

	// 获取target的反射类型和值
	vf := reflect.ValueOf(target)

	// 检查target是否为指针类型
	if vf.Kind() != reflect.Ptr {
		return errors.New("target must be a pointer to a struct or a slice of structs")
	}

	if vf.IsNil() {
		return errors.New("target is nil, please pass a pointer to a struct or a slice of structs")
	}

	vf = vf.Elem() //  解引用指针以获取目标值

	switch vf.Kind() {
	case reflect.Struct:
		// 提取 struct 内的信息
		if err := b.scanStruct(target); err != nil {
			return err
		}
		// target是指向单个struct的指针，扫描一行数据
		if !rows.Next() {
			return sql.ErrNoRows
		}
		return b.scanRow(vf, rows)
	case reflect.Slice, reflect.Array:
		// target是slice或array，扫描多行数据
		sliceElementType := vf.Type().Elem()
		if sliceElementType.Kind() != reflect.Ptr || sliceElementType.Elem().Kind() != reflect.Struct {
			return errors.New("target must be a pointer to a slice of structs")
		}

		// 提取 slice 中 struct 内的信息
		if err := b.scanStruct(reflect.New(sliceElementType.Elem()).Interface()); err != nil { // 这里在抛出 不是指针类型
			return err
		}

		for rows.Next() {
			select {
			case <-ctx.Done():
				vf.SetLen(0) // 清空 slice
				return ctx.Err()
			default:
			}

			// 创建slice元素的新实例
			newStruct := reflect.New(sliceElementType.Elem())
			if err := b.scanRow(newStruct, rows); err != nil {
				return err
			}

			// 将新实例添加到slice中
			vf.Set(reflect.Append(vf, newStruct))
		}
		return nil
	default:
		return errors.New("target must be a pointer to a struct or a slice of structs")
	}
}

// scanRow 用于扫描单行数据到结构体
func (b *Mapper) scanRow(target reflect.Value, rows *sql.Rows) error {
	target = reflect.Indirect(target)

	uniqueTypeName := buildTypeKey(target.Type())
	mate, ok := b.structMateMap.Load(uniqueTypeName)
	if !ok {
		return fmt.Errorf("not found struct type mate: %s", uniqueTypeName)
	}

	// 假设我们有一个函数来获取列的值
	columns, err := rows.Columns()
	if err != nil {
		return err
	}

	// 准备目标结构体的地址
	dest := make([]interface{}, len(columns))
	for i, colName := range columns {
		// 拿到缓存的字段索引
		idx, ok := mate.DBName2IndexCache[colName]
		if !ok {
			return fmt.Errorf("no corresponding field found for column %s", colName)
		}
		dest[i] = target.FieldByIndex(idx).Addr().Interface()
	}

	// 扫描数据
	return rows.Scan(dest...)
}

// ScanRowsToMap 扫描多行数据到map
// @param rows 数据库返回的行数据
// @return 返回map数组和错误
func (b *Mapper) ScanRowsToMap(rows *sql.Rows) ([]map[string]any, error) {
	return b.ScanRowsToMapWithContext(context.Background(), rows)
}

// ScanRowsToMapWithContext 扫描多行数据到map
// @param ctx 上下文
// @param rows 数据库返回的行数据
// @return 返回map数组和错误
func (b *Mapper) ScanRowsToMapWithContext(ctx context.Context, rows *sql.Rows) ([]map[string]any, error) {
	columns, err := rows.Columns()
	if err != nil {
		return nil, err
	}
	result := make([]map[string]any, 0)

	for rows.Next() {
		// 检查上下文是否已取消
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}

		// 创建一个切片来存储列的值
		values := make([]any, len(columns))
		for i := range values {
			values[i] = new(any)
		}

		// 扫描行数据到切片中
		if err := rows.Scan(values...); err != nil {
			return nil, err
		}

		// 构建 结果映射
		rowMap := make(map[string]any, len(columns))
		for i := range values {
			column := columns[i]
			value := *(values[i].(*any))
			rowMap[column] = value
		}
		result = append(result, rowMap)
	}
	return result, nil
}