conn.go

package sqlmock

import (
	"context"
	"database/sql/driver"
	"fmt"
)

// Conn implements driver.Conn and routes operations to queued mock expectations.
type Conn struct {
	mock      *Mock
	txStarted bool
}

// Prepare creates a statement bound to this connection for the provided query.
func (c *Conn) Prepare(query string) (driver.Stmt, error) {
	return &Stmt{
		query: query,
		conn:  c,
	}, nil
}

// extractArgs converts driver.NamedValue arguments into a plain value slice.
func extractArgs(args []driver.NamedValue) []any {
	out := make([]any, len(args))
	for i, a := range args {
		out[i] = a.Value
	}

	return out
}

// QueryContext validates and consumes the next query expectation.
//
// When strict transaction mode is enabled, QueryContext requires an active
// transaction started via Begin/BeginTx.
func (c *Conn) QueryContext(
	ctx context.Context,
	query string,
	args []driver.NamedValue,
) (driver.Rows, error) {
	m := c.mock
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.strictTx && !c.txStarted {
		return nil, errWithoutTx("query")
	}

	op, err := m.next()
	if err != nil {
		return nil, err
	}

	q, ok := op.(*ExpectedQuery)
	if !ok {
		return nil, fmt.Errorf("expected %T, got QUERY", op)
	}

	if err := q.match(runtimeOp{
		kind: "query",
		sql:  query,
		args: extractArgs(args),
	}); err != nil {
		return nil, err
	}

	m.consume()
	return q.rows, nil
}

// ExecContext validates and consumes the next exec expectation.
//
// When strict transaction mode is enabled, ExecContext requires an active
// transaction started via Begin/BeginTx.
func (c *Conn) ExecContext(
	ctx context.Context,
	query string,
	args []driver.NamedValue,
) (driver.Result, error) {
	m := c.mock
	m.mu.Lock()
	defer m.mu.Unlock()

	if m.strictTx && !c.txStarted {
		return nil, errWithoutTx("exec")
	}

	op, err := m.next()
	if err != nil {
		return nil, err
	}

	e, ok := op.(*ExpectedExec)
	if !ok {
		return nil, fmt.Errorf("expected %T, got EXEC", op)
	}

	if err := e.match(runtimeOp{
		kind: "exec",
		sql:  query,
		args: extractArgs(args),
	}); err != nil {
		return nil, err
	}

	m.consume()

	if e.err != nil {
		return nil, e.err
	}

	return e.result, nil
}

// Close closes the mock connection.
func (c *Conn) Close() error { return nil }

// Begin starts a transaction using default options.
func (c *Conn) Begin() (driver.Tx, error) {
	return c.BeginTx(context.Background(), driver.TxOptions{})
}

// BeginTx validates and consumes the next begin expectation.
//
// On success, BeginTx marks the connection transaction as active and returns a
// Tx handle.
func (c *Conn) BeginTx(
	ctx context.Context,
	opts driver.TxOptions,
) (driver.Tx, error) {
	m := c.mock
	m.mu.Lock()
	defer m.mu.Unlock()

	op, err := m.next()
	if err != nil {
		return nil, err
	}

	err = op.match(runtimeOp{kind: "begin"})
	if err != nil {
		return nil, err
	}

	m.consume()
	c.txStarted = true
	return &Tx{conn: c}, nil
}

func errWithoutTx(op string) error {
	return fmt.Errorf(
		"%s without active transaction. Use WithStrictTx() if this should be transactional",
		op,
	)
}