go-sqlmock/rows.go

package sqlmock

import (
	"database/sql/driver"
	"encoding/csv"
	"io"
	"strings"
)

// CSVColumnParser is a function which converts trimmed csv
// column string to a []byte representation. currently
// transforms NULL to nil
var CSVColumnParser = func(s string) []byte {
	switch {
	case strings.ToLower(s) == "null":
		return nil
	}
	return []byte(s)
}

// Rows interface allows to construct rows
// which also satisfies database/sql/driver.Rows interface
type Rows interface {
	// composed interface, supports sql driver.Rows
	driver.Rows

	// AddRow composed from database driver.Value slice
	// return the same instance to perform subsequent actions.
	// Note that the number of values must match the number
	// of columns
	AddRow(columns ...driver.Value) Rows

	// FromCSVString build rows from csv string.
	// return the same instance to perform subsequent actions.
	// Note that the number of values must match the number
	// of columns
	FromCSVString(s string) Rows

	// RowError allows to set an error
	// which will be returned when a given
	// row number is read
	RowError(row int, err error) Rows

	// CloseError allows to set an error
	// which will be returned by rows.Close
	// function.
	//
	// The close error will be triggered only in cases
	// when rows.Next() EOF was not yet reached, that is
	// a default sql library behavior
	CloseError(err error) Rows
}

type rows struct {
	cols     []string
	rows     [][]driver.Value
	pos      int
	nextErr  map[int]error
	closeErr error
}

func (r *rows) Columns() []string {
	return r.cols
}

func (r *rows) Close() error {
	return r.closeErr
}

// advances to next row
func (r *rows) Next(dest []driver.Value) error {
	r.pos++
	if r.pos > len(r.rows) {
		return io.EOF // per interface spec
	}

	for i, col := range r.rows[r.pos-1] {
		dest[i] = col
	}

	return r.nextErr[r.pos-1]
}

// NewRows allows Rows to be created from a
// sql driver.Value slice or from the CSV string and
// to be used as sql driver.Rows
func NewRows(columns []string) Rows {
	return &rows{cols: columns, nextErr: make(map[int]error)}
}

func (r *rows) CloseError(err error) Rows {
	r.closeErr = err
	return r
}

func (r *rows) RowError(row int, err error) Rows {
	r.nextErr[row] = err
	return r
}

func (r *rows) AddRow(values ...driver.Value) Rows {
	if len(values) != len(r.cols) {
		panic("Expected number of values to match number of columns")
	}

	row := make([]driver.Value, len(r.cols))
	for i, v := range values {
		row[i] = v
	}

	r.rows = append(r.rows, row)
	return r
}

func (r *rows) FromCSVString(s string) Rows {
	res := strings.NewReader(strings.TrimSpace(s))
	csvReader := csv.NewReader(res)

	for {
		res, err := csvReader.Read()
		if err != nil || res == nil {
			break
		}

		row := make([]driver.Value, len(r.cols))
		for i, v := range res {
			row[i] = CSVColumnParser(strings.TrimSpace(v))
		}
		r.rows = append(r.rows, row)
	}
	return r
}