From fa235c0a67a7601c3c34e31863bf94d58f8a8bd7 Mon Sep 17 00:00:00 2001
From: sylwiaszunejko <sylwia.szunejko@gmail.com>
Date: Tue, 10 Sep 2024 15:45:52 +0200
Subject: [PATCH] Add option to add ALLOW FILTERING to the update query

---
 qb/update.go      | 11 +++++++++++
 qb/update_test.go |  6 ++++++
 2 files changed, 17 insertions(+)

diff --git a/qb/update.go b/qb/update.go
index d3a404d..a755607 100644
--- a/qb/update.go
+++ b/qb/update.go
@@ -37,6 +37,7 @@ type UpdateBuilder struct {
 	_if         _if
 	using       using
 	exists      bool
+	allowFiltering    bool
 }
 
 // Update returns a new UpdateBuilder with the given table name.
@@ -72,6 +73,10 @@ func (b *UpdateBuilder) ToCql() (stmt string, names []string) {
 		cql.WriteString("IF EXISTS ")
 	}
 
+	if b.allowFiltering {
+		cql.WriteString("ALLOW FILTERING ")
+	}
+
 	stmt = cql.String()
 	return
 }
@@ -196,6 +201,12 @@ func (b *UpdateBuilder) AddFunc(column string, fn *Func) *UpdateBuilder {
 	return b.addValue(column, fn)
 }
 
+// AllowFiltering sets a ALLOW FILTERING clause on the query.
+func (b *UpdateBuilder) AllowFiltering() *UpdateBuilder {
+	b.allowFiltering = true
+	return b
+}
+
 func (b *UpdateBuilder) addValue(column string, value value) *UpdateBuilder {
 	b.assignments = append(b.assignments, assignment{
 		column:      column,
diff --git a/qb/update_test.go b/qb/update_test.go
index 13bb742..4afa761 100644
--- a/qb/update_test.go
+++ b/qb/update_test.go
@@ -161,6 +161,12 @@ func TestUpdateBuilder(t *testing.T) {
 			S: "UPDATE cycling.cyclist_name SET timestamp=timestamp-now() ",
 			N: nil,
 		},
+		// Add ALLOW FILTERING
+		{
+			B: Update("cycling.cyclist_name").Set("id", "user_uuid", "firstname").Where(w).AllowFiltering(),
+			S: "UPDATE cycling.cyclist_name SET id=?,user_uuid=?,firstname=? WHERE id=? ALLOW FILTERING ",
+			N: []string{"id", "user_uuid", "firstname", "expr"},
+		},
 	}
 
 	for _, test := range table {