tdengine-mapper-go/mapping.go

package tdmap

import (
	"fmt"
	"reflect"
	"strings"
)

func NewMapping() *TdMapping {
	return &TdMapping{}
}

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

type TdMapping struct {
	modelMates Map[string, *StructMate]
}

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

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

	return nil
}

// 提取 struct 内的信息
// @param data 包含 struct 类型的数据
// @return 返回一个 map[表名][]*TableRowMateria， 如果超级表名为空，则表示普通表
func (b *TdMapping) 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 import TableName() string func for struct type: %s", uniqueTypeName)
			}
		}

		mate, ok := b.modelMates.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.FiledDBNameCache[name]

			field := vf.FieldByIndex(mate.FieldIndexCache[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.FiledDBNameCache[name]
			field := vf.FieldByIndex(mate.FieldIndexCache[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
}

func (b *TdMapping) ToInsertSQL(data ...any) (string, error) {
	// 扫描 struct 类型数据
	if err := b.scanStruct(data...); err != nil {
		return "", fmt.Errorf("failed to scan struct: %w", err)
	}

	// 提取 struct 内的信息
	tableMap, err := b.extractStructData(data...)
	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
}