Keep the unsafe flag in Tx.Stmtx and Tx.StmtxContext (upstream #979).

sqlx.go

@@ -473,17 +473,20 @@ func (tx *Tx) Preparex(query string) (*Stmt, error) {
 // stmt can be either *sql.Stmt or *sqlx.Stmt.
 func (tx *Tx) Stmtx(stmt interface{}) *Stmt {
 	var s *sql.Stmt
+	unsafe := tx.unsafe
 	switch v := stmt.(type) {
 	case Stmt:
 		s = v.Stmt
+		unsafe = unsafe || v.unsafe
 	case *Stmt:
 		s = v.Stmt
+		unsafe = unsafe || v.unsafe
 	case *sql.Stmt:
 		s = v
 	default:
 		panic(fmt.Sprintf("non-statement type %v passed to Stmtx", reflect.ValueOf(stmt).Type()))
 	}
-	return &Stmt{Stmt: tx.Stmt(s), Mapper: tx.Mapper}
+	return &Stmt{Stmt: tx.Stmt(s), unsafe: unsafe, Mapper: tx.Mapper}
 }
 
 // NamedStmt returns a version of the prepared statement which runs within a transaction.

sqlx_context.go

@@ -281,17 +281,20 @@ func (c *Conn) Rebind(query string) string {
 // transaction. Provided stmt can be either *sql.Stmt or *sqlx.Stmt.
 func (tx *Tx) StmtxContext(ctx context.Context, stmt interface{}) *Stmt {
 	var s *sql.Stmt
+	unsafe := tx.unsafe
 	switch v := stmt.(type) {
 	case Stmt:
 		s = v.Stmt
+		unsafe = unsafe || v.unsafe
 	case *Stmt:
 		s = v.Stmt
+		unsafe = unsafe || v.unsafe
 	case *sql.Stmt:
 		s = v
 	default:
 		panic(fmt.Sprintf("non-statement type %v passed to Stmtx", reflect.ValueOf(stmt).Type()))
 	}
-	return &Stmt{Stmt: tx.StmtContext(ctx, s), Mapper: tx.Mapper}
+	return &Stmt{Stmt: tx.StmtContext(ctx, s), unsafe: unsafe, Mapper: tx.Mapper}
 }
 
 // NamedStmtContext returns a version of the prepared statement which runs

tx_stmtx_test.go

@@ -0,0 +1,64 @@
+package sqlx
+
+import (
+	"context"
+	"testing"
+)
+
+// Tx.Stmtx and Tx.StmtxContext keep the transaction's unsafe flag, and keep
+// the wrapped statement's own unsafe flag (upstream #979).
+func TestTxStmtxKeepsUnsafe(t *testing.T) {
+	RunWithSchema(defaultSchema, t, func(db *DB, t *testing.T, now string) {
+		loadDefaultFixture(db, t)
+		ctx := context.Background()
+		stmt, err := db.Preparex(db.Rebind("SELECT * FROM person WHERE first_name=?"))
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer stmt.Close()
+
+		var partial struct {
+			FirstName string `db:"first_name"`
+		}
+
+		tx, err := db.Unsafe().Beginx()
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer tx.Rollback()
+		if !tx.unsafe {
+			t.Fatal("precondition: tx from db.Unsafe() should be unsafe")
+		}
+		ts := tx.Stmtx(stmt)
+		if !ts.unsafe {
+			t.Error("Stmtx: unsafe Tx returned a safe Stmt")
+		}
+		if err := ts.Get(&partial, "Jason"); err != nil {
+			t.Errorf("Stmtx: Get into partial struct: %v", err)
+		}
+		tsc := tx.StmtxContext(ctx, stmt)
+		if !tsc.unsafe {
+			t.Error("StmtxContext: unsafe Tx returned a safe Stmt")
+		}
+		if err := tsc.GetContext(ctx, &partial, "Jason"); err != nil {
+			t.Errorf("StmtxContext: Get into partial struct: %v", err)
+		}
+
+		// An unsafe statement wrapped in a safe transaction stays unsafe.
+		safeTx, err := db.Beginx()
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer safeTx.Rollback()
+		us := safeTx.Stmtx(stmt.Unsafe())
+		if !us.unsafe {
+			t.Error("Stmtx: unsafe Stmt wrapped in a safe Tx became safe")
+		}
+		if err := us.Get(&partial, "Jason"); err != nil {
+			t.Errorf("Stmtx: Get through unsafe Stmt in safe Tx: %v", err)
+		}
+		if ss := safeTx.StmtxContext(ctx, stmt); ss.unsafe {
+			t.Error("StmtxContext: safe Stmt in safe Tx became unsafe")
+		}
+	})
+}