intake/core/db.go

package core

import (
	"database/sql"
	"runtime"

	_ "github.com/mattn/go-sqlite3"
)

type DB interface {
	Query(query string, args ...any) (*sql.Rows, error)
	QueryRow(query string, args ...any) *sql.Row
	Exec(query string, args ...any) (sql.Result, error)
	Prepare(query string) (*sql.Stmt, error)
	Transact(func(DB) error) error
}

type RoRwDb struct {
	ro *sql.DB
	rw *sql.DB
}

func (db *RoRwDb) Query(query string, args ...any) (*sql.Rows, error) {
	return db.ro.Query(query, args...)
}

func (db *RoRwDb) QueryRow(query string, args ...any) *sql.Row {
	return db.ro.QueryRow(query, args...)
}

func (db *RoRwDb) Exec(query string, args ...any) (sql.Result, error) {
	return db.rw.Exec(query, args...)
}

func (db *RoRwDb) Prepare(query string) (*sql.Stmt, error) {
	return db.rw.Prepare(query)
}

func (db *RoRwDb) Transact(transaction func(DB) error) error {
	tx, err := db.rw.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	_, err = tx.Exec("rollback; begin immediate")
	if err != nil {
		return err
	}
	if err = transaction(&TxDb{tx}); err != nil {
		return err
	}
	if err = tx.Commit(); err != nil {
		return err
	}
	return nil
}

type TxDb struct {
	*sql.Tx
}

func (tx *TxDb) Query(query string, args ...any) (*sql.Rows, error) {
	return tx.Tx.Query(query, args...)
}

func (tx *TxDb) QueryRow(query string, args ...any) *sql.Row {
	return tx.Tx.QueryRow(query, args...)
}

func (tx *TxDb) Exec(query string, args ...any) (sql.Result, error) {
	return tx.Tx.Exec(query, args...)
}

func (tx *TxDb) Prepare(query string) (*sql.Stmt, error) {
	return tx.Tx.Prepare(query)
}

func (tx *TxDb) Transact(transaction func(DB) error) error {
	return transaction(tx)
}

func defaultPragma(db *sql.DB) (sql.Result, error) {
	return db.Exec(`
		pragma journal_mode = WAL;
		pragma busy_timeout = 5000;
		pragma synchronous = NORMAL;
		pragma cache_size = 1000000000;
		pragma foreign_keys = true;
		pragma temp_store = memory;
		pragma mmap_size = 3000000000;
	`)
}

func OpenDb(dataSourceName string) (DB, error) {
	ro, err := sql.Open("sqlite3", dataSourceName)
	if err != nil {
		defer ro.Close()
		return nil, err
	}
	ro.SetMaxOpenConns(max(4, runtime.NumCPU()))
	_, err = defaultPragma(ro)
	if err != nil {
		defer ro.Close()
		return nil, err
	}

	rw, err := sql.Open("sqlite3", dataSourceName)
	if err != nil {
		defer ro.Close()
		defer rw.Close()
		return nil, err
	}
	rw.SetMaxOpenConns(1)
	_, err = defaultPragma(rw)
	if err != nil {
		defer ro.Close()
		defer rw.Close()
		return nil, err
	}

	wrapper := new(RoRwDb)
	wrapper.ro = ro
	wrapper.rw = rw
	return wrapper, nil
}