From d8afaf94f121df5405a9ea6dd8f0291adb1c31cb Mon Sep 17 00:00:00 2001
From: Max Melentyev <max.melentyev@reddit.com>
Date: Mon, 5 Feb 2024 08:02:17 +0000
Subject: [PATCH] Add Query.SerialConsistency()

---
 query.go      | 15 +++++++++++++++
 query_test.go |  1 +
 2 files changed, 16 insertions(+)

diff --git a/query.go b/query.go
index 7083d5b..24a13a6 100644
--- a/query.go
+++ b/query.go
@@ -63,6 +63,13 @@ type Query interface {
 	// level has been set, the default consistency level of the cluster
 	// is used.
 	Consistency(c gocql.Consistency) Query
+
+	// SerialConsistency sets the consistency level for the
+	// serial phase of conditional updates. That consistency can only be
+	// either SERIAL or LOCAL_SERIAL and if not present, it defaults to
+	// SERIAL. This option will be ignored for anything else that a
+	// conditional update/insert.
+	SerialConsistency(cons gocql.SerialConsistency) Query
 }
 
 var (
@@ -133,6 +140,10 @@ func (m QueryMock) Consistency(c gocql.Consistency) Query {
 	return m.Called(c).Get(0).(Query)
 }
 
+func (m QueryMock) SerialConsistency(c gocql.SerialConsistency) Query {
+	return m.Called(c).Get(0).(Query)
+}
+
 type query struct {
 	q *gocql.Query
 }
@@ -181,6 +192,10 @@ func (q query) Consistency(c gocql.Consistency) Query {
 	return &query{q: q.q.Consistency(c)}
 }
 
+func (q query) SerialConsistency(cons gocql.SerialConsistency) Query {
+	return &query{q: q.q.SerialConsistency(cons)}
+}
+
 func (q query) MapScanCAS(dest map[string]interface{}) (applied bool, err error) {
 	return q.q.MapScanCAS(dest)
 }
diff --git a/query_test.go b/query_test.go
index a03e196..4d079be 100644
--- a/query_test.go
+++ b/query_test.go
@@ -79,6 +79,7 @@ func TestQueryMock(t *testing.T) {
 		{"GetConsistency", nil, []interface{}{gocql.Quorum}},
 		{"SetConsistency", []interface{}{gocql.One}, nil},
 		{"Consistency", []interface{}{gocql.One}, []interface{}{m}},
+		{"SerialConsistency", []interface{}{gocql.LocalSerial}, []interface{}{m}},
 	})
 }