wrap sql.Conn

siddontang · siddontang · commit e2189dccf2db · 2019-11-23T22:45:40.000+08:00
diff --git a/go.mod b/go.mod
@@ -5,3 +5,5 @@ require (
 	github.com/lib/pq v1.0.0
 	github.com/mattn/go-sqlite3 v1.9.0
 )
+
+go 1.13
diff --git a/sqlx.go b/sqlx.go
@@ -380,6 +380,14 @@ func (db *DB) PrepareNamed(query string) (*NamedStmt, error) {
 	return prepareNamed(db, query)
 }
 
+// Conn is a wrapper around sql.Conn with extra functionality
+type Conn struct {
+	*sql.Conn
+	driverName string
+	unsafe     bool
+	Mapper     *reflectx.Mapper
+}
+
 // Tx is an sqlx wrapper around sql.Tx with extra functionality
 type Tx struct {
 	*sql.Tx
diff --git a/sqlx_context.go b/sqlx_context.go
@@ -208,6 +208,74 @@ func (db *DB) BeginTxx(ctx context.Context, opts *sql.TxOptions) (*Tx, error) {
 	return &Tx{Tx: tx, driverName: db.driverName, unsafe: db.unsafe, Mapper: db.Mapper}, err
 }
 
+// Connx returns an *sqlx.Conn instead of an *sql.Conn.
+func (db *DB) Connx(ctx context.Context) (*Conn, error) {
+	conn, err := db.DB.Conn(ctx)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Conn{Conn: conn, driverName: db.driverName, unsafe: db.unsafe, Mapper: db.Mapper}, nil
+}
+
+// BeginTxx begins a transaction and returns an *sqlx.Tx instead of an
+// *sql.Tx.
+//
+// The provided context is used until the transaction is committed or rolled
+// back. If the context is canceled, the sql package will roll back the
+// transaction. Tx.Commit will return an error if the context provided to
+// BeginxContext is canceled.
+func (c *Conn) BeginTxx(ctx context.Context, opts *sql.TxOptions) (*Tx, error) {
+	tx, err := c.Conn.BeginTx(ctx, opts)
+	if err != nil {
+		return nil, err
+	}
+	return &Tx{Tx: tx, driverName: c.driverName, unsafe: c.unsafe, Mapper: c.Mapper}, err
+}
+
+// SelectContext using this Conn.
+// Any placeholder parameters are replaced with supplied args.
+func (c *Conn) SelectContext(ctx context.Context, dest interface{}, query string, args ...interface{}) error {
+	return SelectContext(ctx, c, dest, query, args...)
+}
+
+// GetContext using this Conn.
+// Any placeholder parameters are replaced with supplied args.
+// An error is returned if the result set is empty.
+func (c *Conn) GetContext(ctx context.Context, dest interface{}, query string, args ...interface{}) error {
+	return GetContext(ctx, c, dest, query, args...)
+}
+
+// PreparexContext returns an sqlx.Stmt instead of a sql.Stmt.
+//
+// The provided context is used for the preparation of the statement, not for
+// the execution of the statement.
+func (c *Conn) PreparexContext(ctx context.Context, query string) (*Stmt, error) {
+	return PreparexContext(ctx, c, query)
+}
+
+// QueryxContext queries the database and returns an *sqlx.Rows.
+// Any placeholder parameters are replaced with supplied args.
+func (c *Conn) QueryxContext(ctx context.Context, query string, args ...interface{}) (*Rows, error) {
+	r, err := c.Conn.QueryContext(ctx, query, args...)
+	if err != nil {
+		return nil, err
+	}
+	return &Rows{Rows: r, unsafe: c.unsafe, Mapper: c.Mapper}, err
+}
+
+// QueryRowxContext queries the database and returns an *sqlx.Row.
+// Any placeholder parameters are replaced with supplied args.
+func (c *Conn) QueryRowxContext(ctx context.Context, query string, args ...interface{}) *Row {
+	rows, err := c.Conn.QueryContext(ctx, query, args...)
+	return &Row{rows: rows, err: err, unsafe: c.unsafe, Mapper: c.Mapper}
+}
+
+// Rebind a query within a Conn's bindvar type.
+func (c *Conn) Rebind(query string) string {
+	return Rebind(BindType(c.driverName), query)
+}
+
 // StmtxContext returns a version of the prepared statement which runs within a
 // transaction. Provided stmt can be either *sql.Stmt or *sqlx.Stmt.
 func (tx *Tx) StmtxContext(ctx context.Context, stmt interface{}) *Stmt {
diff --git a/sqlx_context_test.go b/sqlx_context_test.go
@@ -1342,3 +1342,85 @@ func TestEmbeddedLiteralsContext(t *testing.T) {
 		}
 	})
 }
+
+func TestConn(t *testing.T) {
+	var schema = Schema{
+		create: `
+			CREATE TABLE tt_conn (
+				id integer,
+				value text NULL DEFAULT NULL
+			);`,
+		drop: "drop table tt_conn;",
+	}
+
+	RunWithSchemaContext(context.Background(), schema, t, func(ctx context.Context, db *DB, t *testing.T) {
+		conn, err := db.Connx(ctx)
+		defer conn.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		_, err = conn.ExecContext(ctx, conn.Rebind(`INSERT INTO tt_conn (id, value) VALUES (?, ?), (?, ?)`), 1, "a", 2, "b")
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		type s struct {
+			ID    int    `db:"id"`
+			Value string `db:"value"`
+		}
+
+		v := []s{}
+
+		err = conn.SelectContext(ctx, &v, "SELECT * FROM tt_conn ORDER BY id ASC")
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		if v[0].ID != 1 {
+			t.Errorf("Expecting ID of 1, got %d", v[0].ID)
+		}
+
+		v1 := s{}
+		err = conn.GetContext(ctx, &v1, conn.Rebind("SELECT * FROM tt_conn WHERE id=?"), 1)
+
+		if err != nil {
+			t.Fatal(err)
+		}
+		if v1.ID != 1 {
+			t.Errorf("Expecting to get back 1, but got %v\n", v1.ID)
+		}
+
+		stmt, err := conn.PreparexContext(ctx, conn.Rebind("SELECT * FROM tt_conn WHERE id=?"))
+		if err != nil {
+			t.Fatal(err)
+		}
+		v1 = s{}
+		tx, err := conn.BeginTxx(ctx, nil)
+		if err != nil {
+			t.Fatal(err)
+		}
+		tstmt := tx.Stmtx(stmt)
+		row := tstmt.QueryRowx(1)
+		err = row.StructScan(&v1)
+		if err != nil {
+			t.Error(err)
+		}
+		tx.Commit()
+		if v1.ID != 1 {
+			t.Errorf("Expecting to get back 1, but got %v\n", v1.ID)
+		}
+
+		rows, err := conn.QueryxContext(ctx, "SELECT * FROM tt_conn")
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		for rows.Next() {
+			err = rows.StructScan(&v1)
+			if err != nil {
+				t.Fatal(err)
+			}
+		}
+	})
+}

Original file line number	Diff line number	Diff line change
`@@ -5,3 +5,5 @@ require (`
`5`	`5`	`github.com/lib/pq v1.0.0`
`6`	`6`	`github.com/mattn/go-sqlite3 v1.9.0`
`7`	`7`	`)`
	`8`	`+`
	`9`	`+go 1.13`