diff --git a/.gitignore b/.gitignore index e43b0f9..3d72576 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ .DS_Store +.idea \ No newline at end of file diff --git a/conn_pool.go b/conn_pool.go index 3a6d311..cbec670 100644 --- a/conn_pool.go +++ b/conn_pool.go @@ -3,6 +3,7 @@ package sharding import ( "context" "database/sql" + "time" "gorm.io/gorm" ) @@ -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 @@ -33,16 +38,31 @@ 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 + pool.sharding.Logger.Trace(ctx, curTime, func() (sql string, rowsAffected int64) { + result, err = pool.ConnPool.ExecContext(ctx, stQuery, args...) + 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) { + var ( + curTime = time.Now() + ) + ftQuery, stQuery, table, err := pool.sharding.resolve(query, args...) if err != nil { return nil, err @@ -53,12 +73,21 @@ func (pool ConnPool) QueryContext(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) { + pool.ConnPool.QueryContext(ctx, ftQuery, args...) + return pool.sharding.Explain(ftQuery, args...), 0 + }, pool.sharding.Error) } } } - return pool.ConnPool.QueryContext(ctx, stQuery, args...) + var rows *sql.Rows + pool.sharding.Logger.Trace(ctx, curTime, func() (sql string, rowsAffected int64) { + rows, err = pool.ConnPool.QueryContext(ctx, stQuery, args...) + return pool.sharding.Explain(stQuery, args...), 0 + }, pool.sharding.Error) + + return rows, err } func (pool ConnPool) QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row {