Skip to content

Commit

Permalink
Merge pull request #367 from uptrace/fix/event-query-template
Browse files Browse the repository at this point in the history 
fix: add Event.QueryTemplate and change Event.Query to be always formatted
  • Loading branch information
@vmihailenco
vmihailenco authored Dec 19, 2021
2 parents ad65a20 + 52b1ccd commit cd7a72e
Show file tree
Hide file tree
Showing 6 changed files with 48 additions and 34 deletions.
51 changes: 30 additions & 21 deletions db.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -227,8 +227,9 @@ 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, nil)
res, err := db.DB.ExecContext(ctx, db.format(query, args))
formattedQuery := db.format(query, args)
ctx, event := db.beforeQuery(ctx, nil, query, args, formattedQuery, nil)
res, err := db.DB.ExecContext(ctx, formattedQuery)
db.afterQuery(ctx, event, res, err)
return res, err
}
Expand All @@ -240,8 +241,9 @@ 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, nil)
rows, err := db.DB.QueryContext(ctx, db.format(query, args))
formattedQuery := db.format(query, args)
ctx, event := db.beforeQuery(ctx, nil, query, args, formattedQuery, nil)
rows, err := db.DB.QueryContext(ctx, formattedQuery)
db.afterQuery(ctx, event, nil, err)
return rows, err
}
Expand All @@ -251,8 +253,9 @@ 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, nil)
row := db.DB.QueryRowContext(ctx, db.format(query, args))
formattedQuery := db.format(query, args)
ctx, event := db.beforeQuery(ctx, nil, query, args, formattedQuery, nil)
row := db.DB.QueryRowContext(ctx, formattedQuery)
db.afterQuery(ctx, event, nil, row.Err())
return row
}
Expand Down Expand Up @@ -282,24 +285,27 @@ 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, nil)
res, err := c.Conn.ExecContext(ctx, c.db.format(query, args))
formattedQuery := c.db.format(query, args)
ctx, event := c.db.beforeQuery(ctx, nil, query, args, formattedQuery, nil)
res, err := c.Conn.ExecContext(ctx, formattedQuery)
c.db.afterQuery(ctx, event, res, err)
return res, err
}

func (c Conn) QueryContext(
ctx context.Context, query string, args ...interface{},
) (*sql.Rows, error) {
ctx, event := c.db.beforeQuery(ctx, nil, query, args, nil)
rows, err := c.Conn.QueryContext(ctx, c.db.format(query, args))
formattedQuery := c.db.format(query, args)
ctx, event := c.db.beforeQuery(ctx, nil, query, args, formattedQuery, nil)
rows, err := c.Conn.QueryContext(ctx, formattedQuery)
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, nil)
row := c.Conn.QueryRowContext(ctx, c.db.format(query, args))
formattedQuery := c.db.format(query, args)
ctx, event := c.db.beforeQuery(ctx, nil, query, args, formattedQuery, nil)
row := c.Conn.QueryRowContext(ctx, formattedQuery)
c.db.afterQuery(ctx, event, nil, row.Err())
return row
}
Expand Down Expand Up @@ -413,7 +419,7 @@ func (db *DB) Begin() (Tx, error) {
}

func (db *DB) BeginTx(ctx context.Context, opts *sql.TxOptions) (Tx, error) {
ctx, event := db.beforeQuery(ctx, nil, "BEGIN", nil, nil)
ctx, event := db.beforeQuery(ctx, nil, "BEGIN", nil, "BEGIN", nil)
tx, err := db.DB.BeginTx(ctx, opts)
db.afterQuery(ctx, event, nil, err)
if err != nil {
Expand All @@ -427,14 +433,14 @@ func (db *DB) BeginTx(ctx context.Context, opts *sql.TxOptions) (Tx, error) {
}

func (tx Tx) Commit() error {
ctx, event := tx.db.beforeQuery(tx.ctx, nil, "COMMIT", nil, nil)
ctx, event := tx.db.beforeQuery(tx.ctx, nil, "COMMIT", nil, "COMMIT", nil)
err := tx.Tx.Commit()
tx.db.afterQuery(ctx, event, nil, err)
return err
}

func (tx Tx) Rollback() error {
ctx, event := tx.db.beforeQuery(tx.ctx, nil, "ROLLBACK", nil, nil)
ctx, event := tx.db.beforeQuery(tx.ctx, nil, "ROLLBACK", nil, "ROLLBACK", nil)
err := tx.Tx.Rollback()
tx.db.afterQuery(ctx, event, nil, err)
return err
Expand All @@ -447,8 +453,9 @@ 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, nil)
res, err := tx.Tx.ExecContext(ctx, tx.db.format(query, args))
formattedQuery := tx.db.format(query, args)
ctx, event := tx.db.beforeQuery(ctx, nil, query, args, formattedQuery, nil)
res, err := tx.Tx.ExecContext(ctx, formattedQuery)
tx.db.afterQuery(ctx, event, res, err)
return res, err
}
Expand All @@ -460,8 +467,9 @@ 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, nil)
rows, err := tx.Tx.QueryContext(ctx, tx.db.format(query, args))
formattedQuery := tx.db.format(query, args)
ctx, event := tx.db.beforeQuery(ctx, nil, query, args, formattedQuery, nil)
rows, err := tx.Tx.QueryContext(ctx, formattedQuery)
tx.db.afterQuery(ctx, event, nil, err)
return rows, err
}
Expand All @@ -471,8 +479,9 @@ 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, nil)
row := tx.Tx.QueryRowContext(ctx, tx.db.format(query, args))
formattedQuery := tx.db.format(query, args)
ctx, event := tx.db.beforeQuery(ctx, nil, query, args, formattedQuery, nil)
row := tx.Tx.QueryRowContext(ctx, formattedQuery)
tx.db.afterQuery(ctx, event, nil, row.Err())
return row
}
Expand Down
14 changes: 8 additions & 6 deletions extra/bunotel/otel.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -147,10 +147,10 @@ func eventQuery(event *bun.QueryEvent) string {

var query string

if len(event.Query) > softQueryLimit {
query = unformattedQuery(event)
} else {
if len(event.Query) <= softQueryLimit {
query = event.Query
} else {
query = unformattedQuery(event)
}

if len(query) > hardQueryLimit {
Expand All @@ -161,10 +161,12 @@ func eventQuery(event *bun.QueryEvent) string {
}

func unformattedQuery(event *bun.QueryEvent) string {
if b, err := event.QueryAppender.AppendQuery(schema.NewNopFormatter(), nil); err == nil {
return bytesToString(b)
if event.IQuery != nil {
if b, err := event.IQuery.AppendQuery(schema.NewNopFormatter(), nil); err == nil {
return bytesToString(b)
}
}
return string(event.Query)
return string(event.QueryTemplate)
}

func dbSystem(db *bun.DB) attribute.KeyValue {
Expand Down
7 changes: 5 additions & 2 deletions hook.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,10 @@ import (
type QueryEvent struct {
DB *DB

QueryAppender schema.QueryAppender // Deprecated: use IQuery instead
QueryAppender schema.QueryAppender // DEPRECATED: use IQuery instead
IQuery Query
Query string
QueryTemplate string
QueryArgs []interface{}
Model Model

Expand Down Expand Up @@ -51,8 +52,9 @@ type QueryHook interface {
func (db *DB) beforeQuery(
ctx context.Context,
iquery Query,
query string,
queryTemplate string,
queryArgs []interface{},
query string,
model Model,
) (context.Context, *QueryEvent) {
atomic.AddUint32(&db.stats.Queries, 1)
Expand All @@ -68,6 +70,7 @@ func (db *DB) beforeQuery(
QueryAppender: iquery,
IQuery: iquery,
Query: query,
QueryTemplate: queryTemplate,
QueryArgs: queryArgs,

StartTime: time.Now(),
Expand Down
2 changes: 1 addition & 1 deletion internal/dbtest/query_hook_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ func testQueryHook(t *testing.T, dbName string, db *bun.DB) {
require.Equal(
t, "SELECT * FROM (SELECT 1) AS t WHERE ('foo' = 'bar')", string(event.Query))

b, err := event.QueryAppender.AppendQuery(schema.NewNopFormatter(), nil)
b, err := event.IQuery.AppendQuery(schema.NewNopFormatter(), nil)
require.NoError(t, err)
require.Equal(t, "SELECT * FROM (SELECT 1) AS t WHERE (? = ?)", string(b))

Expand Down
4 changes: 2 additions & 2 deletions query_base.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -468,7 +468,7 @@ func (q *baseQuery) scan(
model Model,
hasDest bool,
) (sql.Result, error) {
ctx, event := q.db.beforeQuery(ctx, iquery, query, nil, q.model)
ctx, event := q.db.beforeQuery(ctx, iquery, query, nil, query, q.model)

rows, err := q.conn.QueryContext(ctx, query)
if err != nil {
Expand Down Expand Up @@ -498,7 +498,7 @@ func (q *baseQuery) exec(
iquery Query,
query string,
) (sql.Result, error) {
ctx, event := q.db.beforeQuery(ctx, iquery, query, nil, q.model)
ctx, event := q.db.beforeQuery(ctx, iquery, query, nil, query, q.model)
res, err := q.conn.ExecContext(ctx, query)
q.db.afterQuery(ctx, event, nil, err)
return res, err
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 @@ -773,7 +773,7 @@ func (q *SelectQuery) Count(ctx context.Context) (int, error) {
}

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

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

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

var exists bool
err = q.conn.QueryRowContext(ctx, query).Scan(&exists)
Expand Down

0 comments on commit cd7a72e

Please sign in to comment.