Skip to content

Commit

Permalink
feat(bundebug): allow to configure the hook using env var, for exampl…
Browse files Browse the repository at this point in the history
…e, BUNDEBUG={0,1,2}
  • Loading branch information
vmihailenco committed Oct 14, 2021
1 parent b7a275f commit ce92852
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 2 deletions.
6 changes: 6 additions & 0 deletions example/basic/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,4 +6,10 @@ To run this example:
go run .
```

To disable query logging:

```shell
BUNDEBUG=0 go run .
```

See [docs](https://bun.uptrace.dev/) for details.
5 changes: 4 additions & 1 deletion example/basic/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,10 @@ func main() {
sqlite.SetMaxOpenConns(1)

db := bun.NewDB(sqlite, sqlitedialect.New())
db.AddQueryHook(bundebug.NewQueryHook(bundebug.WithVerbose(true)))
db.AddQueryHook(bundebug.NewQueryHook(
bundebug.WithVerbose(true),
bundebug.FromEnv("BUNDEBUG"),
))

// Register models for the fixture.
db.RegisterModel((*User)(nil), (*Story)(nil))
Expand Down
36 changes: 35 additions & 1 deletion extra/bundebug/debug.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"context"
"database/sql"
"fmt"
"os"
"reflect"
"time"

Expand All @@ -14,20 +15,49 @@ import (

type ConfigOption func(*QueryHook)

// WithEnabled enables/disables the hook.
func WithEnabled(on bool) ConfigOption {
return func(h *QueryHook) {
h.enabled = on
}
}

// WithVerbose configures the hook to log all queries
// (by default, only failed queries are logged).
func WithVerbose(on bool) ConfigOption {
return func(h *QueryHook) {
h.verbose = on
}
}

// WithEnv configures the hook using then environment variable value.
// For example, WithEnv("BUNDEBUG"):
// - BUNDEBUG=0 - disables the hook.
// - BUNDEBUG=1 - enables the hook.
// - BUNDEBUG=2 - enables the hook and verbose mode.
func FromEnv(key string) ConfigOption {
if key == "" {
key = "BUNDEBUG"
}
return func(h *QueryHook) {
if env, ok := os.LookupEnv(key); ok {
h.enabled = env != "" && env != "0"
h.verbose = env == "2"
}
}
}

type QueryHook struct {
enabled bool
verbose bool
}

var _ bun.QueryHook = (*QueryHook)(nil)

func NewQueryHook(opts ...ConfigOption) *QueryHook {
h := new(QueryHook)
h := &QueryHook{
enabled: true,
}
for _, opt := range opts {
opt(h)
}
Expand All @@ -41,6 +71,10 @@ func (h *QueryHook) BeforeQuery(
}

func (h *QueryHook) AfterQuery(ctx context.Context, event *bun.QueryEvent) {
if !h.enabled {
return
}

if !h.verbose {
switch event.Err {
case nil, sql.ErrNoRows:
Expand Down

0 comments on commit ce92852

Please sign in to comment.