added ScanQueryCtx (#24)

mailgun · Apr 26, 2022 · df7fa41 · df7fa41
1 parent 945d3c2
commit df7fa41
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 3 deletions.
diff --git a/cassandra.go b/cassandra.go
@@ -16,6 +16,7 @@ type Cassandra interface {
 	ExecuteQuery(string, ...interface{}) error
 	ExecuteBatch(gocql.BatchType, []string, [][]interface{}) error
 	ExecuteUnloggedBatch([]string, [][]interface{}) error
+	ScanQueryCtx(ctx context.Context, queryString string, queryParams []interface{}, outParams ...interface{}) error
 	ScanQuery(string, []interface{}, ...interface{}) error
 	ScanCASQuery(string, []interface{}, ...interface{}) (bool, error)
 	IterQuery(string, []interface{}, ...interface{}) func() (int, bool, error)
@@ -169,9 +170,9 @@ func (c *cassandra) ExecuteUnloggedBatch(queries []string, params [][]interface{
 	return c.ExecuteBatch(gocql.UnloggedBatch, queries, params)
 }
 
-// ScanQuery executes a provided SELECT query at the configured read consistency level.
-func (c *cassandra) ScanQuery(queryString string, queryParams []interface{}, outParams ...interface{}) error {
-	if err := c.Query(c.rcl, queryString, queryParams...).Scan(outParams...); err != nil {
+// ScanQueryCtx executes a provided SELECT query at the configured read consistency level.
+func (c *cassandra) ScanQueryCtx(ctx context.Context, queryString string, queryParams []interface{}, outParams ...interface{}) error {
+	if err := c.Query(c.rcl, queryString, queryParams...).WithContext(ctx).Scan(outParams...); err != nil {
 		if err == gocql.ErrNotFound {
 			return NotFound
 		}
@@ -180,6 +181,11 @@ func (c *cassandra) ScanQuery(queryString string, queryParams []interface{}, out
 	return nil
 }
 
+// ScanQuery executes a provided SELECT query at the configured read consistency level.
+func (c *cassandra) ScanQuery(queryString string, queryParams []interface{}, outParams ...interface{}) error {
+	return c.ScanQueryCtx(context.Background(), queryString, queryParams, outParams...)
+}
+
 // ScanCASQuery executes a lightweight transaction (an UPDATE or INSERT statement containing an IF clause)
 // at the configured write consistency level.
 func (c *cassandra) ScanCASQuery(queryString string, queryParams []interface{}, outParams ...interface{}) (bool, error) {

diff --git a/cassandra_test.go b/cassandra_test.go
@@ -119,6 +119,29 @@ func (s *CassandraSuite) TestScanQueryError(c *C) {
 	c.Assert(err, NotNil)
 }
 
+func (s *CassandraSuite) TestScanQueryCtxSuccess(c *C) {
+	_ = s.cassandra.ExecuteQuery("insert into test (field) values (1)")
+	var field int
+	err := s.cassandra.ScanQueryCtx(context.Background(), "select * from test", []interface{}{}, &field)
+	c.Assert(err, IsNil)
+	c.Assert(field, Equals, 1)
+}
+
+func (s *CassandraSuite) TestScanQueryCtxError(c *C) {
+	var field int
+	err := s.cassandra.ScanQueryCtx(context.Background(), "select * from test where field = 999", []interface{}{}, &field)
+	c.Assert(err, Equals, NotFound)
+}
+
+func (s *CassandraSuite) TestScanQueryCtxCanceled(c *C) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	var field int
+	err := s.cassandra.ScanQueryCtx(ctx, "select * from unknown", []interface{}{}, &field)
+	c.Assert(err, NotNil)
+}
+
 func (s *CassandraSuite) TestScanCASQuerySuccess(c *C) {
 	var field int
 	applied, err := s.cassandra.ScanCASQuery("insert into test (field) values (3) if not exists", []interface{}{}, &field)

diff --git a/testutils.go b/testutils.go
@@ -37,6 +37,10 @@ func (c *TestErrorCassandra) ExecuteUnloggedBatch(queries []string, params [][]i
 	return fmt.Errorf("Error during ExecuteUnloggedBatch")
 }
 
+func (c *TestErrorCassandra) ScanQueryCtx(_ context.Context, queryString string, queryParams []interface{}, outParams ...interface{}) error {
+	return fmt.Errorf("Error during ScanQueryCtx")
+}
+
 func (c *TestErrorCassandra) ScanQuery(queryString string, queryParams []interface{}, outParams ...interface{}) error {
 	return fmt.Errorf("Error during ScanQuery")
 }