Added DB.WithContext

db.go

@@ -60,6 +60,7 @@ type (
 
 		sqlDB      *sql.DB
 		driverName string
+		ctx        context.Context
 	}
 
 	// Errors represents a list of errors.
@@ -129,6 +130,19 @@ func (db *DB) Clone() *DB {
 	return db2
 }
 
+// WithContext returns a new instance of DB associated with the given context.
+func (db *DB) WithContext(ctx context.Context) *DB {
+	db2 := db.Clone()
+	db2.ctx = ctx
+	return db2
+}
+
+// Context returns the context associated with the DB instance.
+// It returns nil if no context is associated.
+func (db *DB) Context() context.Context {
+	return db.ctx
+}
+
 // DB returns the sql.DB instance encapsulated by dbx.DB.
 func (db *DB) DB() *sql.DB {
 	return db.sqlDB
@@ -143,7 +157,13 @@ func (db *DB) Close() error {
 
 // Begin starts a transaction.
 func (db *DB) Begin() (*Tx, error) {
-	tx, err := db.sqlDB.Begin()
+	var tx *sql.Tx
+	var err error
+	if db.ctx != nil {
+		tx, err = db.sqlDB.BeginTx(db.ctx, nil)
+	} else {
+		tx, err = db.sqlDB.Begin()
+	}
 	if err != nil {
 		return nil, err
 	}

db_test.go

@@ -5,6 +5,7 @@
 package dbx
 
 import (
+	"context"
 	"database/sql"
 	"errors"
 	"io/ioutil"
@@ -36,7 +37,13 @@ func TestDB_Open(t *testing.T) {
 		assert.NotNil(t, db.sqlDB)
 		assert.NotNil(t, db.FieldMapper)
 		db2 := db.Clone()
+		assert.NotEqual(t, db, db2)
 		assert.Equal(t, db.driverName, db2.driverName)
+		ctx := context.Background()
+		db3 := db.WithContext(ctx)
+		assert.Equal(t, ctx, db3.ctx)
+		assert.Equal(t, ctx, db3.Context())
+		assert.NotEqual(t, db, db3)
 	}
 
 	_, err = Open("xyz", TestDSN)
@@ -179,6 +186,13 @@ func TestDB_Begin(t *testing.T) {
 			},
 			desc: "Wrap",
 		},
+		{
+			makeTx: func(db *DB) *Tx {
+				tx, _ := db.BeginTx(context.Background(), nil)
+				return tx
+			},
+			desc: "BeginTx",
+		},
 	}
 
 	db := getPreparedDB()

model_query.go

@@ -32,6 +32,7 @@ var (
 func NewModelQuery(model interface{}, fieldMapFunc FieldMapFunc, db *DB, builder Builder) *ModelQuery {
 	q := &ModelQuery{
 		db:      db,
+		ctx:     db.ctx,
 		builder: builder,
 		model:   newStructValue(model, fieldMapFunc),
 	}

query.go

@@ -68,6 +68,7 @@ func NewQuery(db *DB, executor Executor, sql string) *Query {
 		rawSQL:       rawSQL,
 		placeholders: placeholders,
 		params:       Params{},
+		ctx:          db.ctx,
 		FieldMapper:  db.FieldMapper,
 		LogFunc:      db.LogFunc,
 		PerfFunc:     db.PerfFunc,

select.go

@@ -59,6 +59,7 @@ func NewSelectQuery(builder Builder, db *DB) *SelectQuery {
 		union:       []UnionInfo{},
 		limit:       -1,
 		params:      Params{},
+		ctx:         db.ctx,
 		FieldMapper: db.FieldMapper,
 	}
 }