Skip to content

Commit

Permalink
feat: add IQuery and QueryEvent.IQuery
Browse files Browse the repository at this point in the history
  • Loading branch information
vmihailenco committed Sep 27, 2021
1 parent 7688201 commit b762942
Show file tree
Hide file tree
Showing 5 changed files with 19 additions and 42 deletions.
18 changes: 0 additions & 18 deletions extra/bundebug/debug.go
Expand Up @@ -5,7 +5,6 @@ import (
"database/sql"
"fmt"
"reflect"
"strings"
"time"

"github.com/fatih/color"
Expand Down Expand Up @@ -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":
Expand Down
11 changes: 1 addition & 10 deletions extra/bunotel/otel.go
Expand Up @@ -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))
}

Expand Down Expand Up @@ -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 ""
}
19 changes: 14 additions & 5 deletions hook.go
Expand Up @@ -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
Expand All @@ -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)
}
Expand All @@ -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,
Expand All @@ -65,7 +73,8 @@ func (db *DB) beforeQuery(
DB: db,

Model: model,
QueryAppender: queryApp,
QueryAppender: iquery,
IQuery: iquery,
Query: query,
QueryArgs: queryArgs,

Expand Down
8 changes: 4 additions & 4 deletions query_base.go
Expand Up @@ -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 {
Expand All @@ -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 {
Expand Down
5 changes: 0 additions & 5 deletions schema/sqlfmt.go
Expand Up @@ -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)
}
Expand Down

0 comments on commit b762942

Please sign in to comment.