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, *StructMate] // 结构体 元信息缓存 key: struct 的唯一类型名称  value: 元信息
}

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

		if mate, err := scan(datum); err != nil {
			return err
		} else {
			b.structMateMap.Store(mate.UniqueTypeName, mate)
		}
	}

	return nil
}

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

	for _, item := range data {
		tf, vf := getReflectTypeAndValue(item)
		uniqueTypeName := getUniqueTypeName(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)
		}

		materia := &TableRowMateria{
			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.Filed2DBNameCache[name]

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

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

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

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

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

// normalizeData 对输入数据进行统一处理和规范化
// 支持输入为任意数量的参数，参数可以是单个结构体指针或结构体，也可以是结构体切片或指针切片
// 处理逻辑包括:
//   1. 将切片类型的参数拆开放入一个统一的一维切片中（展平）
//   2. 跳过其中的 nil 元素，确保数据有效性
//   3. 将所有非指针类型的元素转换为对应的指针类型，方便后续统一处理
func (b *Mapper) normalizeData(data ...any) []any {
	var normalizedData []any

	for _, item := range data {
		if item == nil {
			// 跳过nil值，避免后续反射调用panic
			continue
		}
		itemType := reflect.TypeOf(item)
		if itemType.Kind() == reflect.Slice {
			// 处理切片类型参数，展平成单个元素
			sliceValue := reflect.ValueOf(item)
			for i := 0; i < sliceValue.Len(); i++ {
				elem := sliceValue.Index(i).Interface()
				if elem == nil {
					// 跳过nil元素
					continue
				}
				normalizedData = append(normalizedData, elem)
			}
		} else {
			// 非切片类型直接追加
			normalizedData = append(normalizedData, item)
		}
	}

	// 转换所有元素为指针类型，便于后续一致性处理
	for i, item := range normalizedData {
		if item == nil {
			continue
		}
		itemType := reflect.TypeOf(item)
		if itemType.Kind() != reflect.Ptr {
			// 创建对应类型的指针
			itemPtr := reflect.New(itemType)
			// 设置指针指向的值为元素本身
			itemPtr.Elem().Set(reflect.ValueOf(item))
			// 替换切片中的元素为指针类型
			normalizedData[i] = itemPtr.Interface()
		}
	}
	return normalizedData
}

func (b *Mapper) ToInsertSQL(data ...any) (string, error) {
	// 统一规范化处理输入数据，将切片拆平并转换元素为指针
	normalizedData := b.normalizeData(data...)

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

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

	// 提取 struct 内的信息
	tableMap, err := b.extractStructData(normalizedData...)
	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 rows == nil {
		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:
		// target是slice的指针
		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 := getUniqueTypeName(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, err
}