tdb/db.go

263 lines
6.1 KiB
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)
}