// 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) }