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 }