package pebbleutil import ( "context" "iter" "slices" "github.com/cockroachdb/pebble/v2" "github.com/cockroachdb/pebble/v2/vfs" ) // Iterate iterates over a pebble reader. func Iterate[T any](src pebble.Reader, iterOptions *pebble.IterOptions, f func(it *pebble.Iterator) (T, error)) iter.Seq2[T, error] { var zero T return func(yield func(T, error) bool) { it, err := src.NewIter(iterOptions) if err != nil { yield(zero, err) return } for it.First(); it.Valid(); it.Next() { value, err := f(it) if err != nil { _ = it.Close() yield(zero, err) return } if !yield(value, nil) { _ = it.Close() return } } err = it.Error() if err != nil { _ = it.Close() yield(zero, err) return } err = it.Close() if err != nil { yield(zero, err) return } } } // IterateKeys yields the keys in a pebble reader. func IterateKeys(src pebble.Reader, iterOptions *pebble.IterOptions) iter.Seq2[[]byte, error] { return Iterate(src, iterOptions, func(it *pebble.Iterator) ([]byte, error) { return slices.Clone(it.Key()), nil }) } // IterateValues yields the values in a pebble reader. func IterateValues(src pebble.Reader, iterOptions *pebble.IterOptions) iter.Seq2[[]byte, error] { return Iterate(src, iterOptions, func(it *pebble.Iterator) ([]byte, error) { value, err := it.ValueAndErr() if err != nil { return nil, err } return slices.Clone(value), nil }) } // MustOpen opens a pebble database. It sets options useful for pomerium and panics if there is an error. func MustOpen(dirname string, options *pebble.Options) *pebble.DB { db, err := Open(dirname, options) if err != nil { panic(err) } return db } // MustOpenMemory opens an in-memory pebble database. It panics if there is an error. func MustOpenMemory(options *pebble.Options) *pebble.DB { if options == nil { options = new(pebble.Options) } options.FS = vfs.NewMem() return MustOpen("", options) } // Open opens a pebble database. It sets options useful for pomerium. func Open(dirname string, options *pebble.Options) (*pebble.DB, error) { if options == nil { options = new(pebble.Options) } options.LoggerAndTracer = pebbleLogger{} return pebble.Open(dirname, options) } // PrefixToUpperBound returns an upper bound for the given prefix. func PrefixToUpperBound(prefix []byte) []byte { upperBound := make([]byte, len(prefix)) copy(upperBound, prefix) for i := len(upperBound) - 1; i >= 0; i-- { upperBound[i] = upperBound[i] + 1 if upperBound[i] != 0 { return upperBound[:i+1] } } return nil // no upper-bound } type pebbleLogger struct{} func (pebbleLogger) Infof(_ string, _ ...any) {} func (pebbleLogger) Errorf(_ string, _ ...any) {} func (pebbleLogger) Fatalf(_ string, _ ...any) {} func (pebbleLogger) Eventf(_ context.Context, _ string, _ ...any) {} func (pebbleLogger) IsTracingEnabled(_ context.Context) bool { return false }