cgimport/sqlite/iter.go

package sqlite

import (
	"context"
	"database/sql"
	"reflect"

	"github.com/wecisecode/util/merrs"
)

type Row struct {
	data map[string]interface{}
}

func (rr *Row) Data() map[string]interface{} {
	return rr.data
}

// 迭代器
type Iter struct {
	ctx  context.Context
	rows *sql.Rows
	cols []*sql.ColumnType
}

// 获取下一数据
func (iter *Iter) NextRow() (row *Row, err error) {
	if iter.ctx != nil {
		err = iter.ctx.Err()
		if err != nil {
			return
		}
	}
	if !iter.rows.Next() {
		iter.Close()
		return
	}
	ftypes := make([]string, len(iter.cols))
	values := make([]interface{}, len(iter.cols))
	for i := range values {
		coltype := iter.cols[i]
		scantype := coltype.ScanType()
		ftypes[i] = coltype.DatabaseTypeName()
		if scantype == nil {
			bs := []byte{}
			values[i] = &bs
		} else {
			v := reflect.New(scantype).Interface()
			values[i] = &v
		}
	}
	if err = iter.rows.Scan(values...); err != nil {
		err = merrs.NormalError.NewCause(err)
		iter.Close()
		return
	}
	m := make(map[string]interface{})
	for i, v := range values {
		fieldname := iter.cols[i].Name()
		ftype := ftypes[i]
		if iter.cols[i].ScanType() == nil {
			v = nil
		} else {
			v = reflect.ValueOf(v).Elem().Interface()
			v = SQLValueDecode(ftype, v)
		}
		m[fieldname] = v
	}
	row = &Row{m}
	return
}

// 关闭迭代器
func (iter *Iter) Close() (err error) {
	return iter.rows.Close()
}

// 一次获取多行数据
func (iter *Iter) NextRows(n int) (rows []*Row, err error) {
	rows = []*Row{}
	for {
		rd, e := iter.NextRow()
		if e != nil {
			err = merrs.NewError(e)
			return
		}
		if rd == nil || n > 0 && len(rows) >= n {
			return
		}
		rows = append(rows, rd)
	}
}

// 一次获取多行map形式的数据
func (iter *Iter) NextMaps(n int) (rds []map[string]interface{}, err error) {
	rds = []map[string]interface{}{}
	for {
		rd, e := iter.NextRow()
		if e != nil {
			err = merrs.NormalError.NewCause(e)
			return
		}
		if rd == nil || n > 0 && len(rds) >= n {
			return
		}
		rds = append(rds, rd.Data())
	}
}

// 一次获取所有数据
func (iter *Iter) AllRows() (rows []*Row, err error) {
	defer iter.Close()
	return iter.NextRows(-1)
}

// 一次获取所有map形式的数据
func (iter *Iter) AllMaps() (rds []map[string]interface{}, err error) {
	defer iter.Close()
	return iter.NextMaps(-1)
}