Benchmarkign sqlite with uuid storage shcmes

bench_test.go

package main

import (
	"database/sql"
	"encoding/binary"
	"fmt"
	"sort"
	"strings"
	"testing"

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

func StorageTypes() []string {
	res := []string{"TEXT", "BLOB"}
	sort.StringSlice(res).Sort()
	return res
}

type UUIDDB interface {
	InsertID(u uuid.UUID) error
	QueryID(u uuid.UUID) error
	// Debug() *sql.DB
	Close() error
}

func NewBasicDB(t testing.TB, storageType string) UUIDDB {
	if sort.SearchStrings(StorageTypes(), storageType) == len(StorageTypes()) {
		t.Fatalf("invalid storage type: %s", storageType)
	}

	var res basicDB
	var err error

	// XXX: Not sure how to do sqlite any other way here
	// res.DB, err = sql.Open("sqlite3", "file::memory:?cache=shared")
	res.DB, err = sql.Open("sqlite3", "file:"+uuid.NewString()+"?mode=memory&cache=shared")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := res.Exec(`CREATE TABLE foo (id ` + storageType + ` PRIMARY KEY);`); err != nil {
		t.Fatal(err)
	}
	res.insertIDStmt, err = res.Prepare(`INSERT INTO foo (id) VALUES (?);`)
	if err != nil {
		t.Fatal(err)
	}
	res.queryIDStmt, err = res.Prepare(`SELECT * FROM foo WHERE id = ?;`)
	if err != nil {
		t.Fatal(err)
	}

	return &res
}

type basicDB struct {
	*sql.DB
	insertIDStmt *sql.Stmt
	queryIDStmt  *sql.Stmt
}

func (db *basicDB) InsertID(u uuid.UUID) error {
	_, err := db.insertIDStmt.Exec(u)
	return err
}
func (db *basicDB) QueryID(u uuid.UUID) error {
	return db.queryIDStmt.QueryRow(&u).Err()
}
func (db *basicDB) Close() error {
	db.DB.Close()
	db.insertIDStmt.Close()
	db.queryIDStmt.Close()
	return nil
}

func NewIntBasedDB(t testing.TB) UUIDDB {
	var res intBasedDB
	var err error

	res.DB, err = sql.Open("sqlite3", "file::memory:?cache=shared")
	if err != nil {
		t.Fatal(err)
	}
	if _, err := res.Exec(`CREATE TABLE foo (id0 INTEGER NOT NULL, id1 INTEGER NOT NULL, PRIMARY KEY(id0, id1));`); err != nil {
		t.Fatal(err)
	}
	res.insertIDStmt, err = res.Prepare(`INSERT INTO foo (id0, id1) VALUES (?, ?);`)
	if err != nil {
		t.Fatal(err)
	}
	res.queryIDStmt, err = res.Prepare(`SELECT id0, id1 FROM foo WHERE id0 = ? AND id1 = ?;`)
	if err != nil {
		t.Fatal(err)
	}

	return &res
}

type intBasedDB struct {
	*sql.DB
	insertIDStmt *sql.Stmt
	queryIDStmt  *sql.Stmt
}

func (db *intBasedDB) InsertID(u uuid.UUID) error {
	id0 := binary.LittleEndian.Uint64(u[:8])
	id1 := binary.LittleEndian.Uint64(u[8:])
	_, err := db.insertIDStmt.Exec(int64(id0), int64(id1))
	return err
}
func (db *intBasedDB) QueryID(u uuid.UUID) error {
	var id0 uint64
	var id1 uint64
	binary.LittleEndian.PutUint64(u[:8], id0)
	binary.LittleEndian.PutUint64(u[8:], id1)
	return db.queryIDStmt.QueryRow(&id0, &id1).Err()
}
func (db *intBasedDB) Close() error {
	db.DB.Close()
	db.insertIDStmt.Close()
	db.queryIDStmt.Close()
	return nil
}

func benchmarkSqliteUUID(b *testing.B, db UUIDDB) {
	defer db.Close()
	// uuids := make([]uuid.UUID, b.N)
	const uuidsLen = int(4096)
	// const uuidsLen = int(1e+6)
	b.Logf("allocating %d uuids...", uuidsLen)
	uuids := make([]uuid.UUID, uuidsLen)
	for i := range uuids {
		uuids[i] = uuid.New()
	}
	b.Logf("allocating %d uuids...done", uuidsLen)

	b.StartTimer()
	for i := range uuids {
		if err := db.InsertID(uuids[i]); err != nil {
			b.Fatalf("uuid %d: %v", i, err)
		}
	}
	b.StopTimer()
	b.Logf("inserts: %v", b.Elapsed())

	b.StartTimer()
	for i := range uuids {
		if err := db.QueryID(uuids[i]); err != nil {
			b.Fatalf("uuid %d: %v", i, err)
		}
	}
	b.StopTimer()
	b.Logf("query: %v", b.Elapsed())
}

func BenchmarkSqliteUUIDText(b *testing.B) {
	benchmarkSqliteUUID(b, NewBasicDB(b, "TEXT"))
}

func BenchmarkSqliteUUIDBlob(b *testing.B) {
	benchmarkSqliteUUID(b, NewBasicDB(b, "BLOB"))
}

func BenchmarkSqliteUUIDInt(b *testing.B) {
	benchmarkSqliteUUID(b, NewIntBasedDB(b))
}

func DebugPrintTable(db *sql.DB, table string, start, limit int) {
	q := fmt.Sprintf("SELECT * FROM %s", table)
	if start >= 0 {
		q = fmt.Sprintf("%s WHERE ROWID >= %d", q, start)
	}
	if limit >= 0 {
		q = fmt.Sprintf("%s %d", q, limit)
	}
	rows, err := db.Query(q)
	if err != nil {
		panic(err)
	}
	defer rows.Close()
	cc, err := rows.Columns()
	if err != nil {
		panic(err)
	}
	header := strings.Join(cc, ",")
	println(header)
	println(strings.Repeat("=", len(header)))
	dest := make([]any, len(cc))
	for i := range dest {
		dest[i] = &cc[i]
	}
	for rows.Next() {
		if err := rows.Scan(dest...); err != nil {
			panic(err)
		}
		println(strings.Join(cc, ","))
	}
	if err := rows.Err(); err != nil {
		panic(err)
	}
}