feat: add IQuery and QueryEvent.IQuery

extra/bundebug/debug.go

@@ -5,7 +5,6 @@ import (
 	"database/sql"
 	"fmt"
 	"reflect"
-	"strings"
 	"time"
 
 	"github.com/fatih/color"
@@ -76,23 +75,6 @@ func formatOperation(event *bun.QueryEvent) string {
 	return operationColor(operation).Sprintf(" %-16s ", operation)
 }
 
-func eventOperation(event *bun.QueryEvent) string {
-	if event.QueryAppender != nil {
-		return event.QueryAppender.Operation()
-	}
-	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 operationColor(operation string) *color.Color {
 	switch operation {
 	case "SELECT":

extra/bunotel/otel.go

@@ -57,7 +57,7 @@ func (h *QueryHook) AfterQuery(ctx context.Context, event *bun.QueryEvent) {
 	dbOperation := attribute.String("db.operation", operation)
 
 	labels := []attribute.KeyValue{dbOperation}
-	if tableName := tableName(event.QueryAppender); tableName != "" {
+	if tableName := event.IQuery.GetTableName(); tableName != "" {
 		labels = append(labels, attribute.String("db.table", tableName))
 	}
 
@@ -168,12 +168,3 @@ func dbSystem(db *bun.DB) string {
 		return ""
 	}
 }
-
-func tableName(query schema.Query) string {
-	if v, ok := query.(interface {
-		GetTableName() string
-	}); ok {
-		return v.GetTableName()
-	}
-	return ""
-}

hook.go

@@ -11,10 +11,18 @@ import (
 	"github.com/uptrace/bun/schema"
 )
 
+type IQuery interface {
+	schema.QueryAppender
+	Operation() string
+	GetModel() Model
+	GetTableName() string
+}
+
 type QueryEvent struct {
 	DB *DB
 
-	QueryAppender schema.Query
+	QueryAppender schema.QueryAppender // Deprecated: use IQuery instead
+	IQuery        IQuery
 	Query         string
 	QueryArgs     []interface{}
 	Model         Model
@@ -27,8 +35,8 @@ type QueryEvent struct {
 }
 
 func (e *QueryEvent) Operation() string {
-	if e.QueryAppender != nil {
-		return e.QueryAppender.Operation()
+	if e.IQuery != nil {
+		return e.IQuery.Operation()
 	}
 	return queryOperation(e.Query)
 }
@@ -50,7 +58,7 @@ type QueryHook interface {
 
 func (db *DB) beforeQuery(
 	ctx context.Context,
-	queryApp schema.Query,
+	iquery IQuery,
 	query string,
 	queryArgs []interface{},
 	model Model,
@@ -65,7 +73,8 @@ func (db *DB) beforeQuery(
 		DB: db,
 
 		Model:         model,
-		QueryAppender: queryApp,
+		QueryAppender: iquery,
+		IQuery:        iquery,
 		Query:         query,
 		QueryArgs:     queryArgs,
 

query_base.go

@@ -456,12 +456,12 @@ func (q *baseQuery) _getFields(omitPK bool) ([]*schema.Field, error) {
 
 func (q *baseQuery) scan(
 	ctx context.Context,
-	queryApp schema.Query,
+	iquery IQuery,
 	query string,
 	model Model,
 	hasDest bool,
 ) (sql.Result, error) {
-	ctx, event := q.db.beforeQuery(ctx, queryApp, query, nil, q.model)
+	ctx, event := q.db.beforeQuery(ctx, iquery, query, nil, q.model)
 
 	rows, err := q.conn.QueryContext(ctx, query)
 	if err != nil {
@@ -488,10 +488,10 @@ func (q *baseQuery) scan(
 
 func (q *baseQuery) exec(
 	ctx context.Context,
-	queryApp schema.Query,
+	iquery IQuery,
 	query string,
 ) (sql.Result, error) {
-	ctx, event := q.db.beforeQuery(ctx, queryApp, query, nil, q.model)
+	ctx, event := q.db.beforeQuery(ctx, iquery, query, nil, q.model)
 
 	res, err := q.conn.ExecContext(ctx, query)
 	if err != nil {

schema/sqlfmt.go

@@ -4,11 +4,6 @@ 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)
 }