From 10750fefd6587a7fa829705b55adafb9f4887a0d Mon Sep 17 00:00:00 2001
From: "mykyta.oleksiienko" <mykyta.oleksiienko@gmail.com>
Date: Tue, 3 Sep 2024 16:28:34 +0300
Subject: [PATCH] Change RetryPolicy so it checks query idempotence
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

RetryPolicy doesn't check the query's idempotency,
but according to the specification queries with false idempotence shouldn't be retried.

patch by Mykyta Oleksiienko; reviewed by João Reis and Jackson Fleming for CASSGO-27
---
 CHANGELOG.md      |  2 ++
 cassandra_test.go | 50 +++++++++++++++++++++++++++++++++++++++++++++++
 conn_test.go      |  2 +-
 query_executor.go |  2 +-
 4 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 20da746a0..06666c491 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
+- Retry policy now takes into account query idempotency (CASSGO-27)
+
 ## [1.7.0] - 2024-09-23
 
 This release is the first after the donation of gocql to the Apache Software Foundation (ASF)
diff --git a/cassandra_test.go b/cassandra_test.go
index 797a7cf7f..3b0c61053 100644
--- a/cassandra_test.go
+++ b/cassandra_test.go
@@ -32,6 +32,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/stretchr/testify/require"
 	"io"
 	"math"
 	"math/big"
@@ -1066,6 +1067,55 @@ func matchSliceMap(t *testing.T, sliceMap []map[string]interface{}, testMap map[
 	}
 }
 
+type MyRetryPolicy struct {
+}
+
+func (*MyRetryPolicy) Attempt(q RetryableQuery) bool {
+	if q.Attempts() > 5 {
+		return false
+	}
+	return true
+}
+
+func (*MyRetryPolicy) GetRetryType(error) RetryType {
+	return Retry
+}
+
+func Test_RetryPolicyIdempotence(t *testing.T) {
+	session := createSession(t)
+	defer session.Close()
+
+	testCases := []struct {
+		name                  string
+		idempotency           bool
+		expectedNumberOfTries int
+	}{
+		{
+			name:                  "with retry",
+			idempotency:           true,
+			expectedNumberOfTries: 6,
+		},
+		{
+			name:                  "without retry",
+			idempotency:           false,
+			expectedNumberOfTries: 1,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			q := session.Query("INSERT INTO  gocql_test.not_existing_table(event_id, time, args) VALUES (?,?,?)", 4, UUIDFromTime(time.Now()), "test")
+
+			q.Idempotent(tc.idempotency)
+			q.RetryPolicy(&MyRetryPolicy{})
+			q.Consistency(All)
+
+			_ = q.Exec()
+			require.Equal(t, tc.expectedNumberOfTries, q.Attempts())
+		})
+	}
+}
+
 func TestSmallInt(t *testing.T) {
 	session := createSession(t)
 	defer session.Close()
diff --git a/conn_test.go b/conn_test.go
index 6cf062d95..8706683ff 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -435,7 +435,7 @@ func TestQueryMultinodeWithMetrics(t *testing.T) {
 	// 1 retry per host
 	rt := &SimpleRetryPolicy{NumRetries: 3}
 	observer := &testQueryObserver{metrics: make(map[string]*hostMetrics), verbose: false, logger: log}
-	qry := db.Query("kill").RetryPolicy(rt).Observer(observer)
+	qry := db.Query("kill").RetryPolicy(rt).Observer(observer).Idempotent(true)
 	if err := qry.Exec(); err == nil {
 		t.Fatalf("expected error")
 	}
diff --git a/query_executor.go b/query_executor.go
index fb68b07f2..03687361a 100644
--- a/query_executor.go
+++ b/query_executor.go
@@ -166,7 +166,7 @@ func (q *queryExecutor) do(ctx context.Context, qry ExecutableQuery, hostIter Ne
 
 		// Exit if the query was successful
 		// or no retry policy defined or retry attempts were reached
-		if iter.err == nil || rt == nil || !rt.Attempt(qry) {
+		if iter.err == nil || !qry.IsIdempotent() || rt == nil || !rt.Attempt(qry) {
 			return iter
 		}
 		lastErr = iter.err