git.sr.ht/~pingoo/stdx@v0.0.0-20240218134121-094174641f6e/sqlitex/db.go

package sqlitex

import (
	"context"
	"database/sql"
	"fmt"
	"runtime"

	"github.com/jmoiron/sqlx"
)

type DB struct {
	writeDB *sqlx.DB
	readDB  *sqlx.DB
}

// ensures at compile time that DB satisfies the Queryer interface
var _ Queryer = (*DB)(nil)

type DBOptions struct {
	BusyTimeout int64
}

type Queryer interface {
	Get(ctx context.Context, dest any, query string, args ...any) error
	Select(ctx context.Context, dest any, query string, args ...any) error
	Query(ctx context.Context, query string, args ...any) (*sql.Rows, error)
	Exec(ctx context.Context, query string, args ...any) (sql.Result, error)
	// Rebind(query string) (ret string)
}

func Open(file string, options *DBOptions) (db *DB, err error) {
	connectionUrl := fmt.Sprintf("%s?_journal_mode=WAL&_busy_timeout=10000&_synchronous=NORMAL&_txlock=immediate&_cache_size=50000000&_foreign_keys=true", file)

	writeDB, err := sqlx.Open("sqlite3", connectionUrl)
	if err != nil {
		err = fmt.Errorf("sqlitex: error openning write connection: %w", err)
		return
	}
	writeDB.SetMaxOpenConns(1)
	err = writeDB.Ping()
	if err != nil {
		err = fmt.Errorf("sqlitex: error establishing write connection: %w", err)
		return
	}

	readDB, err := sqlx.Open("sqlite3", connectionUrl)
	if err != nil {
		err = fmt.Errorf("sqlitex: error openning read connection: %w", err)
		return
	}
	readDB.SetMaxOpenConns(4)
	err = readDB.Ping()
	if err != nil {
		err = fmt.Errorf("sqlitex: error establishing read connection: %w", err)
		return
	}

	db = &DB{
		writeDB: writeDB,
		readDB:  readDB,
	}

	return
}

func (db *DB) Close() (err error) {
	err = db.writeDB.Close()
	if err != nil {
		return
	}

	err = db.readDB.Close()
	if err != nil {
		return
	}

	return
}

// Ping verifies a connection to the database is still alive, establishing a connection if necessary.
// func (db *DB) Ping(ctx context.Context) error {
// 	return db.readDB.PingContext(ctx)
// }

// Stats returns database statistics.
// func (db *DB) Stats() sql.DBStats {
// 	return db.readDB.Stats()
// }

// Begin starts a transaction. The default isolation level is dependent on the driver.
// The provided context is used until the transaction is committed or rolled back. If the context is
// canceled, the sql package will roll back the transaction. Tx.Commit will return an error if the
// context provided to BeginTx is canceled.
// func (db *DB) Begin(ctx context.Context) (Tx, error) {
// 	sqlxTx, err := db.sqlxDB.BeginTxx(ctx, nil)
// 	return &Transaction{sqlxTx}, err
// }

// func (db *DB) Rebind(query string) (ret string) {
// 	return db.sqlxDB.Rebind(query)
// }

// BeginTx starts a transaction.
//
// The provided context is used until the transaction is committed or rolled back. If the context is
// canceled, the sql package will roll back the transaction. Tx.Commit will return an error if the
// context provided to BeginTx is canceled.
//
// The provided TxOptions is optional and may be nil if defaults should be used. If a non-default
// isolation level is used that the driver doesn't support, an error will be returned.
// func (db *DB) BeginTx(ctx context.Context, opts *sql.TxOptions) (Tx, error) {
// 	sqlxTx, err := db.sqlxDB.BeginTxx(ctx, opts)
// 	return &Transaction{sqlxTx}, err
// }

// Exec executes a query without returning any rows. The args are for any placeholder parameters in the query.
func (db *DB) Exec(ctx context.Context, query string, args ...any) (sql.Result, error) {
	return db.writeDB.ExecContext(ctx, query, args...)
}

// Get a single record. Any placeholder parameters are replaced with supplied args. An `ErrNoRows`
// error is returned if the result set is empty.
func (db *DB) Get(ctx context.Context, dest any, query string, args ...any) error {
	return db.readDB.GetContext(ctx, dest, query, args...)
}

// Query executes a query that returns rows, typically a SELECT. The args are for any placeholder
// parameters in the query.
func (db *DB) Query(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
	return db.readDB.QueryContext(ctx, query, args...)
}

// Select an array of records. Any placeholder parameters are replaced with supplied args.
func (db *DB) Select(ctx context.Context, dest any, query string, args ...any) error {
	return db.readDB.SelectContext(ctx, dest, query, args...)
}

func (db *DB) Transaction(ctx context.Context, fn func(tx *Tx) error) (err error) {
	sqlxTx, err := db.writeDB.BeginTxx(ctx, nil)
	if err != nil {
		return err
	}

	tx := &Tx{sqlxTx}

	defer func() {
		if panicErr := recover(); panicErr != nil {
			_, filename, line, _ := runtime.Caller(2)
			rollbackErr := tx.Rollback()
			if rollbackErr != nil {
				err = fmt.Errorf("db: panic (%s:%d) in transaction + rollback: %v, %w", filename, line, panicErr, rollbackErr)
			} else {
				err = fmt.Errorf("db: panic (%s:%d) in transaction: %v", filename, line, panicErr)
			}
		}

		if err == nil {
			err = tx.Commit()
		} else {
			if rollbackErr := tx.Rollback(); rollbackErr != nil {
				err = fmt.Errorf("db: transaction error: %w, rollback err: %w", err, rollbackErr)
			}
		}
	}()

	err = fn(tx)

	return err
}

// Transaction is wrapper of `sqlx.Tx` which implements `Tx`
type Tx struct {
	sqlxTx *sqlx.Tx
}

// ensure at compile time that Tx satisfies the Queryer interface
var _ Queryer = (*Tx)(nil)

// Commit commits the transaction.
func (tx *Tx) Commit() error {
	return tx.sqlxTx.Commit()
}

// Rollback aborts the transaction.
func (tx *Tx) Rollback() error {
	return tx.sqlxTx.Rollback()
}

// Exec executes a query without returning any rows. The args are for any placeholder parameters in the query.
func (tx *Tx) Exec(ctx context.Context, query string, args ...any) (sql.Result, error) {
	return tx.sqlxTx.ExecContext(ctx, query, args...)
}

// Get a single record. Any placeholder parameters are replaced with supplied args. An `ErrNoRows`
// error is returned if the result set is empty.
func (tx *Tx) Get(ctx context.Context, dest any, query string, args ...any) error {
	return tx.sqlxTx.GetContext(ctx, dest, query, args...)
}

// Query executes a query that returns rows, typically a SELECT. The args are for any placeholder
// parameters in the query.
func (tx *Tx) Query(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
	return tx.sqlxTx.QueryContext(ctx, query, args...)
}

// Select an array of records. Any placeholder parameters are replaced with supplied args.
func (tx *Tx) Select(ctx context.Context, dest any, query string, args ...any) error {
	return tx.sqlxTx.SelectContext(ctx, dest, query, args...)
}

// func (tx *Tx) Rebind(query string) (ret string) {
// 	return tx.sqlxTx.Rebind(query)
// }