Add Query.RetryPolicy

query.go

@@ -70,6 +70,9 @@ type Query interface {
 	// SERIAL. This option will be ignored for anything else that a
 	// conditional update/insert.
 	SerialConsistency(cons gocql.SerialConsistency) Query
+
+	// RetryPolicy sets the policy to use when retrying the query.
+	RetryPolicy(r gocql.RetryPolicy) Query
 }
 
 var (
@@ -140,6 +143,10 @@ func (m QueryMock) Consistency(c gocql.Consistency) Query {
 	return m.Called(c).Get(0).(Query)
 }
 
+func (m QueryMock) RetryPolicy(r gocql.RetryPolicy) Query {
+	return m.Called(r).Get(0).(Query)
+}
+
 func (m QueryMock) SerialConsistency(c gocql.SerialConsistency) Query {
 	return m.Called(c).Get(0).(Query)
 }
@@ -199,3 +206,7 @@ func (q query) SerialConsistency(cons gocql.SerialConsistency) Query {
 func (q query) MapScanCAS(dest map[string]interface{}) (applied bool, err error) {
 	return q.q.MapScanCAS(dest)
 }
+
+func (q query) RetryPolicy(r gocql.RetryPolicy) Query {
+	return &query{q: q.q.RetryPolicy(r)}
+}

query_test.go

@@ -59,6 +59,7 @@ func TestQueryMock(t *testing.T) {
 	var m, e = &QueryMock{}, fmt.Errorf("e")
 	ctx := context.Background()
 	it := &IteratorMock{}
+	rt := &gocql.SimpleRetryPolicy{NumRetries: 3}
 	testMock(t, m, &m.Mock, []struct {
 		method    string
 		arguments []interface{}
@@ -80,6 +81,7 @@ func TestQueryMock(t *testing.T) {
 		{"SetConsistency", []interface{}{gocql.One}, nil},
 		{"Consistency", []interface{}{gocql.One}, []interface{}{m}},
 		{"SerialConsistency", []interface{}{gocql.LocalSerial}, []interface{}{m}},
+		{"RetryPolicy", []interface{}{rt}, []interface{}{m}},
 	})
 }