badger_range.go

package main

import (
	"bytes"
	"fmt"
	"log"

	badger "github.com/dgraph-io/badger/v4"
)

func main() {
	// Open the Badger database located in the /tmp/badger directory.
	// It will be created if it doesn't exist.
	db, err := badger.Open(badger.DefaultOptions("./tmp/badger"))
	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()

	ss1 := []byte("table-one")
	ss2 := []byte("table-two")

	if err = Insert(db, ss1); err != nil {
		panic(err)
	}
	if err = Insert(db, ss2); err != nil {
		panic(err)
	}

	fmt.Println("Reading full forwards")
	if err = ReadFull(db, ss1, false); err != nil {
		panic(err)
	}

	fmt.Println("Reading full backwards")
	if err = ReadFull(db, ss1, true); err != nil {
		panic(err)
	}

	fmt.Println("Reading range forwards")
	if err = ReadRange(db, ss1, []byte("row-3"), []byte("row-6"), false); err != nil {
		panic(err)
	}

	fmt.Println("Reading range backwards")
	if err = ReadRange(db, ss1, []byte("row-3"), []byte("row-6"), true); err != nil {
		panic(err)
	}

}

// Reads all the keys in a subspace
// For `reverse=true` we need to seek to find a key that is one larger than the subspace range
// then scan backwards from there
func ReadFull(db *badger.DB, subspace []byte, reverse bool) error {
	return db.View(func(tx *badger.Txn) error {

		iter := tx.NewIterator(badger.IteratorOptions{
			Reverse: reverse,
			Prefix:  subspace,
		})

		startK := subspace
		if reverse {
			startK = End(subspace)
		}

		defer iter.Close()
		for iter.Seek(startK); iter.ValidForPrefix(subspace); iter.Next() {
			val, err := iter.Item().ValueCopy(nil)
			if err != nil {
				return err
			}
			fmt.Printf("%s : %s \n", string(iter.Item().Key()), val)
		}

		return nil
	})
}

// Read between two keys, when going forward if the key is greater than the end key then stop
// and for backwards check if the key is less than the start key and then stop
func ReadRange(db *badger.DB, subspace []byte, start []byte, end []byte, reverse bool) error {
	return db.View(func(tx *badger.Txn) error {

		iter := tx.NewIterator(badger.IteratorOptions{
			Reverse: reverse,
			Prefix:  subspace,
		})

		endK := Pack(subspace, end)
		startK := Pack(subspace, start)
		if reverse {
			startK = Pack(subspace, end)
			endK = Pack(subspace, start)
		}

		defer iter.Close()
		for iter.Seek(startK); iter.ValidForPrefix(subspace); iter.Next() {
			key := iter.Item().Key()
			if reverse {
				if bytes.Compare(key, endK) < 0 {
					break
				}
			} else {
				if bytes.Compare(key, endK) > 0 {
					break
				}
			}

			val, err := iter.Item().ValueCopy(nil)
			if err != nil {
				return err
			}
			fmt.Printf("%s : %s \n", string(iter.Item().Key()), val)
		}

		return nil
	})
}

func Insert(db *badger.DB, subspace []byte) error {
	tx := db.NewTransaction(true)
	defer tx.Discard()

	for i := 0; i < 10; i++ {
		key := Pack(subspace, []byte(fmt.Sprintf("row-%d", i)))
		if err := tx.Set(key, []byte(fmt.Sprintf("value-%d", i))); err != nil {
			return err
		}
	}

	return tx.Commit()
}

// Copy the prefix into a new slice that is one larger than
// the prefix and add an `0xFF` byte to it so
func End(prefix []byte) []byte {
	end := make([]byte, len(prefix)+1)
	copy(end, prefix)

	end[len(end)-1] = 0xFF
	return end
}

// Pack is a simple function to join a key to a subspace
func Pack(subspace []byte, key []byte) []byte {
	keyBits := [][]byte{subspace, key}
	return bytes.Join(keyBits, []byte("-"))
}