Skip to content

Commit

Permalink
fix: allow QueryEvent to better detect operations in raw queries
Browse files Browse the repository at this point in the history
When writing long raw queries, it is common to use multiline strings that
can start with various white space characters (\n, \t, space...). E.g.

	err := handle.NewRaw(`
		SELECT
			...
		FROM
			...
		LEFT JOIN LATERAL (
			...
		) AS sub ON true
		WHERE
			...
	`).Scan(ctx)

In that case, a QueryEvent would report the Operation() as "\n\tSELECT".

This commit gives the ability to report "SELECT" instead.
  • Loading branch information
maximerety committed Nov 24, 2022
1 parent 0b38a1c commit 8e44735
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 7 deletions.
13 changes: 8 additions & 5 deletions hook.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"strings"
"sync/atomic"
"time"
"unicode"

"github.com/uptrace/bun/schema"
)
Expand Down Expand Up @@ -35,13 +36,15 @@ func (e *QueryEvent) Operation() string {
}

func queryOperation(query string) string {
if idx := strings.IndexByte(query, ' '); idx > 0 {
query = query[:idx]
queryOp := strings.TrimLeftFunc(query, unicode.IsSpace)

if idx := strings.IndexByte(queryOp, ' '); idx > 0 {
queryOp = queryOp[:idx]
}
if len(query) > 16 {
query = query[:16]
if len(queryOp) > 16 {
queryOp = queryOp[:16]
}
return query
return queryOp
}

type QueryHook interface {
Expand Down
9 changes: 7 additions & 2 deletions internal/dbtest/query_hook_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ func testQueryHook(t *testing.T, dbName string, db *bun.DB) {
hook.beforeQuery = func(
ctx context.Context, event *bun.QueryEvent,
) context.Context {
require.Equal(t, "SELECT", event.Operation())
require.Equal(
t, "SELECT * FROM (SELECT 1 AS c) AS t WHERE ('foo' = 'bar')", string(event.Query))

Expand All @@ -47,6 +48,7 @@ func testQueryHook(t *testing.T, dbName string, db *bun.DB) {
hook.beforeQuery = func(
ctx context.Context, event *bun.QueryEvent,
) context.Context {
require.Equal(t, "SELECT", event.Operation())
require.Equal(t, "SELECT 1", string(event.Query))
return ctx
}
Expand All @@ -61,12 +63,15 @@ func testQueryHook(t *testing.T, dbName string, db *bun.DB) {
hook.beforeQuery = func(
ctx context.Context, event *bun.QueryEvent,
) context.Context {
require.Equal(t, "SELECT 1", string(event.Query))
require.Equal(t, "SELECT", event.Operation())
require.Equal(t, "\n\t\t\tSELECT 1\n\t\t", string(event.Query))
return ctx
}

var num int
err := db.QueryRow("SELECT 1").Scan(&num)
err := db.QueryRow(`
SELECT 1
`).Scan(&num)
require.NoError(t, err)
require.Equal(t, 1, num)
hook.require(t)
Expand Down

0 comments on commit 8e44735

Please sign in to comment.