From 979397bc5ea0e5b37559b33ef02b7e1e144d0b50 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Matczuk?= <michal@scylladb.com>
Date: Tue, 3 Aug 2021 11:29:54 +0200
Subject: [PATCH] qb: add support for USING TIMEOUT clause

In scylladb/scylla#7781 we added possibility to add timeout as part of USING spec.
This patch adds support for it by adding `Timeout` and `TimeoutNamed` functions to builders.

Fixes #194
---
 qb/batch.go       | 13 +++++++++++++
 qb/batch_test.go  | 10 ++++++++++
 qb/delete.go      | 13 +++++++++++++
 qb/delete_test.go | 11 +++++++++++
 qb/insert.go      | 13 +++++++++++++
 qb/insert_test.go | 33 +++++++++++++++++++++++++++++++++
 qb/select.go      | 18 +++++++++++++++++-
 qb/select_test.go | 12 ++++++++++++
 qb/update.go      | 13 +++++++++++++
 qb/update_test.go | 11 +++++++++++
 qb/using.go       | 26 +++++++++++++++++++++++++-
 qb/using_test.go  | 27 +++++++++++++++++++++++++++
 12 files changed, 198 insertions(+), 2 deletions(-)

diff --git a/qb/batch.go b/qb/batch.go
index 7642fb9..597c4e1 100644
--- a/qb/batch.go
+++ b/qb/batch.go
@@ -145,3 +145,16 @@ func (b *BatchBuilder) TimestampNamed(name string) *BatchBuilder {
 	b.using.TimestampNamed(name)
 	return b
 }
+
+// Timeout adds USING TIMEOUT clause to the query.
+func (b *BatchBuilder) Timeout(d time.Duration) *BatchBuilder {
+	b.using.Timeout(d)
+	return b
+}
+
+// TimeoutNamed adds a USING TIMEOUT clause to the query with a custom
+// parameter name.
+func (b *BatchBuilder) TimeoutNamed(name string) *BatchBuilder {
+	b.using.TimeoutNamed(name)
+	return b
+}
diff --git a/qb/batch_test.go b/qb/batch_test.go
index 6fe6294..027bfc7 100644
--- a/qb/batch_test.go
+++ b/qb/batch_test.go
@@ -72,6 +72,16 @@ func TestBatchBuilder(t *testing.T) {
 			S: "BEGIN BATCH USING TIMESTAMP ? APPLY BATCH ",
 			N: []string{"ts"},
 		},
+		// Add TIMEOUT
+		{
+			B: Batch().Timeout(time.Second),
+			S: "BEGIN BATCH USING TIMEOUT 1s APPLY BATCH ",
+		},
+		{
+			B: Batch().TimeoutNamed("to"),
+			S: "BEGIN BATCH USING TIMEOUT ? APPLY BATCH ",
+			N: []string{"to"},
+		},
 	}
 
 	for _, test := range table {
diff --git a/qb/delete.go b/qb/delete.go
index d0d0070..1c52432 100644
--- a/qb/delete.go
+++ b/qb/delete.go
@@ -92,6 +92,19 @@ func (b *DeleteBuilder) TimestampNamed(name string) *DeleteBuilder {
 	return b
 }
 
+// Timeout adds USING TIMEOUT clause to the query.
+func (b *DeleteBuilder) Timeout(d time.Duration) *DeleteBuilder {
+	b.using.Timeout(d)
+	return b
+}
+
+// TimeoutNamed adds a USING TIMEOUT clause to the query with a custom
+// parameter name.
+func (b *DeleteBuilder) TimeoutNamed(name string) *DeleteBuilder {
+	b.using.TimeoutNamed(name)
+	return b
+}
+
 // Where adds an expression to the WHERE clause of the query. Expressions are
 // ANDed together in the generated CQL.
 func (b *DeleteBuilder) Where(w ...Cmp) *DeleteBuilder {
diff --git a/qb/delete_test.go b/qb/delete_test.go
index caba21d..a95f59e 100644
--- a/qb/delete_test.go
+++ b/qb/delete_test.go
@@ -78,6 +78,17 @@ func TestDeleteBuilder(t *testing.T) {
 			S: "DELETE FROM cycling.cyclist_name USING TIMESTAMP ? WHERE id=? ",
 			N: []string{"ts", "expr"},
 		},
+		// Add TIMEOUT
+		{
+			B: Delete("cycling.cyclist_name").Where(w).Timeout(time.Second),
+			S: "DELETE FROM cycling.cyclist_name USING TIMEOUT 1s WHERE id=? ",
+			N: []string{"expr"},
+		},
+		{
+			B: Delete("cycling.cyclist_name").Where(w).TimeoutNamed("to"),
+			S: "DELETE FROM cycling.cyclist_name USING TIMEOUT ? WHERE id=? ",
+			N: []string{"to", "expr"},
+		},
 		// Add IF EXISTS
 		{
 			B: Delete("cycling.cyclist_name").Where(w).Existing(),
diff --git a/qb/insert.go b/qb/insert.go
index fd11bea..86b4a66 100644
--- a/qb/insert.go
+++ b/qb/insert.go
@@ -183,3 +183,16 @@ func (b *InsertBuilder) TimestampNamed(name string) *InsertBuilder {
 	b.using.TimestampNamed(name)
 	return b
 }
+
+// Timeout adds USING TIMEOUT clause to the query.
+func (b *InsertBuilder) Timeout(d time.Duration) *InsertBuilder {
+	b.using.Timeout(d)
+	return b
+}
+
+// TimeoutNamed adds a USING TIMEOUT clause to the query with a custom
+// parameter name.
+func (b *InsertBuilder) TimeoutNamed(name string) *InsertBuilder {
+	b.using.TimeoutNamed(name)
+	return b
+}
diff --git a/qb/insert_test.go b/qb/insert_test.go
index 6cc0aa7..86d4132 100644
--- a/qb/insert_test.go
+++ b/qb/insert_test.go
@@ -76,6 +76,39 @@ func TestInsertBuilder(t *testing.T) {
 			S: "INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) USING TIMESTAMP ? ",
 			N: []string{"id", "user_uuid", "firstname", "ts"},
 		},
+		// Add TIMESTAMP
+		{
+			B: Insert("cycling.cyclist_name").Columns("id", "user_uuid", "firstname").Timestamp(time.Date(2005, 05, 05, 0, 0, 0, 0, time.UTC)),
+			S: "INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) USING TIMESTAMP 1115251200000000 ",
+			N: []string{"id", "user_uuid", "firstname"},
+		},
+		{
+			B: Insert("cycling.cyclist_name").Columns("id", "user_uuid", "firstname").TimestampNamed("ts"),
+			S: "INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) USING TIMESTAMP ? ",
+			N: []string{"id", "user_uuid", "firstname", "ts"},
+		},
+		// Add TIMESTAMP
+		{
+			B: Insert("cycling.cyclist_name").Columns("id", "user_uuid", "firstname").Timestamp(time.Date(2005, 05, 05, 0, 0, 0, 0, time.UTC)),
+			S: "INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) USING TIMESTAMP 1115251200000000 ",
+			N: []string{"id", "user_uuid", "firstname"},
+		},
+		{
+			B: Insert("cycling.cyclist_name").Columns("id", "user_uuid", "firstname").TimestampNamed("ts"),
+			S: "INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) USING TIMESTAMP ? ",
+			N: []string{"id", "user_uuid", "firstname", "ts"},
+		},
+		// Add TIMEOUT
+		{
+			B: Insert("cycling.cyclist_name").Columns("id", "user_uuid", "firstname").Timeout(time.Second),
+			S: "INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) USING TIMEOUT 1s ",
+			N: []string{"id", "user_uuid", "firstname"},
+		},
+		{
+			B: Insert("cycling.cyclist_name").Columns("id", "user_uuid", "firstname").TimeoutNamed("to"),
+			S: "INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) USING TIMEOUT ? ",
+			N: []string{"id", "user_uuid", "firstname", "to"},
+		},
 		// Add TupleColumn
 		{
 			B: Insert("cycling.cyclist_name").TupleColumn("id", 2),
diff --git a/qb/select.go b/qb/select.go
index 3a9aff2..ac18e29 100644
--- a/qb/select.go
+++ b/qb/select.go
@@ -11,6 +11,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/scylladb/gocqlx/v2"
 )
@@ -37,6 +38,7 @@ type SelectBuilder struct {
 	table             string
 	columns           columns
 	distinct          columns
+	using             using
 	where             where
 	groupBy           columns
 	orderBy           columns
@@ -83,7 +85,8 @@ func (b *SelectBuilder) ToCql() (stmt string, names []string) {
 	cql.WriteString(b.table)
 	cql.WriteByte(' ')
 
-	names = b.where.writeCql(&cql)
+	names = append(names, b.using.writeCql(&cql)...)
+	names = append(names, b.where.writeCql(&cql)...)
 
 	if len(b.groupBy) > 0 {
 		cql.WriteString("GROUP BY ")
@@ -168,6 +171,19 @@ func (b *SelectBuilder) Distinct(columns ...string) *SelectBuilder {
 	return b
 }
 
+// Timeout adds USING TIMEOUT clause to the query.
+func (b *SelectBuilder) Timeout(d time.Duration) *SelectBuilder {
+	b.using.Timeout(d)
+	return b
+}
+
+// TimeoutNamed adds a USING TIMEOUT clause to the query with a custom
+// parameter name.
+func (b *SelectBuilder) TimeoutNamed(name string) *SelectBuilder {
+	b.using.TimeoutNamed(name)
+	return b
+}
+
 // Where adds an expression to the WHERE clause of the query. Expressions are
 // ANDed together in the generated CQL.
 func (b *SelectBuilder) Where(w ...Cmp) *SelectBuilder {
diff --git a/qb/select_test.go b/qb/select_test.go
index 9c7ed9a..be01991 100644
--- a/qb/select_test.go
+++ b/qb/select_test.go
@@ -6,6 +6,7 @@ package qb
 
 import (
 	"testing"
+	"time"
 
 	"github.com/google/go-cmp/cmp"
 )
@@ -77,6 +78,17 @@ func TestSelectBuilder(t *testing.T) {
 			S: "SELECT * FROM cycling.cyclist_name WHERE id=(?,?) AND firstname>(?,?) ",
 			N: []string{"id_0", "id_1", "firstname_0", "firstname_1"},
 		},
+		// Add TIMEOUT
+		{
+			B: Select("cycling.cyclist_name").Where(w, Gt("firstname")).Timeout(time.Second),
+			S: "SELECT * FROM cycling.cyclist_name USING TIMEOUT 1s WHERE id=? AND firstname>? ",
+			N: []string{"expr", "firstname"},
+		},
+		{
+			B: Select("cycling.cyclist_name").Where(w, Gt("firstname")).TimeoutNamed("to"),
+			S: "SELECT * FROM cycling.cyclist_name USING TIMEOUT ? WHERE id=? AND firstname>? ",
+			N: []string{"to", "expr", "firstname"},
+		},
 		// Add GROUP BY
 		{
 			B: Select("cycling.cyclist_name").Columns("MAX(stars) as max_stars").GroupBy("id"),
diff --git a/qb/update.go b/qb/update.go
index b9d63b7..1537a90 100644
--- a/qb/update.go
+++ b/qb/update.go
@@ -117,6 +117,19 @@ func (b *UpdateBuilder) TimestampNamed(name string) *UpdateBuilder {
 	return b
 }
 
+// Timeout adds USING TIMEOUT clause to the query.
+func (b *UpdateBuilder) Timeout(d time.Duration) *UpdateBuilder {
+	b.using.Timeout(d)
+	return b
+}
+
+// TimeoutNamed adds a USING TIMEOUT clause to the query with a custom
+// parameter name.
+func (b *UpdateBuilder) TimeoutNamed(name string) *UpdateBuilder {
+	b.using.TimeoutNamed(name)
+	return b
+}
+
 // Set adds SET clauses to the query.
 // To set a tuple column use SetTuple instead.
 func (b *UpdateBuilder) Set(columns ...string) *UpdateBuilder {
diff --git a/qb/update_test.go b/qb/update_test.go
index 48f771a..2ff61e5 100644
--- a/qb/update_test.go
+++ b/qb/update_test.go
@@ -126,6 +126,17 @@ func TestUpdateBuilder(t *testing.T) {
 			S: "UPDATE cycling.cyclist_name USING TIMESTAMP ? SET id=?,user_uuid=?,firstname=? WHERE id=? ",
 			N: []string{"ts", "id", "user_uuid", "firstname", "expr"},
 		},
+		// Add TIMEOUT
+		{
+			B: Update("cycling.cyclist_name").Set("id", "user_uuid", "firstname").Where(w).Timeout(time.Second),
+			S: "UPDATE cycling.cyclist_name USING TIMEOUT 1s SET id=?,user_uuid=?,firstname=? WHERE id=? ",
+			N: []string{"id", "user_uuid", "firstname", "expr"},
+		},
+		{
+			B: Update("cycling.cyclist_name").Set("id", "user_uuid", "firstname").Where(w).TimeoutNamed("to"),
+			S: "UPDATE cycling.cyclist_name USING TIMEOUT ? SET id=?,user_uuid=?,firstname=? WHERE id=? ",
+			N: []string{"to", "id", "user_uuid", "firstname", "expr"},
+		},
 		// Add IF EXISTS
 		{
 			B: Update("cycling.cyclist_name").Set("id", "user_uuid", "firstname").Where(w).Existing(),
diff --git a/qb/using.go b/qb/using.go
index 2f811a0..827c871 100644
--- a/qb/using.go
+++ b/qb/using.go
@@ -25,7 +25,10 @@ type using struct {
 	ttlName       string
 	timestamp     int64
 	timestampName string
-	using         bool
+	timeout       time.Duration
+	timeoutName   string
+
+	using bool
 }
 
 func (u *using) TTL(d time.Duration) *using {
@@ -55,6 +58,18 @@ func (u *using) TimestampNamed(name string) *using {
 	return u
 }
 
+func (u *using) Timeout(d time.Duration) *using {
+	u.timeout = d
+	u.timeoutName = ""
+	return u
+}
+
+func (u *using) TimeoutNamed(name string) *using {
+	u.timeout = 0
+	u.timeoutName = name
+	return u
+}
+
 func (u *using) writeCql(cql *bytes.Buffer) (names []string) {
 	u.using = false
 
@@ -79,6 +94,15 @@ func (u *using) writeCql(cql *bytes.Buffer) (names []string) {
 		names = append(names, u.timestampName)
 	}
 
+	if u.timeout != 0 {
+		u.writePreamble(cql)
+		fmt.Fprintf(cql, "TIMEOUT %s ", u.timeout)
+	} else if u.timeoutName != "" {
+		u.writePreamble(cql)
+		cql.WriteString("TIMEOUT ? ")
+		names = append(names, u.timeoutName)
+	}
+
 	return
 }
 
diff --git a/qb/using_test.go b/qb/using_test.go
index d9983af..60cbdc2 100644
--- a/qb/using_test.go
+++ b/qb/using_test.go
@@ -52,6 +52,17 @@ func TestUsing(t *testing.T) {
 			S: "USING TIMESTAMP ? ",
 			N: []string{"ts"},
 		},
+		// Timeout
+		{
+			B: new(using).Timeout(time.Second),
+			S: "USING TIMEOUT 1s ",
+		},
+		// TimeoutNamed
+		{
+			B: new(using).TimeoutNamed("to"),
+			S: "USING TIMEOUT ? ",
+			N: []string{"to"},
+		},
 		// TTL Timestamp
 		{
 			B: new(using).TTL(time.Second).Timestamp(time.Date(2005, 05, 05, 0, 0, 0, 0, time.UTC)),
@@ -75,6 +86,22 @@ func TestUsing(t *testing.T) {
 			S: "USING TTL ? AND TIMESTAMP 1115251200000000 ",
 			N: []string{"ttl"},
 		},
+		// TTL Timeout
+		{
+			B: new(using).TTL(time.Second).Timeout(time.Second),
+			S: "USING TTL 1 AND TIMEOUT 1s ",
+		},
+		// TTL TimeoutNamed
+		{
+			B: new(using).TTL(time.Second).TimeoutNamed("to"),
+			S: "USING TTL 1 AND TIMEOUT ? ",
+			N: []string{"to"},
+		},
+		// TTL Timestamp Timeout
+		{
+			B: new(using).TTL(time.Second).Timestamp(time.Date(2005, 05, 05, 0, 0, 0, 0, time.UTC)).Timeout(time.Second),
+			S: "USING TTL 1 AND TIMESTAMP 1115251200000000 AND TIMEOUT 1s ",
+		},
 		// TTL with no duration
 		{
 			B: new(using).TTL(0 * time.Second),