intake/core/db.go

package core

import (
	"database/sql"
	"runtime"

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

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

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

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

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

func (db *DB) Transact(transaction func(*sql.Tx) 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(tx); err != nil {
		return err
	}
	if err = tx.Commit(); err != nil {
		return err
	}
	return nil
}

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(DB)
	wrapper.ro = ro
	wrapper.rw = rw
	return wrapper, nil
}