diff --git a/query.go b/query.go index 41265f8..4413c4b 100644 --- a/query.go +++ b/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) +} diff --git a/query_test.go b/query_test.go index 45dc78a..b5a9759 100644 --- a/query_test.go +++ b/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) + } +}