Replace BadgerDB with bbolt

This commit is contained in:
David Stotijn
2025-01-13 23:15:18 +01:00
parent 24c2ecfa4b
commit fcf0e1c51e
23 changed files with 1153 additions and 1244 deletions

View File

@ -1,57 +0,0 @@
package badger
import (
"fmt"
"github.com/dgraph-io/badger/v3"
)
const (
// Key prefixes. Each prefix value should be unique.
projectPrefix = 0x00
reqLogPrefix = 0x01
resLogPrefix = 0x02
senderReqPrefix = 0x03
// Request log indices.
reqLogProjectIDIndex = 0x00
// Sender request indices.
senderReqProjectIDIndex = 0x00
)
// Database is used to store and retrieve data from an underlying Badger database.
type Database struct {
badger *badger.DB
}
// OpenDatabase opens a new Badger database.
func OpenDatabase(opts badger.Options) (*Database, error) {
db, err := badger.Open(opts)
if err != nil {
return nil, fmt.Errorf("badger: failed to open database: %w", err)
}
return &Database{badger: db}, nil
}
// Close closes the underlying Badger database.
func (db *Database) Close() error {
return db.badger.Close()
}
// DatabaseFromBadgerDB returns a Database with `db` set as the underlying
// Badger database.
func DatabaseFromBadgerDB(db *badger.DB) *Database {
return &Database{badger: db}
}
func entryKey(prefix, index byte, value []byte) []byte {
// Key consists of: | prefix (byte) | index (byte) | value
key := make([]byte, 2+len(value))
key[0] = prefix
key[1] = index
copy(key[2:len(value)+2], value)
return key
}

View File

@ -1,21 +0,0 @@
package badger
import (
"github.com/dgraph-io/badger/v3"
"go.uber.org/zap"
)
// Interface guard.
var _ badger.Logger = (*Logger)(nil)
type Logger struct {
*zap.SugaredLogger
}
func NewLogger(l *zap.SugaredLogger) *Logger {
return &Logger{l}
}
func (l *Logger) Warningf(template string, args ...interface{}) {
l.Warnf(template, args)
}

View File

@ -1,115 +0,0 @@
package badger
import (
"bytes"
"context"
"encoding/gob"
"errors"
"fmt"
"github.com/dgraph-io/badger/v3"
"github.com/oklog/ulid"
"github.com/dstotijn/hetty/pkg/proj"
)
func (db *Database) UpsertProject(ctx context.Context, project proj.Project) error {
buf := bytes.Buffer{}
err := gob.NewEncoder(&buf).Encode(project)
if err != nil {
return fmt.Errorf("badger: failed to encode project: %w", err)
}
err = db.badger.Update(func(txn *badger.Txn) error {
return txn.Set(entryKey(projectPrefix, 0, project.ID[:]), buf.Bytes())
})
if err != nil {
return fmt.Errorf("badger: failed to commit transaction: %w", err)
}
return nil
}
func (db *Database) FindProjectByID(ctx context.Context, projectID ulid.ULID) (project proj.Project, err error) {
err = db.badger.View(func(txn *badger.Txn) error {
item, err := txn.Get(entryKey(projectPrefix, 0, projectID[:]))
if err != nil {
return err
}
err = item.Value(func(rawProject []byte) error {
return gob.NewDecoder(bytes.NewReader(rawProject)).Decode(&project)
})
if err != nil {
return fmt.Errorf("failed to retrieve or parse project: %w", err)
}
return nil
})
if errors.Is(err, badger.ErrKeyNotFound) {
return proj.Project{}, proj.ErrProjectNotFound
}
if err != nil {
return proj.Project{}, fmt.Errorf("badger: failed to commit transaction: %w", err)
}
return project, nil
}
func (db *Database) DeleteProject(ctx context.Context, projectID ulid.ULID) error {
err := db.ClearRequestLogs(ctx, projectID)
if err != nil {
return fmt.Errorf("badger: failed to delete project request logs: %w", err)
}
err = db.DeleteSenderRequests(ctx, projectID)
if err != nil {
return fmt.Errorf("badger: failed to delete project sender requests: %w", err)
}
err = db.badger.Update(func(txn *badger.Txn) error {
return txn.Delete(entryKey(projectPrefix, 0, projectID[:]))
})
if err != nil {
return fmt.Errorf("badger: failed to delete project item: %w", err)
}
return nil
}
func (db *Database) Projects(ctx context.Context) ([]proj.Project, error) {
projects := make([]proj.Project, 0)
err := db.badger.View(func(txn *badger.Txn) error {
var rawProject []byte
prefix := entryKey(projectPrefix, 0, nil)
iterator := txn.NewIterator(badger.DefaultIteratorOptions)
defer iterator.Close()
for iterator.Seek(prefix); iterator.ValidForPrefix(prefix); iterator.Next() {
rawProject, err := iterator.Item().ValueCopy(rawProject)
if err != nil {
return fmt.Errorf("failed to copy value: %w", err)
}
var project proj.Project
err = gob.NewDecoder(bytes.NewReader(rawProject)).Decode(&project)
if err != nil {
return fmt.Errorf("failed to decode project: %w", err)
}
projects = append(projects, project)
}
return nil
})
if err != nil {
return nil, fmt.Errorf("badger: failed to commit transaction: %w", err)
}
return projects, nil
}

View File

@ -1,328 +0,0 @@
package badger
import (
"bytes"
"context"
"encoding/gob"
"errors"
"math/rand"
"regexp"
"testing"
"time"
badgerdb "github.com/dgraph-io/badger/v3"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/oklog/ulid"
"github.com/dstotijn/hetty/pkg/filter"
"github.com/dstotijn/hetty/pkg/proj"
"github.com/dstotijn/hetty/pkg/scope"
)
//nolint:gosec
var ulidEntropy = rand.New(rand.NewSource(time.Now().UnixNano()))
var regexpCompareOpt = cmp.Comparer(func(x, y *regexp.Regexp) bool {
switch {
case x == nil && y == nil:
return true
case x == nil || y == nil:
return false
default:
return x.String() == y.String()
}
})
func TestUpsertProject(t *testing.T) {
t.Parallel()
badgerDB, err := badgerdb.Open(badgerdb.DefaultOptions("").WithInMemory(true))
if err != nil {
t.Fatalf("failed to open badger database: %v", err)
}
database := DatabaseFromBadgerDB(badgerDB)
defer database.Close()
searchExpr, err := filter.ParseQuery("foo AND bar OR NOT baz")
if err != nil {
t.Fatalf("unexpected error (expected: nil, got: %v)", err)
}
exp := proj.Project{
ID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
Name: "foobar",
Settings: proj.Settings{
ReqLogBypassOutOfScope: true,
ReqLogOnlyFindInScope: true,
ReqLogSearchExpr: searchExpr,
ScopeRules: []scope.Rule{
{
URL: regexp.MustCompile("^https://(.*)example.com(.*)$"),
Header: scope.Header{
Key: regexp.MustCompile("^X-Foo(.*)$"),
Value: regexp.MustCompile("^foo(.*)$"),
},
Body: regexp.MustCompile("^foo(.*)"),
},
},
},
}
err = database.UpsertProject(context.Background(), exp)
if err != nil {
t.Fatalf("unexpected error storing project: %v", err)
}
var rawProject []byte
err = badgerDB.View(func(txn *badgerdb.Txn) error {
item, err := txn.Get(entryKey(projectPrefix, 0, exp.ID[:]))
if err != nil {
return err
}
rawProject, err = item.ValueCopy(nil)
return err
})
if err != nil {
t.Fatalf("unexpected error retrieving project from database: %v", err)
}
got := proj.Project{}
err = gob.NewDecoder(bytes.NewReader(rawProject)).Decode(&got)
if err != nil {
t.Fatalf("unexpected error decoding project: %v", err)
}
if diff := cmp.Diff(exp, got, regexpCompareOpt, cmpopts.IgnoreUnexported(proj.Project{})); diff != "" {
t.Fatalf("project not equal (-exp, +got):\n%v", diff)
}
}
func TestFindProjectByID(t *testing.T) {
t.Parallel()
t.Run("existing project", func(t *testing.T) {
t.Parallel()
badgerDB, err := badgerdb.Open(badgerdb.DefaultOptions("").WithInMemory(true))
if err != nil {
t.Fatalf("failed to open badger database: %v", err)
}
database := DatabaseFromBadgerDB(badgerDB)
defer database.Close()
exp := proj.Project{
ID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
Name: "foobar",
Settings: proj.Settings{},
}
buf := bytes.Buffer{}
err = gob.NewEncoder(&buf).Encode(exp)
if err != nil {
t.Fatalf("unexpected error encoding project: %v", err)
}
err = badgerDB.Update(func(txn *badgerdb.Txn) error {
return txn.Set(entryKey(projectPrefix, 0, exp.ID[:]), buf.Bytes())
})
if err != nil {
t.Fatalf("unexpected error setting project: %v", err)
}
got, err := database.FindProjectByID(context.Background(), exp.ID)
if err != nil {
t.Fatalf("unexpected error finding project: %v", err)
}
if diff := cmp.Diff(exp, got, cmpopts.IgnoreUnexported(proj.Project{})); diff != "" {
t.Fatalf("project not equal (-exp, +got):\n%v", diff)
}
})
t.Run("project not found", func(t *testing.T) {
t.Parallel()
database, err := OpenDatabase(badgerdb.DefaultOptions("").WithInMemory(true))
if err != nil {
t.Fatalf("failed to open badger database: %v", err)
}
defer database.Close()
projectID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
_, err = database.FindProjectByID(context.Background(), projectID)
if !errors.Is(err, proj.ErrProjectNotFound) {
t.Fatalf("expected `proj.ErrProjectNotFound`, got: %v", err)
}
})
}
func TestDeleteProject(t *testing.T) {
t.Parallel()
badgerDB, err := badgerdb.Open(badgerdb.DefaultOptions("").WithInMemory(true))
if err != nil {
t.Fatalf("failed to open badger database: %v", err)
}
database := DatabaseFromBadgerDB(badgerDB)
defer database.Close()
// Store fixtures.
projectID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
reqLogID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
senderReqID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
err = badgerDB.Update(func(txn *badgerdb.Txn) error {
// Project item.
if err := txn.Set(entryKey(projectPrefix, 0, projectID[:]), nil); err != nil {
return err
}
// Sender request items.
if err := txn.Set(entryKey(senderReqPrefix, 0, senderReqID[:]), nil); err != nil {
return err
}
if err := txn.Set(entryKey(resLogPrefix, 0, senderReqID[:]), nil); err != nil {
return err
}
err := txn.Set(entryKey(senderReqPrefix, senderReqProjectIDIndex, append(projectID[:], senderReqID[:]...)), nil)
if err != nil {
return err
}
// Request log items.
if err := txn.Set(entryKey(reqLogPrefix, 0, reqLogID[:]), nil); err != nil {
return err
}
if err := txn.Set(entryKey(resLogPrefix, 0, reqLogID[:]), nil); err != nil {
return err
}
err = txn.Set(entryKey(reqLogPrefix, reqLogProjectIDIndex, append(projectID[:], reqLogID[:]...)), nil)
if err != nil {
return err
}
return nil
})
if err != nil {
t.Fatalf("unexpected error creating fixtures: %v", err)
}
err = database.DeleteProject(context.Background(), projectID)
if err != nil {
t.Fatalf("unexpected error deleting project: %v", err)
}
// Assert project key was deleted.
err = badgerDB.View(func(txn *badgerdb.Txn) error {
_, err := txn.Get(entryKey(projectPrefix, 0, projectID[:]))
return err
})
if !errors.Is(err, badgerdb.ErrKeyNotFound) {
t.Fatalf("expected `badger.ErrKeyNotFound`, got: %v", err)
}
// Assert request log item was deleted.
err = badgerDB.View(func(txn *badgerdb.Txn) error {
_, err := txn.Get(entryKey(reqLogPrefix, 0, reqLogID[:]))
return err
})
if !errors.Is(err, badgerdb.ErrKeyNotFound) {
t.Fatalf("expected `badger.ErrKeyNotFound`, got: %v", err)
}
// Assert response log item related to request log was deleted.
err = badgerDB.View(func(txn *badgerdb.Txn) error {
_, err := txn.Get(entryKey(resLogPrefix, 0, reqLogID[:]))
return err
})
if !errors.Is(err, badgerdb.ErrKeyNotFound) {
t.Fatalf("expected `badger.ErrKeyNotFound`, got: %v", err)
}
// Assert request log project ID index key was deleted.
err = badgerDB.View(func(txn *badgerdb.Txn) error {
_, err := txn.Get(entryKey(reqLogPrefix, reqLogProjectIDIndex, append(projectID[:], reqLogID[:]...)))
return err
})
if !errors.Is(err, badgerdb.ErrKeyNotFound) {
t.Fatalf("expected `badger.ErrKeyNotFound`, got: %v", err)
}
// Assert sender request item was deleted.
err = badgerDB.View(func(txn *badgerdb.Txn) error {
_, err := txn.Get(entryKey(senderReqPrefix, 0, senderReqID[:]))
return err
})
if !errors.Is(err, badgerdb.ErrKeyNotFound) {
t.Fatalf("expected `badger.ErrKeyNotFound`, got: %v", err)
}
// Assert response log item related to sender request was deleted.
err = badgerDB.View(func(txn *badgerdb.Txn) error {
_, err := txn.Get(entryKey(resLogPrefix, 0, senderReqID[:]))
return err
})
if !errors.Is(err, badgerdb.ErrKeyNotFound) {
t.Fatalf("expected `badger.ErrKeyNotFound`, got: %v", err)
}
// Assert sender request project ID index key was deleted.
err = badgerDB.View(func(txn *badgerdb.Txn) error {
_, err := txn.Get(entryKey(senderReqPrefix, senderReqProjectIDIndex, append(projectID[:], senderReqID[:]...)))
return err
})
if !errors.Is(err, badgerdb.ErrKeyNotFound) {
t.Fatalf("expected `badger.ErrKeyNotFound`, got: %v", err)
}
}
func TestProjects(t *testing.T) {
t.Parallel()
database, err := OpenDatabase(badgerdb.DefaultOptions("").WithInMemory(true))
if err != nil {
t.Fatalf("failed to open badger database: %v", err)
}
defer database.Close()
exp := []proj.Project{
{
ID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
Name: "one",
},
{
ID: ulid.MustNew(ulid.Timestamp(time.Now())+100, ulidEntropy),
Name: "two",
},
}
// Store fixtures.
for _, project := range exp {
err = database.UpsertProject(context.Background(), project)
if err != nil {
t.Fatalf("unexpected error creating project fixture: %v", err)
}
}
got, err := database.Projects(context.Background())
if err != nil {
t.Fatalf("unexpected error finding projects: %v", err)
}
if len(exp) != len(got) {
t.Fatalf("expected %v projects, got: %v", len(exp), len(got))
}
if diff := cmp.Diff(exp, got, cmpopts.IgnoreUnexported(proj.Project{})); diff != "" {
t.Fatalf("projects not equal (-exp, +got):\n%v", diff)
}
}

View File

@ -1,260 +0,0 @@
package badger
import (
"bytes"
"context"
"encoding/gob"
"errors"
"fmt"
"github.com/dgraph-io/badger/v3"
"github.com/oklog/ulid"
"github.com/dstotijn/hetty/pkg/reqlog"
"github.com/dstotijn/hetty/pkg/scope"
)
func (db *Database) FindRequestLogs(
ctx context.Context,
filter reqlog.FindRequestsFilter,
scope *scope.Scope) ([]reqlog.RequestLog, error,
) {
if filter.ProjectID.Compare(ulid.ULID{}) == 0 {
return nil, reqlog.ErrProjectIDMustBeSet
}
txn := db.badger.NewTransaction(false)
defer txn.Discard()
reqLogIDs, err := findRequestLogIDsByProjectID(txn, filter.ProjectID)
if err != nil {
return nil, fmt.Errorf("badger: failed to find request log IDs: %w", err)
}
reqLogs := make([]reqlog.RequestLog, 0, len(reqLogIDs))
for _, reqLogID := range reqLogIDs {
reqLog, err := getRequestLogWithResponse(txn, reqLogID)
if err != nil {
return nil, fmt.Errorf("badger: failed to get request log (id: %v): %w", reqLogID.String(), err)
}
if filter.OnlyInScope {
if !reqLog.MatchScope(scope) {
continue
}
}
// Filter by search expression.
// TODO: Once pagination is introduced, this filter logic should be done
// as items are retrieved (e.g. when using a `badger.Iterator`).
if filter.SearchExpr != nil {
match, err := reqLog.Matches(filter.SearchExpr)
if err != nil {
return nil, fmt.Errorf(
"badger: failed to match search expression for request log (id: %v): %w",
reqLogID.String(), err,
)
}
if !match {
continue
}
}
reqLogs = append(reqLogs, reqLog)
}
return reqLogs, nil
}
func getRequestLogWithResponse(txn *badger.Txn, reqLogID ulid.ULID) (reqlog.RequestLog, error) {
item, err := txn.Get(entryKey(reqLogPrefix, 0, reqLogID[:]))
switch {
case errors.Is(err, badger.ErrKeyNotFound):
return reqlog.RequestLog{}, reqlog.ErrRequestNotFound
case err != nil:
return reqlog.RequestLog{}, fmt.Errorf("failed to lookup request log item: %w", err)
}
reqLog := reqlog.RequestLog{
ID: reqLogID,
}
err = item.Value(func(rawReqLog []byte) error {
err = gob.NewDecoder(bytes.NewReader(rawReqLog)).Decode(&reqLog)
if err != nil {
return fmt.Errorf("failed to decode request log: %w", err)
}
return nil
})
if err != nil {
return reqlog.RequestLog{}, fmt.Errorf("failed to retrieve or parse request log value: %w", err)
}
item, err = txn.Get(entryKey(resLogPrefix, 0, reqLogID[:]))
if errors.Is(err, badger.ErrKeyNotFound) {
return reqLog, nil
}
if err != nil {
return reqlog.RequestLog{}, fmt.Errorf("failed to get response log: %w", err)
}
err = item.Value(func(rawReslog []byte) error {
var resLog reqlog.ResponseLog
err = gob.NewDecoder(bytes.NewReader(rawReslog)).Decode(&resLog)
if err != nil {
return fmt.Errorf("failed to decode response log: %w", err)
}
reqLog.Response = &resLog
return nil
})
if err != nil {
return reqlog.RequestLog{}, fmt.Errorf("failed to retrieve or parse response log value: %w", err)
}
return reqLog, nil
}
func (db *Database) FindRequestLogByID(ctx context.Context, reqLogID ulid.ULID) (reqLog reqlog.RequestLog, err error) {
txn := db.badger.NewTransaction(false)
defer txn.Discard()
reqLog, err = getRequestLogWithResponse(txn, reqLogID)
if err != nil {
return reqlog.RequestLog{}, fmt.Errorf("badger: failed to get request log: %w", err)
}
return reqLog, nil
}
func (db *Database) StoreRequestLog(ctx context.Context, reqLog reqlog.RequestLog) error {
buf := bytes.Buffer{}
err := gob.NewEncoder(&buf).Encode(reqLog)
if err != nil {
return fmt.Errorf("badger: failed to encode request log: %w", err)
}
entries := []*badger.Entry{
// Request log itself.
{
Key: entryKey(reqLogPrefix, 0, reqLog.ID[:]),
Value: buf.Bytes(),
},
// Index by project ID.
{
Key: entryKey(reqLogPrefix, reqLogProjectIDIndex, append(reqLog.ProjectID[:], reqLog.ID[:]...)),
},
}
err = db.badger.Update(func(txn *badger.Txn) error {
for i := range entries {
err := txn.SetEntry(entries[i])
if err != nil {
return err
}
}
return nil
})
if err != nil {
return fmt.Errorf("badger: failed to commit transaction: %w", err)
}
return nil
}
func (db *Database) StoreResponseLog(ctx context.Context, reqLogID ulid.ULID, resLog reqlog.ResponseLog) error {
buf := bytes.Buffer{}
err := gob.NewEncoder(&buf).Encode(resLog)
if err != nil {
return fmt.Errorf("badger: failed to encode response log: %w", err)
}
err = db.badger.Update(func(txn *badger.Txn) error {
return txn.SetEntry(&badger.Entry{
Key: entryKey(resLogPrefix, 0, reqLogID[:]),
Value: buf.Bytes(),
})
})
if err != nil {
return fmt.Errorf("badger: failed to commit transaction: %w", err)
}
return nil
}
func (db *Database) ClearRequestLogs(ctx context.Context, projectID ulid.ULID) error {
// Note: this transaction is used just for reading; we use the `badger.WriteBatch`
// API to bulk delete items.
txn := db.badger.NewTransaction(false)
defer txn.Discard()
reqLogIDs, err := findRequestLogIDsByProjectID(txn, projectID)
if err != nil {
return fmt.Errorf("badger: failed to find request log IDs: %w", err)
}
writeBatch := db.badger.NewWriteBatch()
defer writeBatch.Cancel()
for _, reqLogID := range reqLogIDs {
// Delete request logs.
err := writeBatch.Delete(entryKey(reqLogPrefix, 0, reqLogID[:]))
if err != nil {
return fmt.Errorf("badger: failed to delete request log: %w", err)
}
// Delete related response log.
err = writeBatch.Delete(entryKey(resLogPrefix, 0, reqLogID[:]))
if err != nil {
return fmt.Errorf("badger: failed to delete request log: %w", err)
}
}
if err := writeBatch.Flush(); err != nil {
return fmt.Errorf("badger: failed to commit batch write: %w", err)
}
err = db.badger.DropPrefix(entryKey(reqLogPrefix, reqLogProjectIDIndex, projectID[:]))
if err != nil {
return fmt.Errorf("badger: failed to drop request log project ID index items: %w", err)
}
return nil
}
func findRequestLogIDsByProjectID(txn *badger.Txn, projectID ulid.ULID) ([]ulid.ULID, error) {
reqLogIDs := make([]ulid.ULID, 0)
opts := badger.DefaultIteratorOptions
opts.PrefetchValues = false
opts.Reverse = true
iterator := txn.NewIterator(opts)
defer iterator.Close()
var projectIndexKey []byte
prefix := entryKey(reqLogPrefix, reqLogProjectIDIndex, projectID[:])
for iterator.Seek(append(prefix, 255)); iterator.ValidForPrefix(prefix); iterator.Next() {
projectIndexKey = iterator.Item().KeyCopy(projectIndexKey)
var id ulid.ULID
// The request log ID starts *after* the first 2 prefix and index bytes
// and the 16 byte project ID.
if err := id.UnmarshalBinary(projectIndexKey[18:]); err != nil {
return nil, fmt.Errorf("failed to parse request log ID: %w", err)
}
reqLogIDs = append(reqLogIDs, id)
}
return reqLogIDs, nil
}

View File

@ -1,240 +0,0 @@
package badger
import (
"bytes"
"context"
"encoding/gob"
"errors"
"fmt"
"github.com/dgraph-io/badger/v3"
"github.com/oklog/ulid"
"github.com/dstotijn/hetty/pkg/reqlog"
"github.com/dstotijn/hetty/pkg/scope"
"github.com/dstotijn/hetty/pkg/sender"
)
func (db *Database) StoreSenderRequest(ctx context.Context, req sender.Request) error {
buf := bytes.Buffer{}
err := gob.NewEncoder(&buf).Encode(req)
if err != nil {
return fmt.Errorf("badger: failed to encode sender request: %w", err)
}
entries := []*badger.Entry{
// Sender request itself.
{
Key: entryKey(senderReqPrefix, 0, req.ID[:]),
Value: buf.Bytes(),
},
// Index by project ID.
{
Key: entryKey(senderReqPrefix, senderReqProjectIDIndex, append(req.ProjectID[:], req.ID[:]...)),
},
}
err = db.badger.Update(func(txn *badger.Txn) error {
for i := range entries {
err := txn.SetEntry(entries[i])
if err != nil {
return err
}
}
return nil
})
if err != nil {
return fmt.Errorf("badger: failed to commit transaction: %w", err)
}
return nil
}
func (db *Database) FindSenderRequestByID(ctx context.Context, senderReqID ulid.ULID) (sender.Request, error) {
txn := db.badger.NewTransaction(false)
defer txn.Discard()
req, err := getSenderRequestWithResponseLog(txn, senderReqID)
if err != nil {
return sender.Request{}, fmt.Errorf("badger: failed to get sender request: %w", err)
}
return req, nil
}
func (db *Database) FindSenderRequests(
ctx context.Context,
filter sender.FindRequestsFilter,
scope *scope.Scope,
) ([]sender.Request, error) {
if filter.ProjectID.Compare(ulid.ULID{}) == 0 {
return nil, sender.ErrProjectIDMustBeSet
}
txn := db.badger.NewTransaction(false)
defer txn.Discard()
senderReqIDs, err := findSenderRequestIDsByProjectID(txn, filter.ProjectID)
if err != nil {
return nil, fmt.Errorf("badger: failed to find sender request IDs: %w", err)
}
senderReqs := make([]sender.Request, 0, len(senderReqIDs))
for _, id := range senderReqIDs {
senderReq, err := getSenderRequestWithResponseLog(txn, id)
if err != nil {
return nil, fmt.Errorf("badger: failed to get sender request (id: %v): %w", id.String(), err)
}
if filter.OnlyInScope {
if !senderReq.MatchScope(scope) {
continue
}
}
// Filter by search expression.
// TODO: Once pagination is introduced, this filter logic should be done
// as items are retrieved (e.g. when using a `badger.Iterator`).
if filter.SearchExpr != nil {
match, err := senderReq.Matches(filter.SearchExpr)
if err != nil {
return nil, fmt.Errorf(
"badger: failed to match search expression for sender request (id: %v): %w",
id.String(), err,
)
}
if !match {
continue
}
}
senderReqs = append(senderReqs, senderReq)
}
return senderReqs, nil
}
func (db *Database) DeleteSenderRequests(ctx context.Context, projectID ulid.ULID) error {
// Note: this transaction is used just for reading; we use the `badger.WriteBatch`
// API to bulk delete items.
txn := db.badger.NewTransaction(false)
defer txn.Discard()
senderReqIDs, err := findSenderRequestIDsByProjectID(txn, projectID)
if err != nil {
return fmt.Errorf("badger: failed to find sender request IDs: %w", err)
}
writeBatch := db.badger.NewWriteBatch()
defer writeBatch.Cancel()
for _, senderReqID := range senderReqIDs {
// Delete sender requests.
err := writeBatch.Delete(entryKey(senderReqPrefix, 0, senderReqID[:]))
if err != nil {
return fmt.Errorf("badger: failed to delete sender requests: %w", err)
}
// Delete related response log.
err = writeBatch.Delete(entryKey(resLogPrefix, 0, senderReqID[:]))
if err != nil {
return fmt.Errorf("badger: failed to delete request log: %w", err)
}
}
if err := writeBatch.Flush(); err != nil {
return fmt.Errorf("badger: failed to commit batch write: %w", err)
}
err = db.badger.DropPrefix(entryKey(senderReqPrefix, senderReqProjectIDIndex, projectID[:]))
if err != nil {
return fmt.Errorf("badger: failed to drop sender request project ID index items: %w", err)
}
return nil
}
func getSenderRequestWithResponseLog(txn *badger.Txn, senderReqID ulid.ULID) (sender.Request, error) {
item, err := txn.Get(entryKey(senderReqPrefix, 0, senderReqID[:]))
switch {
case errors.Is(err, badger.ErrKeyNotFound):
return sender.Request{}, sender.ErrRequestNotFound
case err != nil:
return sender.Request{}, fmt.Errorf("failed to lookup sender request item: %w", err)
}
req := sender.Request{
ID: senderReqID,
}
err = item.Value(func(rawSenderReq []byte) error {
err = gob.NewDecoder(bytes.NewReader(rawSenderReq)).Decode(&req)
if err != nil {
return fmt.Errorf("failed to decode sender request: %w", err)
}
return nil
})
if err != nil {
return sender.Request{}, fmt.Errorf("failed to retrieve or parse sender request value: %w", err)
}
item, err = txn.Get(entryKey(resLogPrefix, 0, senderReqID[:]))
if errors.Is(err, badger.ErrKeyNotFound) {
return req, nil
}
if err != nil {
return sender.Request{}, fmt.Errorf("failed to get response log: %w", err)
}
err = item.Value(func(rawReslog []byte) error {
var resLog reqlog.ResponseLog
err = gob.NewDecoder(bytes.NewReader(rawReslog)).Decode(&resLog)
if err != nil {
return fmt.Errorf("failed to decode response log: %w", err)
}
req.Response = &resLog
return nil
})
if err != nil {
return sender.Request{}, fmt.Errorf("failed to retrieve or parse response log value: %w", err)
}
return req, nil
}
func findSenderRequestIDsByProjectID(txn *badger.Txn, projectID ulid.ULID) ([]ulid.ULID, error) {
senderReqIDs := make([]ulid.ULID, 0)
opts := badger.DefaultIteratorOptions
opts.PrefetchValues = false
opts.Reverse = true
iterator := txn.NewIterator(opts)
defer iterator.Close()
var projectIndexKey []byte
prefix := entryKey(senderReqPrefix, senderReqProjectIDIndex, projectID[:])
for iterator.Seek(append(prefix, 255)); iterator.ValidForPrefix(prefix); iterator.Next() {
projectIndexKey = iterator.Item().KeyCopy(projectIndexKey)
var id ulid.ULID
// The request log ID starts *after* the first 2 prefix and index bytes
// and the 16 byte project ID.
if err := id.UnmarshalBinary(projectIndexKey[18:]); err != nil {
return nil, fmt.Errorf("failed to parse sender request ID: %w", err)
}
senderReqIDs = append(senderReqIDs, id)
}
return senderReqIDs, nil
}

60
pkg/db/bolt/bolt.go Normal file
View File

@ -0,0 +1,60 @@
package bolt
import (
"fmt"
bolt "go.etcd.io/bbolt"
)
// Database is used to store and retrieve data from an underlying Bolt database.
type Database struct {
bolt *bolt.DB
}
// OpenDatabase opens a new Bolt database.
func OpenDatabase(path string, opts *bolt.Options) (*Database, error) {
db, err := bolt.Open(path, 0o600, opts)
if err != nil {
return nil, fmt.Errorf("bolt: failed to open database: %w", err)
}
return DatabaseFromBoltDB(db)
}
// Close closes the underlying Bolt database.
func (db *Database) Close() error {
return db.bolt.Close()
}
// DatabaseFromBoltDB returns a Database with `db` set as the underlying Bolt
// database.
func DatabaseFromBoltDB(db *bolt.DB) (*Database, error) {
err := db.Update(func(tx *bolt.Tx) error {
_, err := tx.CreateBucketIfNotExists(projectsBucketName)
if err != nil {
return err
}
return nil
})
if err != nil {
return nil, fmt.Errorf("bolt: failed to create projects bucket: %w", err)
}
return &Database{bolt: db}, nil
}
func createNestedBucket(tx *bolt.Tx, names ...[]byte) (b *bolt.Bucket, err error) {
for i, name := range names {
if b == nil {
b, err = tx.CreateBucketIfNotExists(name)
} else {
b, err = b.CreateBucketIfNotExists(name)
}
if err != nil {
return nil, fmt.Errorf("bolt: failed to create nested bucket %q: %w", names[:i+1], err)
}
}
return b, nil
}

23
pkg/db/bolt/logger.go Normal file
View File

@ -0,0 +1,23 @@
package bolt
import (
bolt "go.etcd.io/bbolt"
"go.uber.org/zap"
)
// Interface guard.
var _ bolt.Logger = (*Logger)(nil)
type Logger struct {
*zap.SugaredLogger
}
// Warning implements bbolt.Logger.
func (l *Logger) Warning(v ...interface{}) {
l.Warn(v...)
}
// Warningf implements bbolt.Logger.
func (l *Logger) Warningf(format string, v ...interface{}) {
l.Warnf(format, v...)
}

176
pkg/db/bolt/proj.go Normal file
View File

@ -0,0 +1,176 @@
package bolt
import (
"bytes"
"context"
"encoding/gob"
"errors"
"fmt"
"github.com/oklog/ulid"
bolt "go.etcd.io/bbolt"
"github.com/dstotijn/hetty/pkg/proj"
)
var (
ErrProjectsBucketNotFound = errors.New("bolt: projects bucket not found")
ErrProjectBucketNotFound = errors.New("bolt: project bucket not found")
)
var (
projectsBucketName = []byte("projects")
projectKey = []byte("project")
)
func projectsBucket(tx *bolt.Tx) (*bolt.Bucket, error) {
b := tx.Bucket(projectsBucketName)
if b == nil {
return nil, ErrProjectsBucketNotFound
}
return b, nil
}
func projectBucket(tx *bolt.Tx, projectID []byte) (*bolt.Bucket, error) {
pb, err := projectsBucket(tx)
if err != nil {
return nil, err
}
b := pb.Bucket(projectID[:])
if b == nil {
return nil, ErrProjectBucketNotFound
}
return b, nil
}
func (db *Database) UpsertProject(ctx context.Context, project proj.Project) error {
buf := bytes.Buffer{}
err := gob.NewEncoder(&buf).Encode(project)
if err != nil {
return fmt.Errorf("bolt: failed to encode project: %w", err)
}
err = db.bolt.Update(func(tx *bolt.Tx) error {
b, err := createNestedBucket(tx, projectsBucketName, project.ID[:])
if err != nil {
return fmt.Errorf("bolt: failed to create project bucket: %w", err)
}
err = b.Put(projectKey, buf.Bytes())
if err != nil {
return fmt.Errorf("bolt: failed to upsert project: %w", err)
}
_, err = b.CreateBucketIfNotExists(reqLogsBucketName)
if err != nil {
return fmt.Errorf("bolt: failed to create request logs bucket: %w", err)
}
_, err = b.CreateBucketIfNotExists(senderReqsBucketName)
if err != nil {
return fmt.Errorf("bolt: failed to create sender requests bucket: %w", err)
}
return nil
})
if err != nil {
return fmt.Errorf("bolt: failed to commit transaction: %w", err)
}
return nil
}
func (db *Database) FindProjectByID(ctx context.Context, projectID ulid.ULID) (project proj.Project, err error) {
err = db.bolt.View(func(tx *bolt.Tx) error {
bucket, err := projectBucket(tx, projectID[:])
if errors.Is(err, ErrProjectsBucketNotFound) || errors.Is(err, ErrProjectBucketNotFound) {
return proj.ErrProjectNotFound
}
if err != nil {
return err
}
rawProject := bucket.Get(projectKey)
if rawProject == nil {
return proj.ErrProjectNotFound
}
err = gob.NewDecoder(bytes.NewReader(rawProject)).Decode(&project)
if err != nil {
return fmt.Errorf("failed to decode project: %w", err)
}
return nil
})
if err != nil {
return proj.Project{}, fmt.Errorf("bolt: failed to commit transaction: %w", err)
}
return project, nil
}
func (db *Database) DeleteProject(ctx context.Context, projectID ulid.ULID) error {
err := db.bolt.Update(func(tx *bolt.Tx) error {
pb, err := projectsBucket(tx)
if err != nil {
return err
}
err = pb.DeleteBucket(projectID[:])
if err != nil {
return fmt.Errorf("failed to delete project bucket: %w", err)
}
return nil
})
if err != nil {
return fmt.Errorf("bolt: failed to commit transaction: %w", err)
}
return nil
}
func (db *Database) Projects(ctx context.Context) ([]proj.Project, error) {
projects := make([]proj.Project, 0)
err := db.bolt.View(func(tx *bolt.Tx) error {
pb, err := projectsBucket(tx)
if err != nil {
return err
}
err = pb.ForEachBucket(func(projectID []byte) error {
bucket, err := projectBucket(tx, projectID)
if err != nil {
return err
}
rawProject := bucket.Get(projectKey)
if rawProject == nil {
return proj.ErrProjectNotFound
}
var project proj.Project
err = gob.NewDecoder(bytes.NewReader(rawProject)).Decode(&project)
if err != nil {
return fmt.Errorf("bolt: failed to decode project: %w", err)
}
projects = append(projects, project)
return nil
})
if err != nil {
return fmt.Errorf("bolt: failed to iterate over projects: %w", err)
}
return nil
})
if err != nil {
return nil, fmt.Errorf("bolt: failed to commit transaction: %w", err)
}
return projects, nil
}

267
pkg/db/bolt/proj_test.go Normal file
View File

@ -0,0 +1,267 @@
package bolt_test
import (
"bytes"
"context"
"encoding/gob"
"errors"
"math/rand"
"regexp"
"testing"
"time"
"github.com/google/go-cmp/cmp"
"github.com/google/go-cmp/cmp/cmpopts"
"github.com/oklog/ulid"
"go.etcd.io/bbolt"
"github.com/dstotijn/hetty/pkg/db/bolt"
"github.com/dstotijn/hetty/pkg/filter"
"github.com/dstotijn/hetty/pkg/proj"
"github.com/dstotijn/hetty/pkg/scope"
)
//nolint:gosec
var ulidEntropy = rand.New(rand.NewSource(time.Now().UnixNano()))
var regexpCompareOpt = cmp.Comparer(func(x, y *regexp.Regexp) bool {
switch {
case x == nil && y == nil:
return true
case x == nil || y == nil:
return false
default:
return x.String() == y.String()
}
})
func TestUpsertProject(t *testing.T) {
t.Parallel()
path := t.TempDir() + "bolt.db"
boltDB, err := bbolt.Open(path, 0o600, nil)
if err != nil {
t.Fatalf("failed to open bolt database: %v", err)
}
db, err := bolt.DatabaseFromBoltDB(boltDB)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
searchExpr, err := filter.ParseQuery("foo AND bar OR NOT baz")
if err != nil {
t.Fatalf("unexpected error (expected: nil, got: %v)", err)
}
exp := proj.Project{
ID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
Name: "foobar",
Settings: proj.Settings{
ReqLogBypassOutOfScope: true,
ReqLogOnlyFindInScope: true,
ReqLogSearchExpr: searchExpr,
ScopeRules: []scope.Rule{
{
URL: regexp.MustCompile("^https://(.*)example.com(.*)$"),
Header: scope.Header{
Key: regexp.MustCompile("^X-Foo(.*)$"),
Value: regexp.MustCompile("^foo(.*)$"),
},
Body: regexp.MustCompile("^foo(.*)"),
},
},
},
}
err = db.UpsertProject(context.Background(), exp)
if err != nil {
t.Fatalf("unexpected error storing project: %v", err)
}
var rawProject []byte
err = boltDB.View(func(tx *bbolt.Tx) error {
rawProject = tx.Bucket([]byte("projects")).Bucket(exp.ID[:]).Get([]byte("project"))
return nil
})
if err != nil {
t.Fatalf("unexpected error retrieving project from database: %v", err)
}
if rawProject == nil {
t.Fatalf("expected raw project to be retrieved, got: nil")
}
got := proj.Project{}
err = gob.NewDecoder(bytes.NewReader(rawProject)).Decode(&got)
if err != nil {
t.Fatalf("unexpected error decoding project: %v", err)
}
if diff := cmp.Diff(exp, got, regexpCompareOpt, cmpopts.IgnoreUnexported(proj.Project{})); diff != "" {
t.Fatalf("project not equal (-exp, +got):\n%v", diff)
}
}
func TestFindProjectByID(t *testing.T) {
t.Parallel()
t.Run("existing project", func(t *testing.T) {
t.Parallel()
path := t.TempDir() + "bolt.db"
boltDB, err := bbolt.Open(path, 0o600, nil)
if err != nil {
t.Fatalf("failed to open bolt database: %v", err)
}
db, err := bolt.DatabaseFromBoltDB(boltDB)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
exp := proj.Project{
ID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
}
buf := bytes.Buffer{}
err = gob.NewEncoder(&buf).Encode(exp)
if err != nil {
t.Fatalf("unexpected error encoding project: %v", err)
}
err = boltDB.Update(func(tx *bbolt.Tx) error {
b, err := tx.Bucket([]byte("projects")).CreateBucket(exp.ID[:])
if err != nil {
return err
}
return b.Put([]byte("project"), buf.Bytes())
})
if err != nil {
t.Fatalf("unexpected error setting project: %v", err)
}
got, err := db.FindProjectByID(context.Background(), exp.ID)
if err != nil {
t.Fatalf("unexpected error finding project: %v", err)
}
if diff := cmp.Diff(exp, got, cmpopts.IgnoreUnexported(proj.Project{})); diff != "" {
t.Fatalf("project not equal (-exp, +got):\n%v", diff)
}
})
t.Run("project not found", func(t *testing.T) {
t.Parallel()
path := t.TempDir() + "bolt.db"
boltDB, err := bbolt.Open(path, 0o600, nil)
if err != nil {
t.Fatalf("failed to open bolt database: %v", err)
}
db, err := bolt.DatabaseFromBoltDB(boltDB)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
projectID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
_, err = db.FindProjectByID(context.Background(), projectID)
if !errors.Is(err, proj.ErrProjectNotFound) {
t.Fatalf("expected `proj.ErrProjectNotFound`, got: %v", err)
}
})
}
func TestDeleteProject(t *testing.T) {
t.Parallel()
path := t.TempDir() + "bolt.db"
boltDB, err := bbolt.Open(path, 0o600, nil)
if err != nil {
t.Fatalf("failed to open bolt database: %v", err)
}
db, err := bolt.DatabaseFromBoltDB(boltDB)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
// Insert test fixture.
projectID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
err = db.UpsertProject(context.Background(), proj.Project{
ID: projectID,
})
if err != nil {
t.Fatalf("unexpected error storing project: %v", err)
}
err = db.DeleteProject(context.Background(), projectID)
if err != nil {
t.Fatalf("unexpected error deleting project: %v", err)
}
var got *bbolt.Bucket
err = boltDB.View(func(tx *bbolt.Tx) error {
got = tx.Bucket([]byte("projects")).Bucket(projectID[:])
return nil
})
if got != nil {
t.Fatalf("expected bucket to be nil, got: %v", got)
}
}
func TestProjects(t *testing.T) {
t.Parallel()
path := t.TempDir() + "bolt.db"
boltDB, err := bbolt.Open(path, 0o600, nil)
if err != nil {
t.Fatalf("failed to open bolt database: %v", err)
}
db, err := bolt.DatabaseFromBoltDB(boltDB)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
exp := []proj.Project{
{
ID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
Name: "one",
},
{
ID: ulid.MustNew(ulid.Timestamp(time.Now())+100, ulidEntropy),
Name: "two",
},
}
// Store fixtures.
for _, project := range exp {
err = db.UpsertProject(context.Background(), project)
if err != nil {
t.Fatalf("unexpected error creating project fixture: %v", err)
}
}
got, err := db.Projects(context.Background())
if err != nil {
t.Fatalf("unexpected error finding projects: %v", err)
}
if len(exp) != len(got) {
t.Fatalf("expected %v projects, got: %v", len(exp), len(got))
}
if diff := cmp.Diff(exp, got, cmpopts.IgnoreUnexported(proj.Project{})); diff != "" {
t.Fatalf("projects not equal (-exp, +got):\n%v", diff)
}
}

204
pkg/db/bolt/reqlog.go Normal file
View File

@ -0,0 +1,204 @@
package bolt
import (
"bytes"
"context"
"encoding/gob"
"errors"
"fmt"
"github.com/oklog/ulid"
bolt "go.etcd.io/bbolt"
"github.com/dstotijn/hetty/pkg/reqlog"
"github.com/dstotijn/hetty/pkg/scope"
)
var ErrRequestLogsBucketNotFound = errors.New("bolt: request logs bucket not found")
var reqLogsBucketName = []byte("request_logs")
func requestLogsBucket(tx *bolt.Tx, projectID ulid.ULID) (*bolt.Bucket, error) {
pb, err := projectBucket(tx, projectID[:])
if err != nil {
return nil, err
}
b := pb.Bucket(reqLogsBucketName)
if b == nil {
return nil, ErrRequestLogsBucketNotFound
}
return b, nil
}
func (db *Database) FindRequestLogs(ctx context.Context, filter reqlog.FindRequestsFilter, scope *scope.Scope) (reqLogs []reqlog.RequestLog, err error) {
if filter.ProjectID.Compare(ulid.ULID{}) == 0 {
return nil, reqlog.ErrProjectIDMustBeSet
}
tx, err := db.bolt.Begin(false)
if err != nil {
return nil, fmt.Errorf("bolt: failed to begin transaction: %w", err)
}
defer tx.Rollback()
b, err := requestLogsBucket(tx, filter.ProjectID)
if err != nil {
return nil, fmt.Errorf("bolt: failed to get request logs bucket: %w", err)
}
err = b.ForEach(func(reqLogID, rawReqLog []byte) error {
var reqLog reqlog.RequestLog
err = gob.NewDecoder(bytes.NewReader(rawReqLog)).Decode(&reqLog)
if err != nil {
return fmt.Errorf("failed to decode request log: %w", err)
}
if filter.OnlyInScope && !reqLog.MatchScope(scope) {
return nil
}
// Filter by search expression. TODO: Once pagination is introduced,
// this filter logic should be done as items are retrieved.
if filter.SearchExpr != nil {
match, err := reqLog.Matches(filter.SearchExpr)
if err != nil {
return fmt.Errorf("failed to match search expression for request log (id: %v): %w", reqLogID, err)
}
if !match {
return nil
}
}
reqLogs = append(reqLogs, reqLog)
return nil
})
if err != nil {
return nil, fmt.Errorf("bolt: failed to iterate over request logs: %w", err)
}
// Reverse items, so newest requests appear first.
for i, j := 0, len(reqLogs)-1; i < j; i, j = i+1, j-1 {
reqLogs[i], reqLogs[j] = reqLogs[j], reqLogs[i]
}
return reqLogs, nil
}
func (db *Database) FindRequestLogByID(ctx context.Context, projectID, reqLogID ulid.ULID) (reqLog reqlog.RequestLog, err error) {
err = db.bolt.View(func(tx *bolt.Tx) error {
b, err := requestLogsBucket(tx, projectID)
if err != nil {
return fmt.Errorf("bolt: failed to get request logs bucket: %w", err)
}
rawReqLog := b.Get(reqLogID[:])
if rawReqLog == nil {
return reqlog.ErrRequestNotFound
}
err = gob.NewDecoder(bytes.NewReader(rawReqLog)).Decode(&reqLog)
if err != nil {
return fmt.Errorf("failed to decode request log: %w", err)
}
return nil
})
if err != nil {
return reqlog.RequestLog{}, fmt.Errorf("bolt: failed to find request log by ID: %w", err)
}
return reqLog, nil
}
func (db *Database) StoreRequestLog(ctx context.Context, reqLog reqlog.RequestLog) error {
buf := bytes.Buffer{}
err := gob.NewEncoder(&buf).Encode(reqLog)
if err != nil {
return fmt.Errorf("bolt: failed to encode request log: %w", err)
}
err = db.bolt.Update(func(txn *bolt.Tx) error {
b, err := requestLogsBucket(txn, reqLog.ProjectID)
if err != nil {
return fmt.Errorf("failed to get request logs bucket: %w", err)
}
err = b.Put(reqLog.ID[:], buf.Bytes())
if err != nil {
return fmt.Errorf("failed to put request log: %w", err)
}
return nil
})
if err != nil {
return fmt.Errorf("bolt: failed to commit transaction: %w", err)
}
return nil
}
func (db *Database) StoreResponseLog(ctx context.Context, projectID, reqLogID ulid.ULID, resLog reqlog.ResponseLog) error {
buf := bytes.Buffer{}
err := gob.NewEncoder(&buf).Encode(resLog)
if err != nil {
return fmt.Errorf("bolt: failed to encode response log: %w", err)
}
err = db.bolt.Update(func(txn *bolt.Tx) error {
b, err := requestLogsBucket(txn, projectID)
if err != nil {
return fmt.Errorf("failed to get request logs bucket: %w", err)
}
rawReqLog := b.Get(reqLogID[:])
if rawReqLog == nil {
return reqlog.ErrRequestNotFound
}
var reqLog reqlog.RequestLog
err = gob.NewDecoder(bytes.NewReader(rawReqLog)).Decode(&reqLog)
if err != nil {
return fmt.Errorf("failed to decode request log: %w", err)
}
reqLog.Response = &resLog
buf := bytes.Buffer{}
err = gob.NewEncoder(&buf).Encode(reqLog)
if err != nil {
return fmt.Errorf("failed to encode request log: %w", err)
}
err = b.Put(reqLog.ID[:], buf.Bytes())
if err != nil {
return fmt.Errorf("failed to put request log: %w", err)
}
return nil
})
if err != nil {
return fmt.Errorf("bolt: failed to commit transaction: %w", err)
}
return nil
}
func (db *Database) ClearRequestLogs(ctx context.Context, projectID ulid.ULID) error {
err := db.bolt.Update(func(txn *bolt.Tx) error {
pb, err := projectBucket(txn, projectID[:])
if err != nil {
return fmt.Errorf("failed to get project bucket: %w", err)
}
return pb.DeleteBucket(reqLogsBucketName)
})
if err != nil {
return fmt.Errorf("bolt: failed to commit transaction: %w", err)
}
return nil
}

View File

@ -1,4 +1,4 @@
package badger
package bolt_test
import (
"context"
@ -8,10 +8,12 @@ import (
"testing"
"time"
badgerdb "github.com/dgraph-io/badger/v3"
"github.com/google/go-cmp/cmp"
"github.com/oklog/ulid"
"go.etcd.io/bbolt"
"github.com/dstotijn/hetty/pkg/db/bolt"
"github.com/dstotijn/hetty/pkg/proj"
"github.com/dstotijn/hetty/pkg/reqlog"
)
@ -21,15 +23,21 @@ func TestFindRequestLogs(t *testing.T) {
t.Run("without project ID in filter", func(t *testing.T) {
t.Parallel()
database, err := OpenDatabase(badgerdb.DefaultOptions("").WithInMemory(true))
path := t.TempDir() + "bolt.db"
boltDB, err := bbolt.Open(path, 0o600, nil)
if err != nil {
t.Fatalf("failed to open badger database: %v", err)
t.Fatalf("failed to open bolt database: %v", err)
}
defer database.Close()
db, err := bolt.DatabaseFromBoltDB(boltDB)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
filter := reqlog.FindRequestsFilter{}
_, err = database.FindRequestLogs(context.Background(), filter, nil)
_, err = db.FindRequestLogs(context.Background(), filter, nil)
if !errors.Is(err, reqlog.ErrProjectIDMustBeSet) {
t.Fatalf("expected `reqlog.ErrProjectIDMustBeSet`, got: %v", err)
}
@ -38,14 +46,27 @@ func TestFindRequestLogs(t *testing.T) {
t.Run("returns request logs and related response logs", func(t *testing.T) {
t.Parallel()
database, err := OpenDatabase(badgerdb.DefaultOptions("").WithInMemory(true))
path := t.TempDir() + "bolt.db"
boltDB, err := bbolt.Open(path, 0o600, nil)
if err != nil {
t.Fatalf("failed to open badger database: %v", err)
t.Fatalf("failed to open bolt database: %v", err)
}
defer database.Close()
db, err := bolt.DatabaseFromBoltDB(boltDB)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
projectID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
err = db.UpsertProject(context.Background(), proj.Project{
ID: projectID,
})
if err != nil {
t.Fatalf("unexpected error upserting project: %v", err)
}
fixtures := []reqlog.RequestLog{
{
ID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
@ -81,24 +102,17 @@ func TestFindRequestLogs(t *testing.T) {
// Store fixtures.
for _, reqLog := range fixtures {
err = database.StoreRequestLog(context.Background(), reqLog)
err = db.StoreRequestLog(context.Background(), reqLog)
if err != nil {
t.Fatalf("unexpected error creating request log fixture: %v", err)
}
if reqLog.Response != nil {
err = database.StoreResponseLog(context.Background(), reqLog.ID, *reqLog.Response)
if err != nil {
t.Fatalf("unexpected error creating response log fixture: %v", err)
}
}
}
filter := reqlog.FindRequestsFilter{
ProjectID: projectID,
}
got, err := database.FindRequestLogs(context.Background(), filter, nil)
got, err := db.FindRequestLogs(context.Background(), filter, nil)
if err != nil {
t.Fatalf("unexpected error finding request logs: %v", err)
}

172
pkg/db/bolt/sender.go Normal file
View File

@ -0,0 +1,172 @@
package bolt
import (
"bytes"
"context"
"encoding/gob"
"errors"
"fmt"
"github.com/oklog/ulid"
bolt "go.etcd.io/bbolt"
"github.com/dstotijn/hetty/pkg/scope"
"github.com/dstotijn/hetty/pkg/sender"
)
var ErrSenderRequestsBucketNotFound = errors.New("bolt: sender requests bucket not found")
var senderReqsBucketName = []byte("sender_requests")
func senderReqsBucket(tx *bolt.Tx, projectID ulid.ULID) (*bolt.Bucket, error) {
pb, err := projectBucket(tx, projectID[:])
if err != nil {
return nil, err
}
b := pb.Bucket(senderReqsBucketName)
if b == nil {
return nil, ErrSenderRequestsBucketNotFound
}
return b, nil
}
func (db *Database) StoreSenderRequest(ctx context.Context, req sender.Request) error {
buf := bytes.Buffer{}
err := gob.NewEncoder(&buf).Encode(req)
if err != nil {
return fmt.Errorf("bolt: failed to encode sender request: %w", err)
}
err = db.bolt.Update(func(tx *bolt.Tx) error {
senderReqsBucket, err := senderReqsBucket(tx, req.ProjectID)
if err != nil {
return fmt.Errorf("failed to get sender requests bucket: %w", err)
}
err = senderReqsBucket.Put(req.ID[:], buf.Bytes())
if err != nil {
return fmt.Errorf("failed to put sender request: %w", err)
}
return nil
})
if err != nil {
return fmt.Errorf("bolt: failed to commit transaction: %w", err)
}
return nil
}
func (db *Database) FindSenderRequestByID(ctx context.Context, projectID, senderReqID ulid.ULID) (req sender.Request, err error) {
if projectID.Compare(ulid.ULID{}) == 0 {
return sender.Request{}, sender.ErrProjectIDMustBeSet
}
err = db.bolt.View(func(tx *bolt.Tx) error {
senderReqsBucket, err := senderReqsBucket(tx, projectID)
if err != nil {
return fmt.Errorf("failed to get sender requests bucket: %w", err)
}
rawSenderReq := senderReqsBucket.Get(senderReqID[:])
if rawSenderReq == nil {
return sender.ErrRequestNotFound
}
err = gob.NewDecoder(bytes.NewReader(rawSenderReq)).Decode(&req)
if err != nil {
return fmt.Errorf("failed to decode sender request: %w", err)
}
return nil
})
if err != nil {
return sender.Request{}, fmt.Errorf("bolt: failed to commit transaction: %w", err)
}
return req, nil
}
func (db *Database) FindSenderRequests(ctx context.Context, filter sender.FindRequestsFilter, scope *scope.Scope) (reqs []sender.Request, err error) {
if filter.ProjectID.Compare(ulid.ULID{}) == 0 {
return nil, sender.ErrProjectIDMustBeSet
}
tx, err := db.bolt.Begin(false)
if err != nil {
return nil, fmt.Errorf("bolt: failed to begin transaction: %w", err)
}
defer tx.Rollback()
b, err := senderReqsBucket(tx, filter.ProjectID)
if err != nil {
return nil, fmt.Errorf("failed to get sender requests bucket: %w", err)
}
err = b.ForEach(func(senderReqID, rawSenderReq []byte) error {
var req sender.Request
err = gob.NewDecoder(bytes.NewReader(rawSenderReq)).Decode(&req)
if err != nil {
return fmt.Errorf("failed to decode sender request: %w", err)
}
if filter.OnlyInScope {
if !req.MatchScope(scope) {
return nil
}
}
// Filter by search expression. TODO: Once pagination is introduced,
// this filter logic should be done as items are retrieved.
if filter.SearchExpr != nil {
match, err := req.Matches(filter.SearchExpr)
if err != nil {
return fmt.Errorf(
"bolt: failed to match search expression for sender request (id: %v): %w",
senderReqID, err,
)
}
if !match {
return nil
}
}
reqs = append(reqs, req)
return nil
})
if err != nil {
return nil, fmt.Errorf("bolt: failed to commit transaction: %w", err)
}
// Reverse items, so newest requests appear first.
for i, j := 0, len(reqs)-1; i < j; i, j = i+1, j-1 {
reqs[i], reqs[j] = reqs[j], reqs[i]
}
return reqs, nil
}
func (db *Database) DeleteSenderRequests(ctx context.Context, projectID ulid.ULID) error {
err := db.bolt.Update(func(tx *bolt.Tx) error {
senderReqsBucket, err := senderReqsBucket(tx, projectID)
if err != nil {
return fmt.Errorf("failed to get sender requests bucket: %w", err)
}
err = senderReqsBucket.DeleteBucket(senderReqsBucketName)
if err != nil {
return fmt.Errorf("failed to delete sender requests bucket: %w", err)
}
return nil
})
if err != nil {
return fmt.Errorf("bolt: failed to commit transaction: %w", err)
}
return nil
}

View File

@ -1,26 +1,23 @@
package badger_test
package bolt_test
import (
"context"
"errors"
"math/rand"
"net/http"
"net/url"
"testing"
"time"
badgerdb "github.com/dgraph-io/badger/v3"
"github.com/google/go-cmp/cmp"
"github.com/oklog/ulid"
"go.etcd.io/bbolt"
"github.com/dstotijn/hetty/pkg/db/badger"
"github.com/dstotijn/hetty/pkg/db/bolt"
"github.com/dstotijn/hetty/pkg/proj"
"github.com/dstotijn/hetty/pkg/reqlog"
"github.com/dstotijn/hetty/pkg/sender"
)
//nolint:gosec
var ulidEntropy = rand.New(rand.NewSource(time.Now().UnixNano()))
var exampleURL = func() *url.URL {
u, err := url.Parse("https://example.com/foobar")
if err != nil {
@ -33,18 +30,35 @@ var exampleURL = func() *url.URL {
func TestFindRequestByID(t *testing.T) {
t.Parallel()
database, err := badger.OpenDatabase(badgerdb.DefaultOptions("").WithInMemory(true))
path := t.TempDir() + "bolt.db"
boltDB, err := bbolt.Open(path, 0o600, nil)
if err != nil {
t.Fatalf("failed to open badger database: %v", err)
t.Fatalf("failed to open bolt database: %v", err)
}
defer boltDB.Close()
db, err := bolt.DatabaseFromBoltDB(boltDB)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
projectID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
reqID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
err = db.UpsertProject(context.Background(), proj.Project{
ID: projectID,
})
if err != nil {
t.Fatalf("unexpected error upserting project: %v", err)
}
defer database.Close()
// See: https://go.dev/blog/subtests#cleaning-up-after-a-group-of-parallel-tests
t.Run("group", func(t *testing.T) {
t.Run("sender request not found", func(t *testing.T) {
t.Parallel()
_, err := database.FindSenderRequestByID(context.Background(), ulid.ULID{})
_, err := db.FindSenderRequestByID(context.Background(), projectID, reqID)
if !errors.Is(err, sender.ErrRequestNotFound) {
t.Fatalf("expected `sender.ErrRequestNotFound`, got: %v", err)
}
@ -55,7 +69,7 @@ func TestFindRequestByID(t *testing.T) {
exp := sender.Request{
ID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
ProjectID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
ProjectID: projectID,
SourceRequestLogID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
URL: exampleURL,
@ -65,31 +79,23 @@ func TestFindRequestByID(t *testing.T) {
"X-Foo": []string{"bar"},
},
Body: []byte("foo"),
}
err := database.StoreSenderRequest(context.Background(), exp)
if err != nil {
t.Fatalf("unexpected error (expected: nil, got: %v)", err)
}
resLog := reqlog.ResponseLog{
Proto: "HTTP/2.0",
Status: "200 OK",
StatusCode: 200,
Header: http.Header{
"X-Yolo": []string{"swag"},
Response: &reqlog.ResponseLog{
Proto: "HTTP/2.0",
Status: "200 OK",
StatusCode: 200,
Header: http.Header{
"X-Yolo": []string{"swag"},
},
Body: []byte("bar"),
},
Body: []byte("bar"),
}
err = database.StoreResponseLog(context.Background(), exp.ID, resLog)
err := db.StoreSenderRequest(context.Background(), exp)
if err != nil {
t.Fatalf("unexpected error (expected: nil, got: %v)", err)
}
exp.Response = &resLog
got, err := database.FindSenderRequestByID(context.Background(), exp.ID)
got, err := db.FindSenderRequestByID(context.Background(), exp.ProjectID, exp.ID)
if err != nil {
t.Fatalf("unexpected error (expected: nil, got: %v)", err)
}
@ -107,15 +113,22 @@ func TestFindSenderRequests(t *testing.T) {
t.Run("without project ID in filter", func(t *testing.T) {
t.Parallel()
database, err := badger.OpenDatabase(badgerdb.DefaultOptions("").WithInMemory(true))
path := t.TempDir() + "bolt.db"
boltDB, err := bbolt.Open(path, 0o600, nil)
if err != nil {
t.Fatalf("failed to open badger database: %v", err)
t.Fatalf("failed to open bolt database: %v", err)
}
defer database.Close()
defer boltDB.Close()
db, err := bolt.DatabaseFromBoltDB(boltDB)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
filter := sender.FindRequestsFilter{}
_, err = database.FindSenderRequests(context.Background(), filter, nil)
_, err = db.FindSenderRequests(context.Background(), filter, nil)
if !errors.Is(err, sender.ErrProjectIDMustBeSet) {
t.Fatalf("expected `sender.ErrProjectIDMustBeSet`, got: %v", err)
}
@ -124,14 +137,30 @@ func TestFindSenderRequests(t *testing.T) {
t.Run("returns sender requests and related response logs", func(t *testing.T) {
t.Parallel()
database, err := badger.OpenDatabase(badgerdb.DefaultOptions("").WithInMemory(true))
path := t.TempDir() + "bolt.db"
boltDB, err := bbolt.Open(path, 0o600, nil)
if err != nil {
t.Fatalf("failed to open badger database: %v", err)
t.Fatalf("failed to open bolt database: %v", err)
}
defer database.Close()
defer boltDB.Close()
db, err := bolt.DatabaseFromBoltDB(boltDB)
if err != nil {
t.Fatalf("failed to create database: %v", err)
}
defer db.Close()
projectID := ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy)
err = db.UpsertProject(context.Background(), proj.Project{
ID: projectID,
Name: "foobar",
Settings: proj.Settings{},
})
if err != nil {
t.Fatalf("unexpected error creating project (expected: nil, got: %v)", err)
}
fixtures := []sender.Request{
{
ID: ulid.MustNew(ulid.Timestamp(time.Now()), ulidEntropy),
@ -169,24 +198,17 @@ func TestFindSenderRequests(t *testing.T) {
// Store fixtures.
for _, senderReq := range fixtures {
err = database.StoreSenderRequest(context.Background(), senderReq)
err = db.StoreSenderRequest(context.Background(), senderReq)
if err != nil {
t.Fatalf("unexpected error creating request log fixture: %v", err)
}
if senderReq.Response != nil {
err = database.StoreResponseLog(context.Background(), senderReq.ID, *senderReq.Response)
if err != nil {
t.Fatalf("unexpected error creating response log fixture: %v", err)
}
}
}
filter := sender.FindRequestsFilter{
ProjectID: projectID,
}
got, err := database.FindSenderRequests(context.Background(), filter, nil)
got, err := db.FindSenderRequests(context.Background(), filter, nil)
if err != nil {
t.Fatalf("unexpected error finding sender requests: %v", err)
}