Skip to content

Commit

Permalink
feat: support logging sharding stQuery (#164)
Browse files Browse the repository at this point in the history
<!--
Make sure these boxes checked before submitting your pull request.

For significant changes, please open an issue to make an agreement on an
implementation design/plan first before starting it.
-->

- [x] Do only one thing
- [x] Non breaking API changes
- [x] Tested

### What did this pull request do?


support logging sharding sql


### User Case Description

We use go-gorm/sharding to split large tables and record every SQL query
in our log system for quick bug fixes. However, currently, with
go-gorm/sharding, the GORM logger only prints the original SQL and does
not print the sharded SQL.
  • Loading branch information
jinzhu authored Jul 18, 2024
2 parents 1ccc174 + 8d0fcc4 commit 59c3209
Show file tree
Hide file tree
Showing 2 changed files with 30 additions and 11 deletions.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -1 +1,2 @@
.DS_Store
.idea
40 changes: 29 additions & 11 deletions conn_pool.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package sharding
import (
"context"
"database/sql"
"time"

"gorm.io/gorm"
)
Expand All @@ -23,6 +24,10 @@ func (pool ConnPool) PrepareContext(ctx context.Context, query string) (*sql.Stm
}

func (pool ConnPool) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) {
var (
curTime = time.Now()
)

ftQuery, stQuery, table, err := pool.sharding.resolve(query, args...)
if err != nil {
return nil, err
Expand All @@ -33,32 +38,45 @@ func (pool ConnPool) ExecContext(ctx context.Context, query string, args ...any)
if table != "" {
if r, ok := pool.sharding.configs[table]; ok {
if r.DoubleWrite {
pool.ConnPool.ExecContext(ctx, ftQuery, args...)
pool.sharding.Logger.Trace(ctx, curTime, func() (sql string, rowsAffected int64) {
result, _ := pool.ConnPool.ExecContext(ctx, ftQuery, args...)
rowsAffected, _ = result.RowsAffected()
return pool.sharding.Explain(ftQuery, args...), rowsAffected
}, pool.sharding.Error)
}
}
}

return pool.ConnPool.ExecContext(ctx, stQuery, args...)
var result sql.Result
result, err = pool.ConnPool.ExecContext(ctx, stQuery, args...)
pool.sharding.Logger.Trace(ctx, curTime, func() (sql string, rowsAffected int64) {
rowsAffected, _ = result.RowsAffected()
return pool.sharding.Explain(stQuery, args...), rowsAffected
}, pool.sharding.Error)

return result, err
}

// https://github.com/go-gorm/gorm/blob/v1.21.11/callbacks/query.go#L18
func (pool ConnPool) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) {
ftQuery, stQuery, table, err := pool.sharding.resolve(query, args...)
var (
curTime = time.Now()
)

_, stQuery, _, err := pool.sharding.resolve(query, args...)
if err != nil {
return nil, err
}

pool.sharding.querys.Store("last_query", stQuery)

if table != "" {
if r, ok := pool.sharding.configs[table]; ok {
if r.DoubleWrite {
pool.ConnPool.ExecContext(ctx, ftQuery, args...)
}
}
}
var rows *sql.Rows
rows, err = pool.ConnPool.QueryContext(ctx, stQuery, args...)
pool.sharding.Logger.Trace(ctx, curTime, func() (sql string, rowsAffected int64) {
return pool.sharding.Explain(stQuery, args...), 0
}, pool.sharding.Error)

return pool.ConnPool.QueryContext(ctx, stQuery, args...)
return rows, err
}

func (pool ConnPool) QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row {
Expand Down

0 comments on commit 59c3209

Please sign in to comment.