rows.go

package pgxmock

import (
	"encoding/csv"
	"encoding/json"
	"fmt"
	"reflect"
	"strings"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgtype"
)

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

// connRow implements the Row interface for Conn.QueryRow.
type connRow rowSets

func (r *connRow) Scan(dest ...any) (err error) {
	rows := (*rowSets)(r)

	if rows.Err() != nil {
		return rows.Err()
	}

	for _, d := range dest {
		if _, ok := d.(*pgtype.DriverBytes); ok {
			rows.Close()
			return fmt.Errorf("cannot scan into *pgtype.DriverBytes from QueryRow")
		}
	}

	if !rows.Next() {
		if rows.Err() == nil {
			return pgx.ErrNoRows
		}
		return rows.Err()
	}

	_ = rows.Scan(dest...)
	rows.Close()
	return rows.Err()
}

type rowSets struct {
	sets     []*Rows
	RowSetNo int
	ex       *ExpectedQuery
}

func (rs *rowSets) Conn() *pgx.Conn {
	return nil
}

func (rs *rowSets) Err() error {
	r := rs.sets[rs.RowSetNo]
	return r.nextErr[r.recNo-1]
}

func (rs *rowSets) CommandTag() pgconn.CommandTag {
	return rs.sets[rs.RowSetNo].commandTag
}

func (rs *rowSets) FieldDescriptions() []pgconn.FieldDescription {
	return rs.sets[rs.RowSetNo].defs
}

// func (rs *rowSets) Columns() []string {
// 	return rs.sets[rs.pos].cols
// }

func (rs *rowSets) Close() {
	if rs.ex != nil {
		rs.ex.rowsWereClosed = true
	}
	// return rs.sets[rs.pos].closeErr
}

// advances to next row
func (rs *rowSets) Next() bool {
	r := rs.sets[rs.RowSetNo]
	r.recNo++
	return r.recNo <= len(r.rows)
}

// Values returns the decoded row values. As with Scan(), it is an error to
// call Values without first calling Next() and checking that it returned
// true.
func (rs *rowSets) Values() ([]interface{}, error) {
	r := rs.sets[rs.RowSetNo]
	return r.rows[r.recNo-1], r.nextErr[r.recNo-1]
}

func (rs *rowSets) Scan(dest ...interface{}) error {
	r := rs.sets[rs.RowSetNo]
	if len(dest) == 1 {
		if rc, ok := dest[0].(pgx.RowScanner); ok {
			return rc.ScanRow(rs)
		}
	}
	if len(dest) != len(r.defs) {
		return fmt.Errorf("Incorrect argument number %d for columns %d", len(dest), len(r.defs))
	}
	if len(r.rows) == 0 {
		return pgx.ErrNoRows
	}
	for i, col := range r.rows[r.recNo-1] {
		if dest[i] == nil {
			//behave compatible with pgx
			continue
		}
		destVal := reflect.ValueOf(dest[i])
		if destVal.Kind() != reflect.Ptr {
			return fmt.Errorf("Destination argument must be a pointer for column %s", r.defs[i].Name)
		}
		if col == nil {
			dest[i] = nil
			continue
		}
		val := reflect.ValueOf(col)
		if _, ok := dest[i].(*interface{}); ok || val.Type().AssignableTo(destVal.Elem().Type()) {
			if destElem := destVal.Elem(); destElem.CanSet() {
				destElem.Set(val)
			} else {
				return fmt.Errorf("Cannot set destination value for column %s", r.defs[i].Name)
			}
		} else {
			// Try to use Scanner interface
			scanner, ok := destVal.Interface().(interface{ Scan(interface{}) error })

			if !ok {
				return fmt.Errorf("Destination kind '%v' not supported for value kind '%v' of column '%s'",
					destVal.Elem().Kind(), val.Kind(), string(r.defs[i].Name))
			}
			if err := scanner.Scan(val.Interface()); err != nil {
				return fmt.Errorf("Scanning value error for column '%s': %w", string(r.defs[i].Name), err)
			}

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

func (rs *rowSets) RawValues() [][]byte {
	r := rs.sets[rs.RowSetNo]
	dest := make([][]byte, len(r.defs))

	for i, col := range r.rows[r.recNo-1] {
		if b, ok := rawBytes(col); ok {
			dest[i] = b
			continue
		}
		dest[i] = []byte(fmt.Sprintf("%v", col))
	}

	return dest
}

// transforms to debuggable printable string
func (rs *rowSets) String() string {
	if rs.empty() {
		return "\t- returns no data"
	}

	msg := "\t- returns data:\n"
	if len(rs.sets) == 1 {
		for n, row := range rs.sets[0].rows {
			msg += fmt.Sprintf("\t\trow %d - %+v\n", n, row)
		}
		return msg
	}
	for i, set := range rs.sets {
		msg += fmt.Sprintf("\t\tresult set: %d\n", i)
		for n, row := range set.rows {
			msg += fmt.Sprintf("\t\t\trow %d: %+v\n", n, row)
		}
	}
	return msg
}

func (rs *rowSets) empty() bool {
	for _, set := range rs.sets {
		if len(set.rows) > 0 {
			return false
		}
	}
	return true
}

func rawBytes(col interface{}) (_ []byte, ok bool) {
	val, err := json.Marshal(col)
	if err != nil || len(val) == 0 {
		return nil, false
	}
	// Copy the bytes from the mocked row into a shared raw buffer, which we'll replace the content of later
	b := make([]byte, len(val))
	copy(b, val)
	return b, true
}

// Rows is a mocked collection of rows to
// return for Query result
type Rows struct {
	commandTag pgconn.CommandTag
	defs       []pgconn.FieldDescription
	rows       [][]interface{}
	recNo      int
	nextErr    map[int]error
	closeErr   error
}

// NewRows allows Rows to be created from a
// sql interface{} slice or from the CSV string and
// to be used as sql driver.Rows.
// Use pgxmock.NewRows instead if using a custom converter
func NewRows(columns []string) *Rows {
	var coldefs []pgconn.FieldDescription
	for _, column := range columns {
		coldefs = append(coldefs, pgconn.FieldDescription{Name: column})
	}
	return &Rows{
		defs:    coldefs,
		nextErr: make(map[int]error),
	}
}

// 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
func (r *Rows) CloseError(err error) *Rows {
	r.closeErr = err
	return r
}

// RowError allows to set an error
// which will be returned when a given
// row number is read
func (r *Rows) RowError(row int, err error) *Rows {
	r.nextErr[row] = err
	return r
}

// AddRow composed from database interface{} slice
// return the same instance to perform subsequent actions.
// Note that the number of values must match the number
// of columns
func (r *Rows) AddRow(values ...any) *Rows {
	if len(values) != len(r.defs) {
		panic("Expected number of values to match number of columns")
	}

	row := make([]interface{}, len(r.defs))
	copy(row, values)
	r.rows = append(r.rows, row)
	return r
}

// AddRows adds multiple rows composed from any slice and
// returns the same instance to perform subsequent actions.
func (r *Rows) AddRows(values ...[]any) *Rows {
	for _, value := range values {
		r.AddRow(value...)
	}
	return r
}

// AddCommandTag will add a command tag to the result set
func (r *Rows) AddCommandTag(tag pgconn.CommandTag) *Rows {
	r.commandTag = tag
	return r
}

// 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
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([]interface{}, len(r.defs))
		for i, v := range res {
			row[i] = CSVColumnParser(strings.TrimSpace(v))
		}
		r.rows = append(r.rows, row)
	}
	return r
}

// Kind returns rows corresponding to the interface pgx.Rows
// useful for testing entities that implement an interface pgx.RowScanner
func (r *Rows) Kind() pgx.Rows {
	return &rowSets{
		sets: []*Rows{r},
	}
}

// NewRowsWithColumnDefinition return rows with columns metadata
func NewRowsWithColumnDefinition(columns ...pgconn.FieldDescription) *Rows {
	return &Rows{
		defs:    columns,
		nextErr: make(map[int]error),
	}
}