Added Context methods

db.go

@@ -1,6 +1,7 @@
 package nap
 
 import (
+	"context"
 	"database/sql"
 	"database/sql/driver"
 	"errors"
@@ -62,13 +63,30 @@ func (db *DB) Begin() (*sql.Tx, error) {
 	return db.Master().Begin()
 }
 
+// BeginTx starts a transaction.
+// 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 BeginTx is canceled.
+// The provided TxOptions is optional and may be nil if defaults should be used.
+// If a non-default isolation level is used that the driver doesn't support, an error will be returned.
+func (db *DB) BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error) {
+	return db.Master().BeginTx(ctx, opts)
+}
+
 // Exec executes a query without returning any rows.
 // The args are for any placeholder parameters in the query.
 // Exec uses the master as the underlying physical db.
 func (db *DB) Exec(query string, args ...interface{}) (sql.Result, error) {
 	return db.Master().Exec(query, args...)
 }
 
+// ExecContext executes a query without returning any rows.
+// The args are for any placeholder parameters in the query.
+// Exec uses the master as the underlying physical db.
+func (db *DB) ExecContext(ctx context.Context, query string, args ...interface{}) (sql.Result, error) {
+	return db.Master().ExecContext(ctx, query, args...)
+}
+
 // Ping verifies if a connection to each physical database is still alive,
 // establishing a connection if necessary.
 func (db *DB) Ping() error {
@@ -77,6 +95,14 @@ func (db *DB) Ping() error {
 	})
 }
 
+// PingContext verifies if a connection to each physical database is still
+// alive, establishing a connection if necessary.
+func (db *DB) PingContext(ctx context.Context) error {
+	return scatter(len(db.pdbs), func(i int) error {
+		return db.pdbs[i].PingContext(ctx)
+	})
+}
+
 // Prepare creates a prepared statement for later queries or executions
 // on each physical database, concurrently.
 func (db *DB) Prepare(query string) (*Stmt, error) {
@@ -94,13 +120,39 @@ func (db *DB) Prepare(query string) (*Stmt, error) {
 	return &Stmt{db: db, stmts: stmts}, nil
 }
 
+// PrepareContext creates a prepared statement for later queries or executions
+// on each physical database, concurrently.
+// The provided context is used for the preparation of the statement, not for
+// the execution of the statement.
+func (db *DB) PrepareContext(ctx context.Context, query string) (*Stmt, error) {
+	stmts := make([]*sql.Stmt, len(db.pdbs))
+
+	err := scatter(len(db.pdbs), func(i int) (err error) {
+		stmts[i], err = db.pdbs[i].PrepareContext(ctx, query)
+		return err
+	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	return &Stmt{db: db, stmts: stmts}, nil
+}
+
 // Query executes a query that returns rows, typically a SELECT.
 // The args are for any placeholder parameters in the query.
 // Query uses a slave as the physical db.
 func (db *DB) Query(query string, args ...interface{}) (*sql.Rows, error) {
 	return db.Slave().Query(query, args...)
 }
 
+// QueryContext executes a query that returns rows, typically a SELECT.
+// The args are for any placeholder parameters in the query.
+// QueryContext uses a slave as the physical db.
+func (db *DB) QueryContext(ctx context.Context, query string, args ...interface{}) (*sql.Rows, error) {
+	return db.Slave().QueryContext(ctx, query, args...)
+}
+
 // QueryRow executes a query that is expected to return at most one row.
 // QueryRow always return a non-nil value.
 // Errors are deferred until Row's Scan method is called.
@@ -109,6 +161,14 @@ func (db *DB) QueryRow(query string, args ...interface{}) *sql.Row {
 	return db.Slave().QueryRow(query, args...)
 }
 
+// QueryRowContext executes a query that is expected to return at most one row.
+// QueryRowContext always return a non-nil value.
+// Errors are deferred until Row's Scan method is called.
+// QueryRowContext uses a slave as the physical db.
+func (db *DB) QueryRowContext(ctx context.Context, query string, args ...interface{}) *sql.Row {
+	return db.Slave().QueryRowContext(ctx, query, args...)
+}
+
 // SetMaxIdleConns sets the maximum number of connections in the idle
 // connection pool for each underlying physical db.
 // If MaxOpenConns is greater than 0 but less than the new MaxIdleConns then the

stmt.go

@@ -1,6 +1,7 @@
 package nap
 
 import (
+	"context"
 	"database/sql"
 )
 
@@ -26,13 +27,27 @@ func (s *Stmt) Exec(args ...interface{}) (sql.Result, error) {
 	return s.Master().Exec(args...)
 }
 
+// ExecContext executes a prepared statement with the given arguments
+// and returns a Result summarizing the effect of the statement.
+// Exec uses the master as the underlying physical db.
+func (s *Stmt) ExecContext(ctx context.Context, args ...interface{}) (sql.Result, error) {
+	return s.Master().ExecContext(ctx, args...)
+}
+
 // Query executes a prepared query statement with the given
 // arguments and returns the query results as a *sql.Rows.
 // Query uses a slave as the underlying physical db.
 func (s *Stmt) Query(args ...interface{}) (*sql.Rows, error) {
 	return s.Slave().Query(args...)
 }
 
+// QueryContext executes a query that returns rows, typically a SELECT.
+// The args are for any placeholder parameters in the query.
+// QueryContext uses a slave as the physical db.
+func (s *Stmt) QueryContext(ctx context.Context, args ...interface{}) (*sql.Rows, error) {
+	return s.Slave().QueryContext(ctx, args...)
+}
+
 // QueryRow executes a prepared query statement with the given arguments.
 // If an error occurs during the execution of the statement, that error
 // will be returned by a call to Scan on the returned *Row, which is always non-nil.
@@ -43,6 +58,14 @@ func (s *Stmt) QueryRow(args ...interface{}) *sql.Row {
 	return s.Slave().QueryRow(args...)
 }
 
+// QueryRowContext executes a query that is expected to return at most one row.
+// QueryRowContext always return a non-nil value.
+// Errors are deferred until Row's Scan method is called.
+// QueryRowContext uses a slave as the physical db.
+func (s *Stmt) QueryRowContext(ctx context.Context, args ...interface{}) *sql.Row {
+	return s.Slave().QueryRowContext(ctx, args...)
+}
+
 // Master returns the master stmt physical database
 func (s *Stmt) Master() *sql.Stmt {
 	return s.stmts[0]