Propagate context to gocql backend

monzo · Jun 4, 2020 · 25511c9 · 25511c9
1 parent 60ec0b6
commit 25511c9
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/gocql_backend.go b/gocql_backend.go
@@ -17,6 +17,9 @@ func (cb goCQLBackend) QueryWithOptions(opts Options, stmt Statement, scanner Sc
 	if opts.Consistency != nil {
 		qu = qu.Consistency(*opts.Consistency)
 	}
+	if opts.Context != nil {
+		qu.WithContext(opts.Context)
+	}
 
 	iter := qu.Iter()
 	if _, err := scanner.ScanIter(iter.Scanner()); err != nil {
@@ -35,6 +38,9 @@ func (cb goCQLBackend) ExecuteWithOptions(opts Options, stmt Statement) error {
 	if opts.Consistency != nil {
 		qu = qu.Consistency(*opts.Consistency)
 	}
+	if opts.Context != nil {
+		qu.WithContext(opts.Context)
+	}
 	return qu.Exec()
 }
 
@@ -55,6 +61,9 @@ func (cb goCQLBackend) ExecuteAtomicallyWithOptions(opts Options, stmts []Statem
 	if opts.Consistency != nil {
 		batch.Cons = *opts.Consistency
 	}
+	if opts.Context != nil {
+		batch.WithContext(opts.Context)
+	}
 
 	return cb.session.ExecuteBatch(batch)
 }