package bolt import ( "context" "fmt" "os" "path/filepath" "time" "github.com/coreos/bbolt" "go.uber.org/zap" "github.com/influxdata/influxdb/kit/tracing" "github.com/influxdata/influxdb/kv" ) // KVStore is a kv.Store backed by boltdb. type KVStore struct { path string db *bolt.DB logger *zap.Logger } // NewKVStore returns an instance of KVStore with the file at // the provided path. func NewKVStore(path string) *KVStore { return &KVStore{ path: path, logger: zap.NewNop(), } } // Open creates boltDB file it doesn't exists and opens it otherwise. func (s *KVStore) Open(ctx context.Context) error { span, _ := tracing.StartSpanFromContext(ctx) defer span.Finish() // Ensure the required directory structure exists. if err := os.MkdirAll(filepath.Dir(s.path), 0700); err != nil { return fmt.Errorf("unable to create directory %s: %v", s.path, err) } if _, err := os.Stat(s.path); err != nil && !os.IsNotExist(err) { return err } // Open database file. db, err := bolt.Open(s.path, 0600, &bolt.Options{Timeout: 1 * time.Second}) if err != nil { return fmt.Errorf("unable to open boltdb file %v", err) } s.db = db s.logger.Info("Resources opened", zap.String("path", s.path)) return nil } // Close the connection to the bolt database func (s *KVStore) Close() error { if s.db != nil { return s.db.Close() } return nil } // Flush removes all bolt keys within each bucket. func (s *KVStore) Flush(ctx context.Context) { _ = s.db.Update( func(tx *bolt.Tx) error { return tx.ForEach(func(name []byte, b *bolt.Bucket) error { s.cleanBucket(tx, b) return nil }) }, ) } func (s *KVStore) cleanBucket(tx *bolt.Tx, b *bolt.Bucket) { // nested bucket recursion base case: if b == nil { return } c := b.Cursor() for k, v := c.First(); k != nil; k, v = c.Next() { _ = v if err := c.Delete(); err != nil { // clean out nexted buckets s.cleanBucket(tx, b.Bucket(k)) } } } // WithLogger sets the logger on the store. func (s *KVStore) WithLogger(l *zap.Logger) { s.logger = l } // WithDB sets the boltdb on the store. func (s *KVStore) WithDB(db *bolt.DB) { s.db = db } // View opens up a view transaction against the store. func (s *KVStore) View(ctx context.Context, fn func(tx kv.Tx) error) error { span, ctx := tracing.StartSpanFromContext(ctx) defer span.Finish() return s.db.View(func(tx *bolt.Tx) error { return fn(&Tx{ tx: tx, ctx: ctx, }) }) } // Update opens up an update transaction against the store. func (s *KVStore) Update(ctx context.Context, fn func(tx kv.Tx) error) error { span, ctx := tracing.StartSpanFromContext(ctx) defer span.Finish() return s.db.Update(func(tx *bolt.Tx) error { return fn(&Tx{ tx: tx, ctx: ctx, }) }) } // Tx is a light wrapper around a boltdb transaction. It implements kv.Tx. type Tx struct { tx *bolt.Tx ctx context.Context } // Context returns the context for the transaction. func (tx *Tx) Context() context.Context { return tx.ctx } // WithContext sets the context for the transaction. func (tx *Tx) WithContext(ctx context.Context) { tx.ctx = ctx } // createBucketIfNotExists creates a bucket with the provided byte slice. func (tx *Tx) createBucketIfNotExists(b []byte) (*Bucket, error) { bkt, err := tx.tx.CreateBucketIfNotExists(b) if err != nil { return nil, err } return &Bucket{ bucket: bkt, }, nil } // Bucket retrieves the bucket named b. func (tx *Tx) Bucket(b []byte) (kv.Bucket, error) { bkt := tx.tx.Bucket(b) if bkt == nil { return tx.createBucketIfNotExists(b) } return &Bucket{ bucket: bkt, }, nil } // Bucket implements kv.Bucket. type Bucket struct { bucket *bolt.Bucket } // Get retrieves the value at the provided key. func (b *Bucket) Get(key []byte) ([]byte, error) { val := b.bucket.Get(key) if len(val) == 0 { return nil, kv.ErrKeyNotFound } return val, nil } // Put sets the value at the provided key. func (b *Bucket) Put(key []byte, value []byte) error { err := b.bucket.Put(key, value) if err == bolt.ErrTxNotWritable { return kv.ErrTxNotWritable } return err } // Delete removes the provided key. func (b *Bucket) Delete(key []byte) error { err := b.bucket.Delete(key) if err == bolt.ErrTxNotWritable { return kv.ErrTxNotWritable } return err } // Cursor retrieves a cursor for iterating through the entries // in the key value store. func (b *Bucket) Cursor() (kv.Cursor, error) { return &Cursor{ cursor: b.bucket.Cursor(), }, nil } // Cursor is a struct for iterating through the entries // in the key value store. type Cursor struct { cursor *bolt.Cursor } // Seek seeks for the first key that matches the prefix provided. func (c *Cursor) Seek(prefix []byte) ([]byte, []byte) { k, v := c.cursor.Seek(prefix) if len(k) == 0 && len(v) == 0 { return nil, nil } return k, v } // First retrieves the first key value pair in the bucket. func (c *Cursor) First() ([]byte, []byte) { k, v := c.cursor.First() if len(k) == 0 && len(v) == 0 { return nil, nil } return k, v } // Last retrieves the last key value pair in the bucket. func (c *Cursor) Last() ([]byte, []byte) { k, v := c.cursor.Last() if len(k) == 0 && len(v) == 0 { return nil, nil } return k, v } // Next retrieves the next key in the bucket. func (c *Cursor) Next() ([]byte, []byte) { k, v := c.cursor.Next() if len(k) == 0 && len(v) == 0 { return nil, nil } return k, v } // Prev retrieves the previous key in the bucket. func (c *Cursor) Prev() ([]byte, []byte) { k, v := c.cursor.Prev() if len(k) == 0 && len(v) == 0 { return nil, nil } return k, v }