1// Package prefixeddb wraps database operations with a shared key prefix.
  2package prefixeddb
  3
  4import (
  5	"bytes"
  6	"slices"
  7
  8	"github.com/arnaucube/arbo/db"
  9)
 10
 11// PrefixedDatabase wraps a db.Database prefixing all keys with `prefix`.
 12type PrefixedDatabase struct {
 13	prefix []byte
 14	db     db.Database
 15}
 16
 17// check that PrefixedDatabase implements the db.Database interface
 18var _ db.Database = (*PrefixedDatabase)(nil)
 19
 20func prefixSlice(prefix, v []byte) []byte {
 21	// Copy prefix to ensure we don't reuse the slice's backing array.
 22	jointPrefix := make([]byte, 0, len(prefix)+len(v))
 23	jointPrefix = append(jointPrefix, prefix...)
 24	jointPrefix = append(jointPrefix, v...)
 25	// Make sure the prefix's capacity is fixed, to prevent future appends from sharing memory.
 26	return slices.Clip(jointPrefix)
 27}
 28
 29// NewPrefixedDatabase creates a new PrefixedDatabase.  If the db is already a
 30// PrefixedDatabase, instead of wrapping again, the prefixes are appended to
 31// avoid unnecessary layers.
 32func NewPrefixedDatabase(db db.Database, prefix []byte) *PrefixedDatabase {
 33	if pdb, ok := db.(*PrefixedDatabase); ok {
 34		return &PrefixedDatabase{prefixSlice(pdb.prefix, prefix), pdb.db}
 35	}
 36	return &PrefixedDatabase{prefix, db}
 37}
 38
 39// Close implements the db.Database.Close interface method.  Notice that this
 40// method also closes the wrapped db.Database.
 41func (d *PrefixedDatabase) Close() error {
 42	return d.db.Close()
 43}
 44
 45// Compact implements the db.Database.Compact interface method.
 46func (d *PrefixedDatabase) Compact() error {
 47	return d.db.Compact()
 48}
 49
 50// WriteTx returns a db.WriteTx
 51func (d *PrefixedDatabase) WriteTx() db.WriteTx {
 52	return NewPrefixedWriteTx(d.db.WriteTx(), d.prefix)
 53}
 54
 55// Get implements the db.Database.Get interface method
 56func (d *PrefixedDatabase) Get(key []byte) ([]byte, error) {
 57	return d.db.Get(prefixSlice(d.prefix, key))
 58}
 59
 60// Iterate implements the db.Database.Iterate interface method
 61func (d *PrefixedDatabase) Iterate(prefix []byte, callback func(key, value []byte) bool) error {
 62	return d.db.Iterate(prefixSlice(d.prefix, prefix), func(key, value []byte) bool {
 63		return callback(bytes.TrimPrefix(key, d.prefix), value)
 64	})
 65}
 66
 67// PrefixedReader wraps a Reader tx prefixing all keys with `prefix`.
 68type PrefixedReader struct {
 69	prefix []byte
 70	rd     db.Reader
 71}
 72
 73// check that PrefixedDatabase implements the db.Database interface
 74var _ db.Reader = (*PrefixedReader)(nil)
 75
 76// NewPrefixedReader creates a new PrefixedReader.  If the rd is already a
 77// PrefixedReader, instead of wrapping again, the prefixes are appended to
 78// avoid unnecessary layers.
 79func NewPrefixedReader(rd db.Reader, prefix []byte) *PrefixedReader {
 80	if pdb, ok := rd.(*PrefixedReader); ok {
 81		return &PrefixedReader{prefixSlice(pdb.prefix, prefix), pdb.rd}
 82	}
 83	return &PrefixedReader{prefix, rd}
 84}
 85
 86// Get retrieves a value using the configured key prefix.
 87func (d *PrefixedReader) Get(key []byte) ([]byte, error) {
 88	return d.rd.Get(prefixSlice(d.prefix, key))
 89}
 90
 91// Iterate implements the db.Database.Iterate interface method
 92func (d *PrefixedReader) Iterate(prefix []byte, callback func(key, value []byte) bool) error {
 93	return d.rd.Iterate(prefixSlice(d.prefix, prefix), func(key, value []byte) bool {
 94		return callback(bytes.TrimPrefix(key, d.prefix), value)
 95	})
 96}
 97
 98// PrefixedWriteTx wraps a WriteTx prefixing all keys with `prefix`.
 99type PrefixedWriteTx struct {
100	prefix []byte
101	tx     db.WriteTx
102}
103
104// check that PrefixedWriteTx implements the db.WriteTx interface
105var _ db.WriteTx = (*PrefixedWriteTx)(nil)
106
107// NewPrefixedWriteTx creates a new db.WriteTx.  If the tx is already a
108// PrefixedWriteTx, instead of wrapping again, the prefixes are appended to avoid
109// unnecessary layers.
110func NewPrefixedWriteTx(tx db.WriteTx, prefix []byte) *PrefixedWriteTx {
111	if ptx, ok := tx.(*PrefixedWriteTx); ok {
112		return &PrefixedWriteTx{prefixSlice(ptx.prefix, prefix), ptx.tx}
113	}
114	return &PrefixedWriteTx{prefix, tx}
115}
116
117// Get implements the db.WriteTx.Get interface method
118func (t *PrefixedWriteTx) Get(key []byte) ([]byte, error) {
119	return t.tx.Get(prefixSlice(t.prefix, key))
120}
121
122// Iterate visits matching entries with the configured prefix removed.
123func (t *PrefixedWriteTx) Iterate(prefix []byte, callback func(key, value []byte) bool) error {
124	return t.tx.Iterate(prefixSlice(t.prefix, prefix), func(key, value []byte) bool {
125		return callback(bytes.TrimPrefix(key, t.prefix), value)
126	})
127}
128
129// Discard implements the db.WriteTx.Discard interface method.  Notice that this
130// method also discards the wrapped db.WriteTx.
131func (t *PrefixedWriteTx) Discard() {
132	t.tx.Discard()
133}
134
135// Set implements the db.WriteTx.Set interface method
136func (t *PrefixedWriteTx) Set(key, value []byte) error {
137	return t.tx.Set(prefixSlice(t.prefix, key), value)
138}
139
140// Delete implements the db.WriteTx.Delete interface method
141func (t *PrefixedWriteTx) Delete(key []byte) error {
142	return t.tx.Delete(prefixSlice(t.prefix, key))
143}
144
145// Apply implements the db.WriteTx.Apply interface method
146func (t *PrefixedWriteTx) Apply(other db.WriteTx) error {
147	return t.tx.Apply(other)
148}
149
150// Unwrap returns the wrapped WriteTx
151func (t *PrefixedWriteTx) Unwrap() db.WriteTx {
152	return t.tx
153}
154
155// Commit implements the db.WriteTx.Commit interface method.  Notice that this
156// method also commits the wrapped db.WriteTx.
157func (t *PrefixedWriteTx) Commit() error {
158	return t.tx.Commit()
159}