Adding capabilities to change the error returned by the driver

Author: ext_christleo

conn.go

@@ -17,13 +17,12 @@ type Conn struct {
 // It will trigger PrePing, Ping, PostPing hooks.
 //
 // If the original connection does not satisfy "database/sql/driver".Pinger, it does nothing.
-func (conn *Conn) Ping(c context.Context) error {
-	var err error
+func (conn *Conn) Ping(c context.Context) (err error) {
 	var ctx interface{}
 	hooks := conn.Proxy.getHooks(c)
 
 	if hooks != nil {
-		defer func() { hooks.postPing(c, ctx, conn, err) }()
+		defer func() { err = hooks.postPing(c, ctx, conn, err) }()
 		if ctx, err = hooks.prePing(c, conn); err != nil {
 			return err
 		}
@@ -49,31 +48,30 @@ func (conn *Conn) Prepare(query string) (driver.Stmt, error) {
 }
 
 // PrepareContext returns a prepared statement which is wrapped by Stmt.
-func (conn *Conn) PrepareContext(c context.Context, query string) (driver.Stmt, error) {
+func (conn *Conn) PrepareContext(c context.Context, query string) (stmt driver.Stmt, err error) {
 	var ctx interface{}
-	var stmt = &Stmt{
+	var stmtAux = &Stmt{
 		QueryString: query,
 		Proxy:       conn.Proxy,
 		Conn:        conn,
 	}
-	var err error
 	hooks := conn.Proxy.getHooks(c)
 	if hooks != nil {
-		defer func() { hooks.postPrepare(c, ctx, stmt, err) }()
-		if ctx, err = hooks.prePrepare(c, stmt); err != nil {
+		defer func() { err = hooks.postPrepare(c, ctx, stmtAux, err) }()
+		if ctx, err = hooks.prePrepare(c, stmtAux); err != nil {
 			return nil, err
 		}
 	}
 
 	if connCtx, ok := conn.Conn.(driver.ConnPrepareContext); ok {
-		stmt.Stmt, err = connCtx.PrepareContext(c, stmt.QueryString)
+		stmtAux.Stmt, err = connCtx.PrepareContext(c, stmtAux.QueryString)
 	} else {
-		stmt.Stmt, err = conn.Conn.Prepare(stmt.QueryString)
+		stmtAux.Stmt, err = conn.Conn.Prepare(stmtAux.QueryString)
 		if err == nil {
 			select {
 			default:
 			case <-c.Done():
-				stmt.Stmt.Close()
+				stmtAux.Stmt.Close()
 				return nil, c.Err()
 			}
 		}
@@ -83,21 +81,20 @@ func (conn *Conn) PrepareContext(c context.Context, query string) (driver.Stmt,
 	}
 
 	if hooks != nil {
-		if err = hooks.prepare(c, ctx, stmt); err != nil {
+		if err = hooks.prepare(c, ctx, stmtAux); err != nil {
 			return nil, err
 		}
 	}
-	return stmt, nil
+	return stmtAux, nil
 }
 
 // Close calls the original Close method.
-func (conn *Conn) Close() error {
+func (conn *Conn) Close() (err error) {
 	ctx := context.Background()
-	var err error
 	var myctx interface{}
 
 	if hooks := conn.Proxy.hooks; hooks != nil {
-		defer func() { hooks.postClose(ctx, myctx, conn, err) }()
+		defer func() { err = hooks.postClose(ctx, myctx, conn, err) }()
 		if myctx, err = hooks.preClose(ctx, conn); err != nil {
 			return err
 		}
@@ -123,14 +120,12 @@ func (conn *Conn) Begin() (driver.Tx, error) {
 
 // BeginTx starts and returns a new transaction which is wrapped by Tx.
 // It will trigger PreBegin, Begin, PostBegin hooks.
-func (conn *Conn) BeginTx(c context.Context, opts driver.TxOptions) (driver.Tx, error) {
+func (conn *Conn) BeginTx(c context.Context, opts driver.TxOptions) (tx driver.Tx, err error) {
 	// set the hooks.
-	var err error
 	var ctx interface{}
-	var tx driver.Tx
 	hooks := conn.Proxy.getHooks(c)
 	if hooks != nil {
-		defer func() { hooks.postBegin(c, ctx, conn, err) }()
+		defer func() { err = hooks.postBegin(c, ctx, conn, err) }()
 		if ctx, err = hooks.preBegin(c, conn); err != nil {
 			return nil, err
 		}
@@ -255,7 +250,7 @@ func (conn *Conn) Query(query string, args []driver.Value) (driver.Rows, error)
 // It wil trigger PreQuery, Query, PostQuery hooks.
 //
 // If the original connection does not satisfy "database/sql/driver".QueryerContext nor "database/sql/driver".Queryer, it return ErrSkip error.
-func (conn *Conn) QueryContext(c context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
+func (conn *Conn) QueryContext(c context.Context, query string, args []driver.NamedValue) (rows driver.Rows, err error) {
 	queryer, qok := conn.Conn.(driver.Queryer)
 	queryerCtx, qCtxOk := conn.Conn.(driver.QueryerContext)
 	if !qok && !qCtxOk {
@@ -268,11 +263,9 @@ func (conn *Conn) QueryContext(c context.Context, query string, args []driver.Na
 		Conn:        conn,
 	}
 	var ctx interface{}
-	var err error
-	var rows driver.Rows
 	hooks := conn.Proxy.getHooks(c)
 	if hooks != nil {
-		defer func() { hooks.postQuery(c, ctx, stmt, args, rows, err) }()
+		defer func() { err = hooks.postQuery(c, ctx, stmt, args, rows, err) }()
 		if ctx, err = hooks.preQuery(c, stmt, args); err != nil {
 			return nil, err
 		}
@@ -340,13 +333,12 @@ type sessionResetter interface {
 }
 
 // ResetSession resets the state of Conn.
-func (conn *Conn) ResetSession(ctx context.Context) error {
-	var err error
+func (conn *Conn) ResetSession(ctx context.Context) (err error) {
 	var myctx interface{}
 	hooks := conn.Proxy.getHooks(ctx)
 
 	if hooks != nil {
-		defer func() { hooks.postResetSession(ctx, myctx, conn, err) }()
+		defer func() { err = hooks.postResetSession(ctx, myctx, conn, err) }()
 		if myctx, err = hooks.preResetSession(ctx, conn); err != nil {
 			return err
 		}