db.go

/*
Package txdb is a single transaction based [database/sql/driver] implementation.
When the connection is opened, it starts a transaction and all operations
performed on the returned [database/sql.DB] will be within that transaction. If
concurrent actions are performed, the lock is acquired and connection is always
released the statements and rows are not holding the connection.

Why is it useful? A very basic use case would be if you want to make functional
tests, you can prepare a test database and within each test you do not have to
reload a database. All tests are isolated within a transaction and execute fast.
And you do not have to interface your [database/sql.DB] reference in your code,
txdb is like a standard [database/sql/driver.Driver].

This driver supports any [database/sql/driver.Driver] connection to be opened.
You can register txdb for different drivers and have it under different driver
names. Under the hood whenever a txdb driver is opened, it attempts to open a
real connection and starts transaction. When close is called, it rollbacks
transaction leaving your prepared test database in the same state as before.

Example, assuming you have a mysql database called txdb_test and a table users with a
username:

	package main

	import (
		"database/sql"
		"log"

		"github.com/DATA-DOG/go-txdb"
		_ "github.com/go-sql-driver/mysql"
	)

	func init() {
		// we register an sql driver named "txdb"
		txdb.Register("txdb", "mysql", "root@/txdb_test")
	}

	func main() {
		// dsn serves as an unique identifier for connection pool
		db, err := sql.Open("txdb", "identifier")
		if err != nil {
			log.Fatal(err)
		}
		defer db.Close()

		if _, err := db.Exec(`INSERT INTO users(username) VALUES("gopher")`); err != nil {
			log.Fatal(err)
		}
	}

Every time you will run this application, it will remain in the same state as before.
*/
package txdb

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"fmt"
	"io"
	"reflect"
	"sync"
)

// New returns a [database/sql/driver.Connector], which can be passed to
// [database/sql.OpenDB]. This can be used in place of [Register].
// It takes the same arguments as [Register], with the omission of name.
func New(drv, dsn string, options ...func(*conn) error) driver.Connector {
	return &txConnector{
		driver: &TxDriver{
			dsn:     dsn,
			drv:     drv,
			conns:   make(map[string]*conn),
			options: options,
		},
		name: "connector",
	}
}

// Register registers a txdb sql driver under the given sql driver name
// which can be used to open a single transaction based database connection.
//
// When Open is called any number of times it returns the same transaction
// connection.
//
// Any Begin, Commit calls will not start or close the transaction.
// Instead the savepoint will be created, released, or rolled back.
// If your SQL driver does not support save points, use nil
// for the SavePointOption argument. If driver has non-default
// save point logic, you can override the default with SavePointOption.
//
// When [Close] is called, the transaction is rolled back.
//
// The drv dsn are passed to [databse/sql.Open].
//
// Note: if you open a secondary database, make sure to differentiate
// the dsn string when opening the [driver/sql.DB]. The transaction will be
// isolated within that dsn.
func Register(name, drv, dsn string, options ...func(*conn) error) {
	sql.Register(name, &TxDriver{
		dsn:     dsn,
		drv:     drv,
		conns:   make(map[string]*conn),
		options: options,
	})
}

type conn struct {
	sync.Mutex
	tx        *sql.Tx
	dsn       string
	opened    uint
	drv       *TxDriver
	saves     uint
	savePoint SavePoint

	cancel func()
	ctx    interface{ Done() <-chan struct{} }
}

// TxDriver is a [database/sql/driver.Driver] implementation which runs on
// single transaction. When [database/sql.DB.Close] is called, transaction is
// rolled back.
type TxDriver struct {
	sync.Mutex
	db       *sql.DB
	realConn driver.Conn // Meant to be used as NamedValueChecker
	conns    map[string]*conn
	options  []func(*conn) error

	drv string
	dsn string
}

var (
	_ driver.Driver        = (*TxDriver)(nil)
	_ driver.DriverContext = (*TxDriver)(nil)
)

type txConnector struct {
	driver *TxDriver
	name   string
}

var _ driver.Connector = (*txConnector)(nil)

// Connect satisfies the [database/sql/driver.Connector] interface.
func (c *txConnector) Connect(context.Context) (driver.Conn, error) {
	// The DSN passed here doesn't matter, since it's only used to disambiguate
	// connections, but that disambiguation happens in the call to New() when
	// used through the driver.Connector interface.
	return c.driver.Open(c.name)
}

// Driver satisfies the [database/sql/driver.Connector] interface.
func (c *txConnector) Driver() driver.Driver {
	return c.driver
}

func (d *TxDriver) DB() *sql.DB {
	return d.db
}

// OpenConnector satisfies the [database/sql/driver.DriverContext] interface.
func (d *TxDriver) OpenConnector(name string) (driver.Connector, error) {
	return &txConnector{
		driver: d,
		name:   name,
	}, nil
}

func (d *TxDriver) Open(dsn string) (driver.Conn, error) {
	d.Lock()
	defer d.Unlock()
	// first open a real database connection
	if d.db == nil {
		db, err := sql.Open(d.drv, d.dsn)
		if err != nil {
			return nil, err
		}
		d.db = db

		realConn, err := db.Driver().Open(d.dsn)
		if err != nil {
			return nil, err
		}
		d.realConn = realConn
	}
	c, ok := d.conns[dsn]
	if !ok {
		c = &conn{
			dsn:       dsn,
			drv:       d,
			savePoint: &defaultSavePoint{},
			cancel:    func() {},
			ctx:       stubCtx{},
		}
		for _, opt := range d.options {
			if e := opt(c); e != nil {
				return c, e
			}
		}
		d.conns[dsn] = c
	}
	c.opened++ // safe since conn.Close() must acquire driver lock first
	return c, nil
}

func (d *TxDriver) deleteConn(dsn string) error {
	// d must be locked before call
	delete(d.conns, dsn)
	if len(d.conns) == 0 && d.db != nil {
		if err := d.db.Close(); err != nil {
			return err
		}
		d.db = nil
		if err := d.realConn.Close(); err != nil {
			return err
		}
	}
	return nil
}

func (c *conn) beginOnce() (*sql.Tx, error) {
	if c.tx == nil {
		tx, err := c.drv.db.Begin()
		if err != nil {
			return nil, err
		}
		c.tx = tx
	}
	return c.tx, nil
}

func (c *conn) Close() (err error) {
	c.drv.Lock()
	defer c.drv.Unlock()

	c.opened--
	if c.opened == 0 {
		if c.tx != nil {
			err := c.tx.Rollback()
			if err != nil {
				return err
			}
			c.cancel()
			c.tx = nil
		}
		return c.drv.deleteConn(c.dsn)
	}
	return
}

type tx struct {
	id   string
	conn *conn
}

func (c *conn) Begin() (driver.Tx, error) {
	if c.savePoint == nil {
		return &tx{"_", c}, nil // save point is not supported
	}

	c.Lock()
	defer c.Unlock()

	connTx, err := c.beginOnce()
	if err != nil {
		return nil, err
	}

	c.saves++
	id := fmt.Sprintf("tx_%d", c.saves)
	_, err = connTx.Exec(c.savePoint.Create(id))
	if err != nil {
		return nil, err
	}
	return &tx{id, c}, nil
}

func (tx *tx) Commit() error {
	if tx.conn.savePoint == nil {
		return nil // save point is not supported
	}

	tx.conn.Lock()
	defer tx.conn.Unlock()

	connTx, err := tx.conn.beginOnce()
	if err != nil {
		return err
	}

	_, err = connTx.Exec(tx.conn.savePoint.Release(tx.id))
	return err
}

func (tx *tx) Rollback() error {
	if tx.conn.savePoint == nil {
		return nil // save point is not supported
	}

	tx.conn.Lock()
	defer tx.conn.Unlock()

	connTx, err := tx.conn.beginOnce()
	if err != nil {
		return err
	}

	_, err = connTx.Exec(tx.conn.savePoint.Rollback(tx.id))
	return err
}

func (c *conn) Prepare(query string) (driver.Stmt, error) {
	c.Lock()
	defer c.Unlock()

	tx, err := c.beginOnce()
	if err != nil {
		return nil, err
	}

	st, err := tx.Prepare(query)
	if err != nil {
		return nil, err
	}
	return &stmt{st: st}, nil
}

func (c *conn) Exec(query string, args []driver.Value) (driver.Result, error) {
	c.Lock()
	defer c.Unlock()

	tx, err := c.beginOnce()
	if err != nil {
		return nil, err
	}

	return tx.Exec(query, mapArgs(args)...)
}

func mapArgs(args []driver.Value) (res []interface{}) {
	res = make([]interface{}, len(args))
	for i := range args {
		res[i] = args[i]
	}
	return
}

func (c *conn) Query(query string, args []driver.Value) (driver.Rows, error) {
	c.Lock()
	defer c.Unlock()

	tx, err := c.beginOnce()
	if err != nil {
		return nil, err
	}

	// query rows
	rs, err := tx.Query(query, mapArgs(args)...)
	if err != nil {
		return nil, err
	}
	defer rs.Close()

	return buildRows(rs)
}

// Implement the NamedValueChecker interface
func (c *conn) CheckNamedValue(nv *driver.NamedValue) error {
	if nvc, ok := c.drv.realConn.(driver.NamedValueChecker); ok {
		return nvc.CheckNamedValue(nv)
	}

	switch nv.Value.(type) {
	case sql.Out:
		return nil
	default:
		return driver.ErrSkip
	}
}

type stmt struct {
	mu   sync.Mutex
	st   *sql.Stmt
	done chan bool
}

func (s *stmt) Exec(args []driver.Value) (driver.Result, error) {
	dr, err := s.st.Exec(mapArgs(args)...)
	if err != nil {
		s.closeDone(true)
	}
	return dr, err
}

func (s *stmt) NumInput() int {
	return -1
}

func (s *stmt) Close() error {
	s.closeDone(false)
	return s.st.Close()
}

func (s *stmt) Query(args []driver.Value) (driver.Rows, error) {
	rows, err := s.st.Query(mapArgs(args)...)
	if err != nil {
		s.closeDone(true)
		return nil, err
	}
	return buildRows(rows)
}

func (s *stmt) closeDone(withErr bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.done == nil {
		return
	}

	select {
	case s.done <- withErr:
	default:
	}

	close(s.done)
	s.done = nil
}

type rows struct {
	rows     [][]driver.Value
	pos      int
	cols     []string
	colTypes []*sql.ColumnType
}

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

func (r *rows) ColumnTypeDatabaseTypeName(index int) string {
	return r.colTypes[index].DatabaseTypeName()
}

func (r *rows) Next(dest []driver.Value) error {
	r.pos++
	if r.pos > len(r.rows) {
		return io.EOF
	}

	for i, val := range r.rows[r.pos-1] {
		dest[i] = *(val.(*interface{}))
	}

	return nil
}

func (r *rows) Close() error {
	return nil
}

func (r *rows) read(rs *sql.Rows) error {
	var err error
	r.cols, err = rs.Columns()
	if err != nil {
		return err
	}

	r.colTypes, err = rs.ColumnTypes()
	if err != nil {
		return err
	}

	for rs.Next() {
		values := make([]interface{}, len(r.cols))
		for i := range values {
			values[i] = new(interface{})
		}
		if err := rs.Scan(values...); err != nil {
			return err
		}
		row := make([]driver.Value, len(r.cols))
		for i, v := range values {
			row[i] = driver.Value(v)
		}
		r.rows = append(r.rows, row)
	}
	return rs.Err()
}

type rowSets struct {
	sets []*rows
	pos  int
}

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

func (rs *rowSets) ColumnTypeDatabaseTypeName(index int) string {
	return rs.sets[rs.pos].ColumnTypeDatabaseTypeName(index)
}

func (rs *rowSets) ColumnTypeScanType(index int) reflect.Type {
	return rs.sets[rs.pos].colTypes[index].ScanType()
}

func (rs *rowSets) Close() error {
	return nil
}

// advances to next row
func (rs *rowSets) Next(dest []driver.Value) error {
	return rs.sets[rs.pos].Next(dest)
}

type stubCtx struct{}

func (s stubCtx) Done() <-chan struct{} {
	return nil
}