From ce928524cab9a83395f3772ae9dd5d7732af281d Mon Sep 17 00:00:00 2001
From: Vladimir Mihailenco <vladimir.webdev@gmail.com>
Date: Thu, 14 Oct 2021 16:13:05 +0300
Subject: [PATCH] feat(bundebug): allow to configure the hook using env var,
 for example, BUNDEBUG={0,1,2}

---
 example/basic/README.md |  6 ++++++
 example/basic/main.go   |  5 ++++-
 extra/bundebug/debug.go | 36 +++++++++++++++++++++++++++++++++++-
 3 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/example/basic/README.md b/example/basic/README.md
index 162603b9..2155a8be 100644
--- a/example/basic/README.md
+++ b/example/basic/README.md
@@ -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.
diff --git a/example/basic/main.go b/example/basic/main.go
index 67a9a8be..505ec54b 100644
--- a/example/basic/main.go
+++ b/example/basic/main.go
@@ -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))
diff --git a/extra/bundebug/debug.go b/extra/bundebug/debug.go
index a4d4a9e0..d5897016 100644
--- a/extra/bundebug/debug.go
+++ b/extra/bundebug/debug.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"fmt"
+	"os"
 	"reflect"
 	"time"
 
@@ -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)
 	}
@@ -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: