diff --git a/extra/bundebug/debug.go b/extra/bundebug/debug.go index 9ef270710..cc11aca93 100644 --- a/extra/bundebug/debug.go +++ b/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) } diff --git a/extra/bunotel/otel.go b/extra/bunotel/otel.go index 7c6a73718..805bfbde9 100644 --- a/extra/bunotel/otel.go +++ b/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 diff --git a/hook.go b/hook.go index 4cfa68fa6..88f8adbf2 100644 --- a/hook.go +++ b/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,6 +25,23 @@ 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) @@ -31,7 +49,7 @@ type QueryHook interface { func (db *DB) beforeQuery( ctx context.Context, - queryApp schema.QueryAppender, + queryApp schema.Query, query string, queryArgs []interface{}, ) (context.Context, *QueryEvent) { diff --git a/query_base.go b/query_base.go index 83cbd2605..4e1151dbe 100644 --- a/query_base.go +++ b/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) diff --git a/query_column_add.go b/query_column_add.go index ce2f60bf0..30d76e610 100644 --- a/query_column_add.go +++ b/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 diff --git a/query_column_drop.go b/query_column_drop.go index 5684beeb3..9fa54610d 100644 --- a/query_column_drop.go +++ b/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 diff --git a/query_delete.go b/query_delete.go index bcdf9c501..97b6dec7b 100644 --- a/query_delete.go +++ b/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 diff --git a/query_index_create.go b/query_index_create.go index de7eb7aa0..2b2b1edd7 100644 --- a/query_index_create.go +++ b/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 diff --git a/query_index_drop.go b/query_index_drop.go index c922ff04f..2a641a281 100644 --- a/query_index_drop.go +++ b/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 diff --git a/query_insert.go b/query_insert.go index d02633bf6..37daaf34a 100644 --- a/query_insert.go +++ b/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 diff --git a/query_select.go b/query_select.go index 7ff93366f..42db84613 100644 --- a/query_select.go +++ b/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) } diff --git a/query_table_create.go b/query_table_create.go index 0a4b3567c..a5fa8e231 100644 --- a/query_table_create.go +++ b/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 diff --git a/query_table_drop.go b/query_table_drop.go index 2c30171c1..8192bd548 100644 --- a/query_table_drop.go +++ b/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 diff --git a/query_table_truncate.go b/query_table_truncate.go index 1e4bef7f6..765afde92 100644 --- a/query_table_truncate.go +++ b/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) { diff --git a/query_update.go b/query_update.go index 1e032c548..ed387ecad 100644 --- a/query_update.go +++ b/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 diff --git a/query_values.go b/query_values.go index 075db4d0d..b6a78476d 100644 --- a/query_values.go +++ b/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 diff --git a/schema/sqlfmt.go b/schema/sqlfmt.go index bbdb0a01f..ba64e07ed 100644 --- a/schema/sqlfmt.go +++ b/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) }