db/exec/scanner.go

package exec

import (
	"database/sql"
	"encoding/json"
	"reflect"

	"git.fsdpf.net/go/db/internal/errors"
	"git.fsdpf.net/go/db/internal/util"
)

type (
	// Scanner knows how to scan sql.Rows into structs.
	Scanner interface {
		Next() bool
		ScanStruct(i interface{}) error
		ScanStructs(i interface{}) error
		ScanVal(i interface{}) error
		ScanVals(i interface{}) error
		GetRecord() (map[string]any, error)
		GetRecords() ([]map[string]any, error)
		Close() error
		Err() error
	}

	scanner struct {
		rows      *sql.Rows
		columnMap util.ColumnMap
		columns   []string
	}
)

func unableToFindFieldError(col string) error {
	return errors.New(`unable to find corresponding field to column "%s" returned by query`, col)
}

// NewScanner returns a scanner that can be used for scanning rows into structs.
func NewScanner(rows *sql.Rows) Scanner {
	return &scanner{rows: rows}
}

// Next prepares the next row for Scanning. See sql.Rows#Next for more
// information.
func (s *scanner) Next() bool {
	return s.rows.Next()
}

// Err returns the error, if any that was encountered during iteration. See
// sql.Rows#Err for more information.
func (s *scanner) Err() error {
	return s.rows.Err()
}

func (s *scanner) GetRecords() ([]map[string]any, error) {
	records := []map[string]any{}
	for s.Next() {
		row, err := s.GetRecord()
		if err != nil {
			return records, err
		}
		records = append(records, row)
	}

	return records, s.Err()
}

func (s *scanner) GetRecord() (record map[string]any, err error) {
	// Setup columns, but only once.
	if s.columns == nil || s.columnMap == nil {
		colsType, err := s.rows.ColumnTypes()
		if err != nil {
			return nil, err
		}

		s.columns = make([]string, len(colsType))
		s.columnMap = util.ColumnMap{}
		for i := range colsType {
			s.columns[i] = colsType[i].Name()

			typ := colsType[i].ScanType()

			// 强制日期时间为字符串
			if typ == reflect.TypeOf(sql.NullTime{}) {
				typ = reflect.TypeOf(sql.Null[[]uint8]{})
			}

			s.columnMap[colsType[i].Name()] = util.ColumnData{
				ColumnName: colsType[i].Name(),
				GoType:     typ,
			}
		}
	}

	scans := make([]interface{}, len(s.columns))

	for i, col := range s.columns {
		scans[i] = reflect.New(s.columnMap[col].GoType).Interface()
	}

	if err := s.rows.Scan(scans...); err != nil {
		return nil, err
	}

	record = make(map[string]interface{}, len(s.columns))
	for i, col := range s.columns {
		var vv any
		switch v := scans[i].(type) {
		case *sql.Null[[]uint8]: // 强制日期时间为字符串
			if vv, err = v.Value(); err == nil && vv != nil {
				vv = string(vv.([]uint8))
			}
		case *sql.RawBytes:
			if len(*v) > 1 {
				if rune((*v)[0]) == rune('[') || rune((*v)[0]) == rune('{') {
					err = json.Unmarshal(*v, &vv)
				} else {
					vv = string(*v)
				}
			} else {
				vv = string(*v)
			}
		default:
			vv = reflect.Indirect(reflect.ValueOf(v)).Interface()
		}
		if err != nil {
			return
		}
		record[col] = vv
	}

	return record, s.Err()
}

// ScanStruct will scan the current row into i.
func (s *scanner) ScanStruct(i interface{}) error {
	// Setup columnMap and columns, but only once.
	if s.columnMap == nil || s.columns == nil {
		cm, err := util.GetColumnMap(i)
		if err != nil {
			return err
		}

		cols, err := s.rows.Columns()
		if err != nil {
			return err
		}

		// 补全未知字段类型
		if len(cols) != len(cm) {
			colTypes, err := s.rows.ColumnTypes()
			if err != nil {
				return err
			}
			for _, t := range colTypes {
				if _, ok := cm[t.Name()]; !ok {
					cm[t.Name()] = util.ColumnData{
						ColumnName: t.Name(),
						GoType:     t.ScanType(),
					}
				}
			}
		}

		s.columnMap = cm
		s.columns = cols
	}

	scans, err := createColumnScans(s.columns, s.columnMap)

	if err != nil {
		return err
	}

	if err := s.rows.Scan(scans...); err != nil {
		return err
	}

	record := map[string]interface{}{}
	for index, col := range s.columns {
		record[col] = scans[index]
	}

	util.AssignStructVals(i, record, s.columnMap)

	return s.Err()
}

// ScanStructs scans results in slice of structs
func (s *scanner) ScanStructs(i interface{}) error {
	val, err := checkScanStructsTarget(i)
	if err != nil {
		return err
	}
	return s.scanIntoSlice(val, func(i interface{}) error {
		return s.ScanStruct(i)
	})
}

// ScanVal will scan the current row and column into i.
func (s *scanner) ScanVal(i interface{}) error {
	if err := s.rows.Scan(i); err != nil {
		return err
	}

	return s.Err()
}

// ScanStructs scans results in slice of values
func (s *scanner) ScanVals(i interface{}) error {
	val, err := checkScanValsTarget(i)
	if err != nil {
		return err
	}
	return s.scanIntoSlice(val, func(i interface{}) error {
		return s.ScanVal(i)
	})
}

// Close closes the Rows, preventing further enumeration. See sql.Rows#Close
// for more info.
func (s *scanner) Close() error {
	return s.rows.Close()
}

func (s *scanner) scanIntoSlice(val reflect.Value, it func(i interface{}) error) error {
	elemType := util.GetSliceElementType(val)

	for s.Next() {
		row := reflect.New(elemType)
		if rowErr := it(row.Interface()); rowErr != nil {
			return rowErr
		}
		util.AppendSliceElement(val, row)
	}

	return s.Err()
}

func checkScanStructsTarget(i interface{}) (reflect.Value, error) {
	val := reflect.ValueOf(i)
	if !util.IsPointer(val.Kind()) {
		return val, errUnsupportedScanStructsType
	}
	val = reflect.Indirect(val)
	if !util.IsSlice(val.Kind()) {
		return val, errUnsupportedScanStructsType
	}
	return val, nil
}

func checkScanValsTarget(i interface{}) (reflect.Value, error) {
	val := reflect.ValueOf(i)
	if !util.IsPointer(val.Kind()) {
		return val, errUnsupportedScanValsType
	}
	val = reflect.Indirect(val)
	if !util.IsSlice(val.Kind()) {
		return val, errUnsupportedScanValsType
	}
	return val, nil
}

func createColumnScans(cols []string, cm util.ColumnMap) (scans []interface{}, err error) {
	scans = make([]interface{}, 0, len(cols))

	for _, col := range cols {
		data, ok := cm[col]

		if !ok {
			return scans, unableToFindFieldError(col)
		}
		// 处理 converting NULL to string is unsupported
		// 前面将 string 和 int 类型转为了 *string 和 *int
		switch data.GoType.Kind() {
		case reflect.String,
			reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64,
			reflect.Float32, reflect.Float64,
			reflect.Bool:
			scans = append(scans, reflect.New(reflect.PointerTo(data.GoType)).Interface())
		case reflect.Map, reflect.Slice, reflect.Struct:
			scans = append(scans, reflect.New(reflect.PointerTo(reflect.TypeOf(json.RawMessage{}))).Interface())
		default:
			scans = append(scans, reflect.New(data.GoType).Interface())
		}
	}

	return scans, nil
}