Add configurable consistency to query

query.go

@@ -41,6 +41,13 @@ type Query interface {
 	// Release releases a query back into a pool of queries. Released queries
 	// cannot be reused.
 	Release()
+
+	// GetConsistency returns the currently configured consistency level for
+	// the query.
+	GetConsistency() gocql.Consistency
+
+	// SetConsistency sets the consistency level for this query.
+	SetConsistency(c gocql.Consistency)
 }
 
 var (
@@ -93,6 +100,14 @@ func (m QueryMock) Release() {
 	m.Called()
 }
 
+func (m QueryMock) GetConsistency() gocql.Consistency {
+	return m.Called().Get(0).(gocql.Consistency)
+}
+
+func (m QueryMock) SetConsistency(c gocql.Consistency) {
+	m.Called(c)
+}
+
 type query struct {
 	q *gocql.Query
 }
@@ -128,3 +143,11 @@ func (q query) Scan(dest ...interface{}) error {
 func (q query) Release() {
 	q.q.Release()
 }
+
+func (q query) GetConsistency() gocql.Consistency {
+	return q.q.GetConsistency()
+}
+
+func (q query) SetConsistency(c gocql.Consistency) {
+	q.q.SetConsistency(c)
+}

query_test.go

@@ -3,6 +3,7 @@ package gockle
 import (
 	"context"
 	"fmt"
+	"github.com/gocql/gocql"
 	"reflect"
 	"testing"
 )
@@ -73,5 +74,23 @@ func TestQueryMock(t *testing.T) {
 		{"MapScan", []interface{}{map[string]interface{}(nil)}, []interface{}{nil}},
 		{"MapScan", []interface{}{map[string]interface{}{"a": 1}}, []interface{}{e}},
 		{"Release", nil, nil},
+		{"GetConsistency", nil, []interface{}{gocql.Quorum}},
+		{"SetConsistency", []interface{}{gocql.One}, nil},
 	})
 }
+
+func TestQueryConsistency(t *testing.T) {
+	var s = newSession(t)
+	defer s.Close()
+	q := s.Query("select * from gockle_test.test")
+	actual := q.GetConsistency()
+	if gocql.Quorum != actual {
+		t.Errorf("Actual consistency %s, expected %s", actual, gocql.Quorum)
+	}
+
+	q.SetConsistency(gocql.One)
+	actual = q.GetConsistency()
+	if gocql.One != actual {
+		t.Errorf("Actual consistency %s, expected %s", actual, gocql.One)
+	}
+}