Add QueryContext methods for Jet objects

db.go

@@ -1,6 +1,7 @@
 package jet
 
 import (
+	"context"
 	"database/sql"
 )
 
@@ -69,5 +70,10 @@ func (db *Db) Begin() (*Tx, error) {
 
 // Query creates a prepared query that can be run with Rows or Run.
 func (db *Db) Query(query string, args ...interface{}) Runnable {
-	return newQuery(db, db, query, args...)
+	return db.QueryContext(context.Background(), query, args...)
+}
+
+// QueryContext creates a prepared query that can be run with Rows or Run.
+func (db *Db) QueryContext(ctx context.Context, query string, args ...interface{}) Runnable {
+	return newQuery(ctx, db, db, query, args...)
 }

query.go

@@ -1,6 +1,7 @@
 package jet
 
 import (
+	"context"
 	"database/sql"
 	"sync"
 )
@@ -12,16 +13,18 @@ type jetQuery struct {
 	id    string
 	query string
 	args  []interface{}
+	ctx   context.Context
 }
 
 // newQuery initiates a new query for the provided query object (either *sql.Tx or *sql.DB)
-func newQuery(qo queryObject, db *Db, query string, args ...interface{}) *jetQuery {
+func newQuery(ctx context.Context, qo queryObject, db *Db, query string, args ...interface{}) *jetQuery {
 	return &jetQuery{
 		qo:    qo,
 		db:    db,
 		id:    newQueryId(),
 		query: query,
 		args:  args,
+		ctx:   ctx,
 	}
 }
 
@@ -33,6 +36,10 @@ func (q *jetQuery) Rows(v interface{}) (err error) {
 	q.m.Lock()
 	defer q.m.Unlock()
 
+	if q.ctx == nil {
+		q.ctx = context.Background()
+	}
+
 	// disable lru in transactions
 	useLru := true
 	switch q.qo.(type) {
@@ -82,12 +89,12 @@ func (q *jetQuery) Rows(v interface{}) (err error) {
 
 	// If no rows need to be unpacked use Exec
 	if v == nil {
-		_, err := stmt.Exec(args...)
+		_, err := stmt.ExecContext(q.ctx, args...)
 		return err
 	}
 
 	// run query
-	rows, err := stmt.Query(args...)
+	rows, err := stmt.QueryContext(q.ctx, args...)
 	if err != nil {
 		return err
 	}

tx.go

@@ -1,6 +1,7 @@
 package jet
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 )
@@ -15,7 +16,12 @@ type Tx struct {
 
 // Query creates a prepared query that can be run with Rows or Run.
 func (tx *Tx) Query(query string, args ...interface{}) Runnable {
-	q := newQuery(tx.tx, tx.db, query, args...)
+	return tx.QueryContext(context.Background(), query, args...)
+}
+
+// QueryContext creates a prepared query that can be run with Rows or Run.
+func (tx *Tx) QueryContext(ctx context.Context, query string, args ...interface{}) Runnable {
+	q := newQuery(ctx, tx.tx, tx.db, query, args...)
 	q.id = tx.qid
 	return q
 }