tdb/db.go

// Package tdb provides a terrible database built on top of bolt.
// It does all sorts of too-smart things with reflection that will
// either be great and make your life easier, or suck and you just
// shouldn't use this package.
package tdb

import (
	"fmt"
	"log"
	"reflect"

	"github.com/golang/protobuf/proto"
	bolt "go.etcd.io/bbolt"
)

const logPrefix = "[tdb] "

type debugLogger interface {
	debugLog(message string)
	debugLogf(f string, args ...interface{})
}

type DB interface {
	debugLogger
	Transactable
	Close() error
	GetTable(name string) (Table, error)
	GetTableOrPanic(name string) Table
}

type DBSetup interface {
	debugLogger
	AddTable(thing proto.Message, createSchema CreateTableSchema) error
	AddTableOrPanic(thing proto.Message, createSchema CreateTableSchema)
	AddIndex(options SimpleIndexOptions) error
	AddIndexOrPanic(options SimpleIndexOptions)
	SetDebug(enabled bool)
}

type CreateDBSchema func(DBSetup) error

type db struct {
	ready  bool
	closed bool
	debug  bool
	b      *bolt.DB
	tables map[string]*table
}

func New(b *bolt.DB, tableBucket interface{}, createSchema CreateDBSchema) (DB, error) {
	tdb := &db{
		b:      b,
		tables: make(map[string]*table),
	}
	err := createSchema(tdb)
	if err != nil {
		return nil, err
	}
	tdb.debugLog("Schema creation completed successfuly, initializing...")
	err = tdb.b.Update(func(tx *bolt.Tx) error {
		return tdb.initialize(convertTx(tx))
	})
	if err != nil {
		return nil, err
	}

	tdb.debugLog("Initialization complete, populating table bucket...")
	tdb.populateTableBucket(tableBucket)

	tdb.debugLog("Setup of new tdb complete... returning")
	return tdb, err
}

func NewOrPanic(b *bolt.DB, tableBucket interface{}, createSchema CreateDBSchema) DB {
	tdb, err := New(b, tableBucket, createSchema)
	if err != nil {
		panic(err)
	}
	return tdb
}

func (db *db) Close() error {
	db.closed = true
	return db.b.Close()
}

func (db *db) populateTableBucket(tableBucket interface{}) {
	if tableBucket == nil {
		db.debugLog("[populate] no table bucket")
		return
	}

	bucketPtrVal := reflect.ValueOf(tableBucket)

	if bucketPtrVal.Kind() != reflect.Ptr {
		db.debugLog("[populate] tableBucket is not a pointer")
		return
	}

	bucketVal := bucketPtrVal.Elem()

	if bucketVal.Kind() != reflect.Struct {
		db.debugLog("[populate] tableBucket is not a ptr to a struct")
		return
	}

	tableBucketType := bucketVal.Type()
	fieldCount := tableBucketType.NumField()
	for i := 0; i < fieldCount; i++ {
		db.populateField(tableBucketType.Field(i), bucketVal)
	}
}

func (db *db) populateField(field reflect.StructField, bucketVal reflect.Value) {
	table, ok := db.tables[field.Name]
	if !ok {
		db.debugLogf("[populate] no such table '%s'", field.Name)
		return
	}

	tableType := reflect.TypeOf((*Table)(nil)).Elem()
	if field.Type != tableType {
		db.debugLogf("[populate] wrong types for '%s', got '%s', expected '%s'", field.Name, field.Type.String(), tableType.String())
		return
	}
	// maybe check CanSet()?
	bucketValField := bucketVal.FieldByName(field.Name)
	if !bucketValField.CanSet() {
		db.debugLogf("[populate] cannot set field '%s'", field.Name)
		return
	}

	db.debugLogf("[populate] set field '%s'", field.Name)
	bucketValField.Set(reflect.ValueOf(table))
}

func (db *db) SetDebug(debug bool) {
	db.debug = debug
}

func (db *db) AddTable(thing proto.Message, createSchema CreateTableSchema) error {
	t := dbTypeOf(thing)
	db.debugLogf("AddTable invoked for type %s", t.Name)

	if _, has := db.tables[t.Name]; has {
		return fmt.Errorf("Database already has table with name '%s'", t.Name)
	}

	idField := t.IdField()

	table, err := newTable(db, t, idField, createSchema)
	if err != nil {
		return err
	}
	db.debugLogf("Table schema creation for '%s' completed successfuly", t.Name)
	db.tables[t.Name] = table
	return nil
}

func (db *db) AddTableOrPanic(thing proto.Message, createSchema CreateTableSchema) {
	if err := db.AddTable(thing, createSchema); err != nil {
		panic(err)
	}
}

func (db *db) AddIndex(options SimpleIndexOptions) error {
	table, ok := db.tables[options.Table]
	if !ok {
		return fmt.Errorf("No such table '%s'", options.Table)
	}
	return table.AddIndex(options)
}

func (db *db) AddIndexOrPanic(options SimpleIndexOptions) {
	if err := db.AddIndex(options); err != nil {
		panic(err)
	}
}

func (db *db) debugLog(message string) {
	if db.debug {
		log.Print(logPrefix + message)
	}
}

func (db *db) debugLogf(f string, args ...interface{}) {
	if db.debug {
		log.Printf(logPrefix + fmt.Sprintf(f, args...))
	}
}

func (db *db) initialize(tx *Tx) error {
	err := db.initializeTables(tx)
	if err != nil {
		return err
	}
	db.debugLog("Initialization complete")
	return nil
}

func (db *db) initializeTables(tx *Tx) error {
	for name, table := range db.tables {
		db.debugLogf("Initializating table '%s'...", name)
		err := table.initialize(tx)
		if err != nil {
			return err
		}
		db.debugLogf("Initialized table '%s'", name)
	}
	return nil
}

func (db *db) GetTable(name string) (Table, error) {
	table, ok := db.tables[name]
	if !ok {
		return nil, fmt.Errorf("No such table '%s'", name)
	}
	return table, nil
}

func (db *db) GetTableOrPanic(name string) Table {
	table, err := db.GetTable(name)
	if err != nil {
		panic(err)
	}
	return table
}

func (db *db) ReadTx(t Transaction) error {
	return db.b.View(func(tx *bolt.Tx) error {
		return t(convertTx(tx))
	})
}

func (db *db) readTxHelper(t Transaction, txs ...*Tx) error {
	txlen := len(txs)
	if txlen > 1 {
		db.debugLogf("[db.readTxHelper] Got %d transactions, can only handle 1.", txlen)
		return fmt.Errorf("Got %d transactions, can only handle 1.", txlen)
	} else if txlen == 1 {
		db.debugLogf("[db.readTxHelper] Found existing transaction: %#v", txs)
		return t(txs[0])
	}
	return db.ReadTx(t)
}

func (db *db) WriteTx(t Transaction) error {
	return db.b.Update(func(tx *bolt.Tx) error {
		return t(convertTx(tx))
	})
}

func (db *db) writeTxHelper(t Transaction, txs ...*Tx) error {
	txlen := len(txs)
	if txlen > 1 {
		db.debugLogf("[db.readTxHelper] Got %d transactions, can only handle 1.", txlen)
		return fmt.Errorf("Got %d transactions, can only handle 1.", txlen)
	} else if txlen == 1 {
		return t(txs[0])
	}
	return db.WriteTx(t)
}