Merge pull request #11 from cocorambo/add-context-funcs

Add database/sql context funcs

README.md

@@ -2,6 +2,8 @@
 
 Nap is a library that abstracts access to master-slave physical SQL servers topologies as a single logical database mimicking the standard `sql.DB` APIs.
 
+Nap requires Go version 1.8 or greater.
+
 ## Install
 ```shell
 $ go get github.com/tsenart/nap

db.go

@@ -1,6 +1,7 @@
 package nap
 
 import (
+	"context"
 	"database/sql"
 	"database/sql/driver"
 	"strings"
@@ -44,19 +45,35 @@ func (db *DB) Close() error {
 
 // Driver returns the physical database's underlying driver.
 func (db *DB) Driver() driver.Driver {
-	return db.pdbs[0].Driver()
+	return db.Master().Driver()
 }
 
 // Begin starts a transaction on the master. The isolation level is dependent on the driver.
 func (db *DB) Begin() (*sql.Tx, error) {
-	return db.pdbs[0].Begin()
+	return db.Master().Begin()
+}
+
+// BeginTx starts a transaction with the provided context on the master.
+//
+// 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.pdbs[0].Exec(query, args...)
+	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,
@@ -67,6 +84,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) {
@@ -84,19 +109,53 @@ 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.pdbs[db.slave(len(db.pdbs))].Query(query, args...)
+	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.
 // QueryRow uses a slave as the physical db.
 func (db *DB) QueryRow(query string, args ...interface{}) *sql.Row {
-	return db.pdbs[db.slave(len(db.pdbs))].QueryRow(query, args...)
+	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