feat: initial commit

chdebug/debug.go

@@ -0,0 +1,133 @@
+package chdebug
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+	"io"
+	"os"
+	"reflect"
+	"time"
+
+	"github.com/fatih/color"
+	"github.com/uptrace/go-clickhouse/ch"
+)
+
+type Option func(*QueryHook)
+
+// WithEnabled enables/disables the hook.
+func WithEnabled(on bool) Option {
+	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) Option {
+	return func(h *QueryHook) {
+		h.verbose = on
+	}
+}
+
+// WithWriter sets the log output to an io.Writer
+// the default is os.Stderr
+func WithWriter(w io.Writer) Option {
+	return func(h *QueryHook) {
+		h.writer = w
+	}
+}
+
+// FromEnv configures the hook using the environment variable value.
+// For example, WithEnv("CHDEBUG"):
+//    - CHDEBUG=0 - disables the hook.
+//    - CHDEBUG=1 - enables the hook.
+//    - CHDEBUG=2 - enables the hook and verbose mode.
+func FromEnv(key string) Option {
+	if key == "" {
+		key = "CHDEBUG"
+	}
+	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
+	writer  io.Writer
+}
+
+var _ ch.QueryHook = (*QueryHook)(nil)
+
+func NewQueryHook(opts ...Option) *QueryHook {
+	h := &QueryHook{
+		enabled: true,
+		writer:  os.Stderr,
+	}
+	for _, opt := range opts {
+		opt(h)
+	}
+	return h
+}
+
+func (h *QueryHook) BeforeQuery(ctx context.Context, evt *ch.QueryEvent) context.Context {
+	return ctx
+}
+
+func (h *QueryHook) AfterQuery(ctx context.Context, event *ch.QueryEvent) {
+	if !h.enabled {
+		return
+	}
+
+	if !h.verbose {
+		switch event.Err {
+		case nil, sql.ErrNoRows:
+			return
+		}
+	}
+
+	now := time.Now()
+	dur := now.Sub(event.StartTime)
+
+	args := []any{
+		"[ch]",
+		now.Format(" 15:04:05.000 "),
+		formatOperation(event),
+		fmt.Sprintf(" %10s ", dur.Round(time.Microsecond)),
+		event.Query,
+	}
+
+	if event.Err != nil {
+		typ := reflect.TypeOf(event.Err).String()
+		args = append(args,
+			"\t",
+			color.New(color.BgRed).Sprintf(" %s ", typ+": "+event.Err.Error()),
+		)
+	}
+
+	fmt.Fprintln(h.writer, args...)
+}
+
+func formatOperation(event *ch.QueryEvent) string {
+	operation := event.Operation()
+	return operationColor(operation).Sprintf(" %-16s ", operation)
+}
+
+func operationColor(operation string) *color.Color {
+	switch operation {
+	case "SELECT":
+		return color.New(color.BgGreen, color.FgHiWhite)
+	case "INSERT":
+		return color.New(color.BgBlue, color.FgHiWhite)
+	case "UPDATE":
+		return color.New(color.BgYellow, color.FgHiBlack)
+	case "DELETE":
+		return color.New(color.BgMagenta, color.FgHiWhite)
+	default:
+		return color.New(color.BgWhite, color.FgHiBlack)
+	}
+}