diff --git a/db.go b/db.go index 0699256..00bda4f 100644 --- a/db.go +++ b/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,11 +109,37 @@ 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. @@ -96,7 +147,15 @@ func (db *DB) Query(query string, args ...interface{}) (*sql.Rows, error) { // 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