feat: support tx context

conn.go

@@ -120,7 +120,7 @@ func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (dd driver.Tx
 		return nil, err
 	}
 
-	return &tx{Tx: t, TxHook: c.ConnHook.(TxHook)}, nil
+	return &tx{Tx: t, TxHook: c.ConnHook.(TxHook), txContext: ctx}, nil
 }
 
 func namedValueToValue(named []driver.NamedValue) ([]driver.Value, error) {

stmt.go

@@ -55,7 +55,6 @@ func (s *stmt) QueryContext(ctx context.Context, args []driver.NamedValue) (rows
 }
 
 func (s *stmt) ExecContext(ctx context.Context, args []driver.NamedValue) (r driver.Result, err error) {
-
 	ctx, args, err = s.BeforeStmtExecContext(ctx, s.query, args, nil)
 	defer func() {
 		_, r, err = s.AfterStmtExecContext(ctx, s.query, args, r, err)

tx.go

@@ -5,13 +5,24 @@ import (
 	"database/sql/driver"
 )
 
+type txContextKey struct{}
 type tx struct {
 	driver.Tx
 	TxHook
+	txContext context.Context
+}
+
+func TxContextFromContext(ctx context.Context) context.Context {
+	value := ctx.Value(txContextKey{})
+	if value != nil {
+		return value.(context.Context)
+	}
+
+	return nil
 }
 
 func (t *tx) Commit() (err error) {
-	ctx := context.Background()
+	ctx := context.WithValue(context.Background(), txContextKey{}, t.txContext)
 	ctx, err = t.BeforeCommit(ctx, nil)
 	defer func() {
 		_, err = t.AfterCommit(ctx, err)
@@ -29,7 +40,7 @@ func (t *tx) Commit() (err error) {
 }
 
 func (t *tx) Rollback() (err error) {
-	ctx := context.Background()
+	ctx := context.WithValue(context.Background(), txContextKey{}, t.txContext)
 	ctx, err = t.BeforeRollback(ctx, nil)
 	defer func() {
 		_, err = t.AfterRollback(ctx, err)