Add context to ExecuteQuery (#20)

Makefile

@@ -0,0 +1,3 @@
+.PHONY: integration-test
+integration-test:
+	go test ./... -race -timeout 1m --count=1

cassandra.go

@@ -1,6 +1,7 @@
 package sandra
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -10,6 +11,7 @@ import (
 
 type Cassandra interface {
 	Query(gocql.Consistency, string, ...interface{}) *gocql.Query
+	ExecuteQueryCtx(ctx context.Context, queryString string, queryParams ...interface{}) error
 	ExecuteQuery(string, ...interface{}) error
 	ExecuteBatch(gocql.BatchType, []string, [][]interface{}) error
 	ExecuteUnloggedBatch([]string, [][]interface{}) error
@@ -131,9 +133,15 @@ func (c *cassandra) Query(consistency gocql.Consistency, queryString string, que
 	return c.session.Query(queryString, queryParams...).Consistency(consistency)
 }
 
+// ExecuteQueryCtx executes a single DML/DDL statement at the configured write consistency level.
+func (c *cassandra) ExecuteQueryCtx(ctx context.Context, queryString string, queryParams ...interface{}) error {
+	return c.Query(c.wcl, queryString, queryParams...).WithContext(ctx).Exec()
+}
+
 // ExecuteQuery executes a single DML/DDL statement at the configured write consistency level.
+// Deprecated: ExecuteQuery is deprecated. Switch to ExecuteQueryCtx.
 func (c *cassandra) ExecuteQuery(queryString string, queryParams ...interface{}) error {
-	return c.Query(c.wcl, queryString, queryParams...).Exec()
+	return c.ExecuteQueryCtx(context.Background(), queryString, queryParams...)
 }
 
 // ExecuteBatch executes a batch of DML/DDL statements at the configured write consistency level.

cassandra_test.go

@@ -1,6 +1,7 @@
 package sandra
 
 import (
+	"context"
 	"testing"
 
 	"github.com/gocql/gocql"
@@ -61,6 +62,25 @@ func (s *CassandraSuite) TestExecuteQueryError(c *C) {
 	c.Assert(err, NotNil)
 }
 
+func (s *CassandraSuite) TestExecuteQueryCtxSuccess(c *C) {
+	err := s.cassandra.ExecuteQueryCtx(context.Background(), "insert into test (field) values (1)")
+	c.Assert(err, IsNil)
+}
+
+func (s *CassandraSuite) TestExecuteQueryCtxError(c *C) {
+	err := s.cassandra.ExecuteQueryCtx(context.Background(), "drop table unknown")
+	c.Assert(err, NotNil)
+}
+
+func (s *CassandraSuite) TestExecuteQueryCtxCanceled(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	err := s.cassandra.ExecuteQueryCtx(ctx, "insert into test (field) values (1)")
+	c.Log(err)
+	c.Assert(err, NotNil)
+}
+
 func (s *CassandraSuite) TestExecuteBatchSuccess(c *C) {
 	queries := []string{
 		"insert into test (field) values (?)",

testutils.go → testutils_test.go

@@ -1,10 +1,10 @@
 package sandra
 
 import (
+	"context"
 	"fmt"
-	"time"
-
 	"sync"
+	"time"
 
 	"github.com/gocql/gocql"
 	"github.com/pkg/errors"
@@ -24,6 +24,10 @@ func (c *TestErrorCassandra) Session() *gocql.Session {
 	return nil
 }
 
+func (c *TestErrorCassandra) ExecuteQueryCtx(_ context.Context, queryString string, queryParams ...interface{}) error {
+	return fmt.Errorf("error during ExecuteQueryCtx")
+}
+
 func (c *TestErrorCassandra) ExecuteQuery(queryString string, queryParams ...interface{}) error {
 	return fmt.Errorf("Error during ExecuteQuery")
 }