Skip to content

Commit

Permalink
Merge pull request #219 from uptrace/feat/event-model
Browse files Browse the repository at this point in the history 
feat: add QueryEvent.Model
  • Loading branch information
@vmihailenco
vmihailenco committed Sep 27, 2021
2 parents 8fdc713 + b762942 commit 2abf3ba
Show file tree
Hide file tree
Showing 8 changed files with 34 additions and 53 deletions.
1 change: 1 addition & 0 deletions .github/workflows/build.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ jobs:
name: build
runs-on: ubuntu-latest
strategy:
fail-fast: false
matrix:
go-version: [1.16.x, 1.17.x]

Expand Down
18 changes: 9 additions & 9 deletions db.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -203,7 +203,7 @@ func (db *DB) Exec(query string, args ...interface{}) (sql.Result, error) {
func (db *DB) ExecContext(
ctx context.Context, query string, args ...interface{},
) (sql.Result, error) {
ctx, event := db.beforeQuery(ctx, nil, query, args)
ctx, event := db.beforeQuery(ctx, nil, query, args, nil)
res, err := db.DB.ExecContext(ctx, db.format(query, args))
db.afterQuery(ctx, event, res, err)
return res, err
Expand All @@ -216,7 +216,7 @@ func (db *DB) Query(query string, args ...interface{}) (*sql.Rows, error) {
func (db *DB) QueryContext(
ctx context.Context, query string, args ...interface{},
) (*sql.Rows, error) {
ctx, event := db.beforeQuery(ctx, nil, query, args)
ctx, event := db.beforeQuery(ctx, nil, query, args, nil)
rows, err := db.DB.QueryContext(ctx, db.format(query, args))
db.afterQuery(ctx, event, nil, err)
return rows, err
Expand All @@ -227,7 +227,7 @@ func (db *DB) QueryRow(query string, args ...interface{}) *sql.Row {
}

func (db *DB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
ctx, event := db.beforeQuery(ctx, nil, query, args)
ctx, event := db.beforeQuery(ctx, nil, query, args, nil)
row := db.DB.QueryRowContext(ctx, db.format(query, args))
db.afterQuery(ctx, event, nil, row.Err())
return row
Expand Down Expand Up @@ -258,7 +258,7 @@ func (db *DB) Conn(ctx context.Context) (Conn, error) {
func (c Conn) ExecContext(
ctx context.Context, query string, args ...interface{},
) (sql.Result, error) {
ctx, event := c.db.beforeQuery(ctx, nil, query, args)
ctx, event := c.db.beforeQuery(ctx, nil, query, args, nil)
res, err := c.Conn.ExecContext(ctx, c.db.format(query, args))
c.db.afterQuery(ctx, event, res, err)
return res, err
Expand All @@ -267,14 +267,14 @@ func (c Conn) ExecContext(
func (c Conn) QueryContext(
ctx context.Context, query string, args ...interface{},
) (*sql.Rows, error) {
ctx, event := c.db.beforeQuery(ctx, nil, query, args)
ctx, event := c.db.beforeQuery(ctx, nil, query, args, nil)
rows, err := c.Conn.QueryContext(ctx, c.db.format(query, args))
c.db.afterQuery(ctx, event, nil, err)
return rows, err
}

func (c Conn) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
ctx, event := c.db.beforeQuery(ctx, nil, query, args)
ctx, event := c.db.beforeQuery(ctx, nil, query, args, nil)
row := c.Conn.QueryRowContext(ctx, c.db.format(query, args))
c.db.afterQuery(ctx, event, nil, row.Err())
return row
Expand Down Expand Up @@ -392,7 +392,7 @@ func (tx Tx) Exec(query string, args ...interface{}) (sql.Result, error) {
func (tx Tx) ExecContext(
ctx context.Context, query string, args ...interface{},
) (sql.Result, error) {
ctx, event := tx.db.beforeQuery(ctx, nil, query, args)
ctx, event := tx.db.beforeQuery(ctx, nil, query, args, nil)
res, err := tx.Tx.ExecContext(ctx, tx.db.format(query, args))
tx.db.afterQuery(ctx, event, res, err)
return res, err
Expand All @@ -405,7 +405,7 @@ func (tx Tx) Query(query string, args ...interface{}) (*sql.Rows, error) {
func (tx Tx) QueryContext(
ctx context.Context, query string, args ...interface{},
) (*sql.Rows, error) {
ctx, event := tx.db.beforeQuery(ctx, nil, query, args)
ctx, event := tx.db.beforeQuery(ctx, nil, query, args, nil)
rows, err := tx.Tx.QueryContext(ctx, tx.db.format(query, args))
tx.db.afterQuery(ctx, event, nil, err)
return rows, err
Expand All @@ -416,7 +416,7 @@ func (tx Tx) QueryRow(query string, args ...interface{}) *sql.Row {
}

func (tx Tx) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
ctx, event := tx.db.beforeQuery(ctx, nil, query, args)
ctx, event := tx.db.beforeQuery(ctx, nil, query, args, nil)
row := tx.Tx.QueryRowContext(ctx, tx.db.format(query, args))
tx.db.afterQuery(ctx, event, nil, row.Err())
return row
Expand Down
18 changes: 0 additions & 18 deletions extra/bundebug/debug.go
View file
Original file line number Diff line number Diff line change
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
View file
Original file line number Diff line number Diff line change
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 ""
}
22 changes: 17 additions & 5 deletions hook.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,21 @@ 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

StartTime time.Time
Result sql.Result
Expand All @@ -26,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 @@ -49,9 +58,10 @@ type QueryHook interface {

func (db *DB) beforeQuery(
ctx context.Context,
queryApp schema.Query,
iquery IQuery,
query string,
queryArgs []interface{},
model Model,
) (context.Context, *QueryEvent) {
atomic.AddUint32(&db.stats.Queries, 1)

Expand All @@ -62,7 +72,9 @@ func (db *DB) beforeQuery(
event := &QueryEvent{
DB: db,

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

Expand Down
8 changes: 4 additions & 4 deletions query_base.go
View file
Original file line number Diff line number Diff line change
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)
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)
ctx, event := q.db.beforeQuery(ctx, iquery, query, nil, q.model)

res, err := q.conn.ExecContext(ctx, query)
if err != nil {
Expand Down
4 changes: 2 additions & 2 deletions query_select.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -744,7 +744,7 @@ func (q *SelectQuery) Count(ctx context.Context) (int, error) {
}

query := internal.String(queryBytes)
ctx, event := q.db.beforeQuery(ctx, qq, query, nil)
ctx, event := q.db.beforeQuery(ctx, qq, query, nil, q.model)

var num int
err = q.conn.QueryRowContext(ctx, query).Scan(&num)
Expand Down Expand Up @@ -803,7 +803,7 @@ func (q *SelectQuery) Exists(ctx context.Context) (bool, error) {
}

query := internal.String(queryBytes)
ctx, event := q.db.beforeQuery(ctx, qq, query, nil)
ctx, event := q.db.beforeQuery(ctx, qq, query, nil, q.model)

var exists bool
err = q.conn.QueryRowContext(ctx, query).Scan(&exists)
Expand Down
5 changes: 0 additions & 5 deletions schema/sqlfmt.go
View file
Original file line number Diff line number Diff line change
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 2abf3ba

Please sign in to comment.