feat: add Event.Operation

extra/bundebug/debug.go

@@ -72,24 +72,13 @@ func (h *QueryHook) AfterQuery(ctx context.Context, event *bun.QueryEvent) {
 }
 
 func formatOperation(event *bun.QueryEvent) string {
-	operation := eventOperation(event)
+	operation := event.Operation()
 	return operationColor(operation).Sprintf(" %-16s ", operation)
 }
 
 func eventOperation(event *bun.QueryEvent) string {
-	switch event.QueryAppender.(type) {
-	case *bun.SelectQuery:
-		return "SELECT"
-	case *bun.InsertQuery:
-		return "INSERT"
-	case *bun.UpdateQuery:
-		return "UPDATE"
-	case *bun.DeleteQuery:
-		return "DELETE"
-	case *bun.CreateTableQuery:
-		return "CREATE TABLE"
-	case *bun.DropTableQuery:
-		return "DROP TABLE"
+	if event.QueryAppender != nil {
+		return event.QueryAppender.Operation()
 	}
 	return queryOperation(event.Query)
 }

extra/bunotel/otel.go

@@ -37,7 +37,7 @@ func (h *QueryHook) BeforeQuery(ctx context.Context, event *bun.QueryEvent) cont
 		return ctx
 	}
 
-	operation := eventOperation(event)
+	operation := event.Operation()
 	ctx, span := tracer.Start(ctx, operation)
 	span.SetAttributes(attribute.String("db.operation", operation))
 
@@ -109,34 +109,6 @@ func funcFileLine(pkg string) (string, string, int) {
 	return fn, file, line
 }
 
-func eventOperation(event *bun.QueryEvent) string {
-	switch event.QueryAppender.(type) {
-	case *bun.SelectQuery:
-		return "SELECT"
-	case *bun.InsertQuery:
-		return "INSERT"
-	case *bun.UpdateQuery:
-		return "UPDATE"
-	case *bun.DeleteQuery:
-		return "DELETE"
-	case *bun.CreateTableQuery:
-		return "CREATE TABLE"
-	case *bun.DropTableQuery:
-		return "DROP TABLE"
-	}
-	return queryOperation(event.Query)
-}
-
-func queryOperation(name string) string {
-	if idx := strings.IndexByte(name, ' '); idx > 0 {
-		name = name[:idx]
-	}
-	if len(name) > 16 {
-		name = name[:16]
-	}
-	return name
-}
-
 func eventQuery(event *bun.QueryEvent) string {
 	const softQueryLimit = 5000
 	const hardQueryLimit = 10000

hook.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"database/sql"
 	"reflect"
+	"strings"
 	"sync/atomic"
 	"time"
 
@@ -13,7 +14,7 @@ import (
 type QueryEvent struct {
 	DB *DB
 
-	QueryAppender schema.QueryAppender
+	QueryAppender schema.Query
 	Query         string
 	QueryArgs     []interface{}
 
@@ -24,14 +25,31 @@ type QueryEvent struct {
 	Stash map[interface{}]interface{}
 }
 
+func (e *QueryEvent) Operation() string {
+	if e.QueryAppender != nil {
+		return e.QueryAppender.Operation()
+	}
+	return queryOperation(e.Query)
+}
+
+func queryOperation(query string) string {
+	if idx := strings.IndexByte(query, ' '); idx > 0 {
+		query = query[:idx]
+	}
+	if len(query) > 16 {
+		query = query[:16]
+	}
+	return query
+}
+
 type QueryHook interface {
 	BeforeQuery(context.Context, *QueryEvent) context.Context
 	AfterQuery(context.Context, *QueryEvent)
 }
 
 func (db *DB) beforeQuery(
 	ctx context.Context,
-	queryApp schema.QueryAppender,
+	queryApp schema.Query,
 	query string,
 	queryArgs []interface{},
 ) (context.Context, *QueryEvent) {

query_base.go

@@ -427,7 +427,7 @@ func (q *baseQuery) _getFields(omitPK bool) ([]*schema.Field, error) {
 
 func (q *baseQuery) scan(
 	ctx context.Context,
-	queryApp schema.QueryAppender,
+	queryApp schema.Query,
 	query string,
 	model model,
 	hasDest bool,
@@ -459,7 +459,7 @@ func (q *baseQuery) scan(
 
 func (q *baseQuery) exec(
 	ctx context.Context,
-	queryApp schema.QueryAppender,
+	queryApp schema.Query,
 	query string,
 ) (sql.Result, error) {
 	ctx, event := q.db.beforeQuery(ctx, queryApp, query, nil)

query_column_add.go

@@ -61,6 +61,10 @@ func (q *AddColumnQuery) ColumnExpr(query string, args ...interface{}) *AddColum
 
 //------------------------------------------------------------------------------
 
+func (q *AddColumnQuery) Operation() string {
+	return "ADD COLUMN"
+}
+
 func (q *AddColumnQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
 	if q.err != nil {
 		return nil, q.err

query_column_drop.go

@@ -68,6 +68,10 @@ func (q *DropColumnQuery) ColumnExpr(query string, args ...interface{}) *DropCol
 
 //------------------------------------------------------------------------------
 
+func (q *DropColumnQuery) Operation() string {
+	return "DROP COLUMN"
+}
+
 func (q *DropColumnQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
 	if q.err != nil {
 		return nil, q.err

query_delete.go

@@ -129,6 +129,10 @@ func (q *DeleteQuery) hasReturning() bool {
 
 //------------------------------------------------------------------------------
 
+func (q *DeleteQuery) Operation() string {
+	return "DELETE"
+}
+
 func (q *DeleteQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
 	if q.err != nil {
 		return nil, q.err

query_index_create.go

@@ -142,6 +142,10 @@ func (q *CreateIndexQuery) WhereOr(query string, args ...interface{}) *CreateInd
 
 //------------------------------------------------------------------------------
 
+func (q *CreateIndexQuery) Operation() string {
+	return "CREATE INDEX"
+}
+
 func (q *CreateIndexQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
 	if q.err != nil {
 		return nil, q.err

query_index_drop.go

@@ -62,6 +62,10 @@ func (q *DropIndexQuery) Index(query string, args ...interface{}) *DropIndexQuer
 
 //------------------------------------------------------------------------------
 
+func (q *DropIndexQuery) Operation() string {
+	return "DROP INDEX"
+}
+
 func (q *DropIndexQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
 	if q.err != nil {
 		return nil, q.err

query_insert.go

@@ -142,6 +142,10 @@ func (q *InsertQuery) Replace() *InsertQuery {
 
 //------------------------------------------------------------------------------
 
+func (q *InsertQuery) Operation() string {
+	return "INSERT"
+}
+
 func (q *InsertQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
 	if q.err != nil {
 		return nil, q.err

query_select.go

@@ -359,6 +359,10 @@ func (q *SelectQuery) selectJoins(ctx context.Context, joins []relationJoin) err
 
 //------------------------------------------------------------------------------
 
+func (q *SelectQuery) Operation() string {
+	return "SELECT"
+}
+
 func (q *SelectQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
 	return q.appendQuery(fmter, b, false)
 }

query_table_create.go

@@ -85,6 +85,10 @@ func (q *CreateTableQuery) ForeignKey(query string, args ...interface{}) *Create
 	return q
 }
 
+func (q *CreateTableQuery) Operation() string {
+	return "CREATE TABLE"
+}
+
 func (q *CreateTableQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
 	if q.err != nil {
 		return nil, q.err

query_table_drop.go

@@ -68,6 +68,10 @@ func (q *DropTableQuery) Restrict() *DropTableQuery {
 
 //------------------------------------------------------------------------------
 
+func (q *DropTableQuery) Operation() string {
+	return "DROP TABLE"
+}
+
 func (q *DropTableQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
 	if q.err != nil {
 		return nil, q.err

query_table_truncate.go

@@ -64,6 +64,10 @@ func (q *TruncateTableQuery) Restrict() *TruncateTableQuery {
 
 //------------------------------------------------------------------------------
 
+func (q *TruncateTableQuery) Operation() string {
+	return "TRUNCATE TABLE"
+}
+
 func (q *TruncateTableQuery) AppendQuery(
 	fmter schema.Formatter, b []byte,
 ) (_ []byte, err error) {

query_update.go

@@ -160,6 +160,10 @@ func (q *UpdateQuery) hasReturning() bool {
 
 //------------------------------------------------------------------------------
 
+func (q *UpdateQuery) Operation() string {
+	return "SELECT"
+}
+
 func (q *UpdateQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
 	if q.err != nil {
 		return nil, q.err

query_values.go

@@ -94,6 +94,10 @@ func (q *ValuesQuery) AppendColumns(fmter schema.Formatter, b []byte) (_ []byte,
 	return nil, fmt.Errorf("bun: Values does not support %T", q.model)
 }
 
+func (q *ValuesQuery) Operation() string {
+	return "SELECT"
+}
+
 func (q *ValuesQuery) AppendQuery(fmter schema.Formatter, b []byte) (_ []byte, err error) {
 	if q.err != nil {
 		return nil, q.err

schema/sqlfmt.go

@@ -4,6 +4,11 @@ type QueryAppender interface {
 	AppendQuery(fmter Formatter, b []byte) ([]byte, error)
 }
 
+type Query interface {
+	QueryAppender
+	Operation() string
+}
+
 type ColumnsAppender interface {
 	AppendColumns(fmter Formatter, b []byte) ([]byte, error)
 }