add missing Context variant methods

Author: joao-r-reis

session.go

@@ -786,7 +786,16 @@ func (s *Session) ExecuteBatchCAS(batch *Batch, dest ...interface{}) (applied bo
 // Further scans on the interator must also remember to include
 // the applied boolean as the first argument to *Iter.Scan
 func (b *Batch) ExecCAS(dest ...interface{}) (applied bool, iter *Iter, err error) {
-	iter = b.session.executeBatch(b, nil)
+	return b.ExecCASContext(nil, dest...)
+}
+
+// ExecCASContext executes a batch operation with the provided context and returns true if successful and
+// an iterator (to scan additional rows if more than one conditional statement)
+// was sent.
+// Further scans on the interator must also remember to include
+// the applied boolean as the first argument to *Iter.Scan
+func (b *Batch) ExecCASContext(ctx context.Context, dest ...interface{}) (applied bool, iter *Iter, err error) {
+	iter = b.session.executeBatch(b, ctx)
 	if err := iter.checkErrAndNotFound(); err != nil {
 		iter.Close()
 		return false, nil, err
@@ -814,7 +823,14 @@ func (s *Session) MapExecuteBatchCAS(batch *Batch, dest map[string]interface{})
 // however it accepts a map rather than a list of arguments for the initial
 // scan.
 func (b *Batch) MapExecCAS(dest map[string]interface{}) (applied bool, iter *Iter, err error) {
-	iter = b.session.executeBatch(b, nil)
+	return b.MapExecCASContext(nil, dest)
+}
+
+// MapExecCASContext executes a batch operation with the provided context much like ExecuteBatchCAS,
+// however it accepts a map rather than a list of arguments for the initial
+// scan.
+func (b *Batch) MapExecCASContext(ctx context.Context, dest map[string]interface{}) (applied bool, iter *Iter, err error) {
+	iter = b.session.executeBatch(b, ctx)
 	if err := iter.checkErrAndNotFound(); err != nil {
 		iter.Close()
 		return false, nil, err
@@ -1297,7 +1313,14 @@ func (q *Query) iterInternal(c *Conn, ctx context.Context) *Iter {
 // row into the map pointed at by m and discards the rest. If no rows
 // were selected, ErrNotFound is returned.
 func (q *Query) MapScan(m map[string]interface{}) error {
-	iter := q.Iter()
+	return q.MapScanContext(nil, m)
+}
+
+// MapScanContext executes the query with the provided context, copies the columns of the first selected
+// row into the map pointed at by m and discards the rest. If no rows
+// were selected, ErrNotFound is returned.
+func (q *Query) MapScanContext(ctx context.Context, m map[string]interface{}) error {
+	iter := q.IterContext(ctx)
 	if err := iter.checkErrAndNotFound(); err != nil {
 		return err
 	}
@@ -1309,7 +1332,14 @@ func (q *Query) MapScan(m map[string]interface{}) error {
 // row into the values pointed at by dest and discards the rest. If no rows
 // were selected, ErrNotFound is returned.
 func (q *Query) Scan(dest ...interface{}) error {
-	iter := q.Iter()
+	return q.ScanContext(nil, dest...)
+}
+
+// ScanContext executes the query with the provided context, copies the columns of the first selected
+// row into the values pointed at by dest and discards the rest. If no rows
+// were selected, ErrNotFound is returned.
+func (q *Query) ScanContext(ctx context.Context, dest ...interface{}) error {
+	iter := q.IterContext(ctx)
 	if err := iter.checkErrAndNotFound(); err != nil {
 		return err
 	}
@@ -1326,8 +1356,20 @@ func (q *Query) Scan(dest ...interface{}) error {
 // SELECT * FROM. So using ScanCAS with INSERT is inherently prone to
 // column mismatching. Use MapScanCAS to capture them safely.
 func (q *Query) ScanCAS(dest ...interface{}) (applied bool, err error) {
+	return q.ScanCASContext(nil, dest...)
+}
+
+// ScanCASContext executes a lightweight transaction (i.e. an UPDATE or INSERT
+// statement containing an IF clause) with the provided context. If the transaction fails because
+// the existing values did not match, the previous values will be stored
+// in dest.
+//
+// As for INSERT .. IF NOT EXISTS, previous values will be returned as if
+// SELECT * FROM. So using ScanCAS with INSERT is inherently prone to
+// column mismatching. Use MapScanCAS to capture them safely.
+func (q *Query) ScanCASContext(ctx context.Context, dest ...interface{}) (applied bool, err error) {
 	q.disableSkipMetadata = true
-	iter := q.Iter()
+	iter := q.IterContext(ctx)
 	if err := iter.checkErrAndNotFound(); err != nil {
 		return false, err
 	}
@@ -1349,8 +1391,20 @@ func (q *Query) ScanCAS(dest ...interface{}) (applied bool, err error) {
 // SELECT * FROM. So using ScanCAS with INSERT is inherently prone to
 // column mismatching. MapScanCAS is added to capture them safely.
 func (q *Query) MapScanCAS(dest map[string]interface{}) (applied bool, err error) {
+	return q.MapScanCASContext(nil, dest)
+}
+
+// MapScanCASContext executes a lightweight transaction (i.e. an UPDATE or INSERT
+// statement containing an IF clause) with the provided context. If the transaction fails because
+// the existing values did not match, the previous values will be stored
+// in dest map.
+//
+// As for INSERT .. IF NOT EXISTS, previous values will be returned as if
+// SELECT * FROM. So using ScanCAS with INSERT is inherently prone to
+// column mismatching. MapScanCAS is added to capture them safely.
+func (q *Query) MapScanCASContext(ctx context.Context, dest map[string]interface{}) (applied bool, err error) {
 	q.disableSkipMetadata = true
-	iter := q.Iter()
+	iter := q.IterContext(ctx)
 	if err := iter.checkErrAndNotFound(); err != nil {
 		return false, err
 	}