286 lines
7.4 KiB
Go
286 lines
7.4 KiB
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 import 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
|
||
}
|
||
|
||
func (b *Mapper) 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
|
||
}
|
||
|
||
// 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...)
|
||
}
|