qb: add named limit and per partition limit clauses

qb/limit.go

@@ -0,0 +1,45 @@
+// Copyright (C) 2017 ScyllaDB
+// Use of this source code is governed by a ALv2-style
+// license that can be found in the LICENSE file.
+
+package qb
+
+import (
+	"bytes"
+	"strconv"
+)
+
+type limit struct {
+	value        value
+	perPartition bool
+}
+
+func limitLit(l uint, perPartition bool) limit {
+	val := strconv.FormatUint(uint64(l), 10)
+	return limit{
+		value:        lit(val),
+		perPartition: perPartition,
+	}
+}
+
+func limitNamed(name string, perPartition bool) limit {
+	return limit{
+		value:        param(name),
+		perPartition: perPartition,
+	}
+}
+
+func (l limit) writeCql(cql *bytes.Buffer) (names []string) {
+	if l.value == nil {
+		return nil
+	}
+
+	if l.perPartition {
+		cql.WriteString("PER PARTITION ")
+	}
+	cql.WriteString("LIMIT ")
+
+	names = l.value.writeCql(cql)
+	cql.WriteByte(' ')
+	return
+}

qb/select.go

@@ -10,7 +10,6 @@ package qb
 import (
 	"bytes"
 	"context"
-	"fmt"
 	"time"
 
 	"github.com/scylladb/gocqlx/v2"
@@ -42,8 +41,8 @@ type SelectBuilder struct {
 	where             where
 	groupBy           columns
 	orderBy           columns
-	limit             uint
-	limitPerPartition uint
+	limit             limit
+	limitPerPartition limit
 	allowFiltering    bool
 	bypassCache       bool
 	json              bool
@@ -100,17 +99,8 @@ func (b *SelectBuilder) ToCql() (stmt string, names []string) {
 		cql.WriteByte(' ')
 	}
 
-	if b.limit != 0 {
-		cql.WriteString("LIMIT ")
-		cql.WriteString(fmt.Sprint(b.limit))
-		cql.WriteByte(' ')
-	}
-
-	if b.limitPerPartition != 0 {
-		cql.WriteString("PER PARTITION LIMIT ")
-		cql.WriteString(fmt.Sprint(b.limitPerPartition))
-		cql.WriteByte(' ')
-	}
+	names = append(names, b.limitPerPartition.writeCql(&cql)...)
+	names = append(names, b.limit.writeCql(&cql)...)
 
 	if b.allowFiltering {
 		cql.WriteString("ALLOW FILTERING ")
@@ -214,13 +204,25 @@ func (b *SelectBuilder) OrderBy(column string, o Order) *SelectBuilder {
 
 // Limit sets a LIMIT clause on the query.
 func (b *SelectBuilder) Limit(limit uint) *SelectBuilder {
-	b.limit = limit
+	b.limit = limitLit(limit, false)
+	return b
+}
+
+// LimitNamed produces LIMIT ? clause with a custom parameter name.
+func (b *SelectBuilder) LimitNamed(name string) *SelectBuilder {
+	b.limit = limitNamed(name, false)
 	return b
 }
 
 // LimitPerPartition sets a PER PARTITION LIMIT clause on the query.
 func (b *SelectBuilder) LimitPerPartition(limit uint) *SelectBuilder {
-	b.limitPerPartition = limit
+	b.limitPerPartition = limitLit(limit, true)
+	return b
+}
+
+// LimitPerPartitionNamed produces PER PARTITION LIMIT ? clause with a custom parameter name.
+func (b *SelectBuilder) LimitPerPartitionNamed(name string) *SelectBuilder {
+	b.limitPerPartition = limitNamed(name, true)
 	return b
 }
 

qb/select_test.go

@@ -128,12 +128,36 @@ func TestSelectBuilder(t *testing.T) {
 			S: "SELECT * FROM cycling.cyclist_name WHERE id=? LIMIT 10 ",
 			N: []string{"expr"},
 		},
+		// Add named LIMIT
+		{
+			B: Select("cycling.cyclist_name").Where(w).LimitNamed("limit"),
+			S: "SELECT * FROM cycling.cyclist_name WHERE id=? LIMIT ? ",
+			N: []string{"expr", "limit"},
+		},
 		// Add PER PARTITION LIMIT
 		{
 			B: Select("cycling.cyclist_name").Where(w).LimitPerPartition(10),
 			S: "SELECT * FROM cycling.cyclist_name WHERE id=? PER PARTITION LIMIT 10 ",
 			N: []string{"expr"},
 		},
+		// Add named PER PARTITION LIMIT
+		{
+			B: Select("cycling.cyclist_name").Where(w).LimitPerPartitionNamed("partition_limit"),
+			S: "SELECT * FROM cycling.cyclist_name WHERE id=? PER PARTITION LIMIT ? ",
+			N: []string{"expr", "partition_limit"},
+		},
+		// Add PER PARTITION LIMIT and LIMIT
+		{
+			B: Select("cycling.cyclist_name").Where(w).LimitPerPartition(2).Limit(10),
+			S: "SELECT * FROM cycling.cyclist_name WHERE id=? PER PARTITION LIMIT 2 LIMIT 10 ",
+			N: []string{"expr"},
+		},
+		// Add named PER PARTITION LIMIT and LIMIT
+		{
+			B: Select("cycling.cyclist_name").Where(w).LimitPerPartitionNamed("partition_limit").LimitNamed("limit"),
+			S: "SELECT * FROM cycling.cyclist_name WHERE id=? PER PARTITION LIMIT ? LIMIT ? ",
+			N: []string{"expr", "partition_limit", "limit"},
+		},
 		// Add ALLOW FILTERING
 		{
 			B: Select("cycling.cyclist_name").Where(w).AllowFiltering(),