qb: batch builder

example_test.go

@@ -155,6 +155,37 @@ func TestExample(t *testing.T) {
 			}
 		}
 
+		// Batch
+		{
+			i := qb.Insert("gocqlx_test.person").Columns("first_name", "last_name", "email")
+
+			stmt, names := qb.Batch().
+				Add("a.", i).
+				Add("b.", i).
+				ToCql()
+			q := gocqlx.Query(session.Query(stmt), names)
+
+			b := struct {
+				A Person
+				B Person
+			}{
+				A: Person{
+					"Ian",
+					"Citizen",
+					[]string{"ian.citzen@gocqlx_test.com"},
+				},
+				B: Person{
+					"Igy",
+					"Citizen",
+					[]string{"igy.citzen@gocqlx_test.com"},
+				},
+			}
+
+			if err := q.BindStruct(&b).Exec(); err != nil {
+				t.Fatal(err)
+			}
+		}
+
 		// Select
 		{
 			stmt, names := qb.Select("gocqlx_test.person").Where(qb.In("first_name")).ToCql()

qb/batch.go

@@ -0,0 +1,91 @@
+package qb
+
+import (
+	"bytes"
+	"fmt"
+)
+
+// BATCH reference:
+// https://cassandra.apache.org/doc/latest/cql/dml.html#batch
+
+// builder is interface implemented by other builders.
+type builder interface {
+	ToCql() (stmt string, names []string)
+}
+
+// BatchBuilder builds CQL BATCH statements.
+type BatchBuilder struct {
+	unlogged bool
+	counter  bool
+	using    using
+	stmts    []string
+	names    []string
+}
+
+// Batch returns a new BatchBuilder.
+func Batch() *BatchBuilder {
+	return &BatchBuilder{}
+}
+
+// ToCql builds the query into a CQL string and named args.
+func (b *BatchBuilder) ToCql() (stmt string, names []string) {
+	cql := bytes.Buffer{}
+
+	cql.WriteString("BEGIN ")
+	if b.unlogged {
+		cql.WriteString("UNLOGGED ")
+	}
+	if b.counter {
+		cql.WriteString("COUNTER ")
+	}
+	cql.WriteString("BATCH ")
+
+	names = append(names, b.using.writeCql(&cql)...)
+
+	for _, stmt := range b.stmts {
+		cql.WriteString(stmt)
+		cql.WriteByte(';')
+		cql.WriteByte(' ')
+	}
+	names = append(names, b.names...)
+
+	cql.WriteString("APPLY BATCH ")
+
+	stmt = cql.String()
+	return
+}
+
+// UnLogged sets a UNLOGGED BATCH clause on the query.
+func (b *BatchBuilder) UnLogged() *BatchBuilder {
+	b.unlogged = true
+	return b
+}
+
+// Counter sets a COUNTER BATCH clause on the query.
+func (b *BatchBuilder) Counter() *BatchBuilder {
+	b.counter = true
+	return b
+}
+
+// Timestamp sets a USING TIMESTAMP clause on the query.
+func (b *BatchBuilder) Timestamp() *BatchBuilder {
+	b.using.timestamp = true
+	return b
+}
+
+// TTL sets a USING TTL clause on the query.
+func (b *BatchBuilder) TTL() *BatchBuilder {
+	b.using.ttl = true
+	return b
+}
+
+// Add adds another batch statement from a builder.
+func (b *BatchBuilder) Add(prefix string, builder Builder) *BatchBuilder {
+	stmt, names := builder.ToCql()
+
+	b.stmts = append(b.stmts, stmt)
+	for _, name := range names {
+		b.names = append(b.names, fmt.Sprint(prefix, name))
+	}
+	return b
+}

qb/batch_test.go

@@ -0,0 +1,80 @@
+package qb
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+type mockBuilder struct {
+	stmt  string
+	names []string
+}
+
+func (b mockBuilder) ToCql() (stmt string, names []string) {
+	return b.stmt, b.names
+}
+
+func TestBatchBuilder(t *testing.T) {
+	m := mockBuilder{"INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) ", []string{"id", "user_uuid", "firstname"}}
+
+	table := []struct {
+		B *BatchBuilder
+		N []string
+		S string
+	}{
+		// Basic test for Batch
+		{
+			B: Batch().Add("a.", m),
+			S: "BEGIN BATCH INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) ; APPLY BATCH ",
+			N: []string{"a.id", "a.user_uuid", "a.firstname"},
+		},
+		// Add statement
+		{
+			B: Batch().
+				Add("a.", m).
+				Add("b.", m),
+			S: "BEGIN BATCH INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) ; INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) ; APPLY BATCH ",
+			N: []string{"a.id", "a.user_uuid", "a.firstname", "b.id", "b.user_uuid", "b.firstname"},
+		},
+		// Add UNLOGGED
+		{
+			B: Batch().UnLogged(),
+			S: "BEGIN UNLOGGED BATCH APPLY BATCH ",
+		},
+		// Add COUNTER
+		{
+			B: Batch().Counter(),
+			S: "BEGIN COUNTER BATCH APPLY BATCH ",
+		},
+		// Add TTL
+		{
+			B: Batch().TTL(),
+			S: "BEGIN BATCH USING TTL ? APPLY BATCH ",
+			N: []string{"_ttl"},
+		},
+		// Add TIMESTAMP
+		{
+			B: Batch().Timestamp(),
+			S: "BEGIN BATCH USING TIMESTAMP ? APPLY BATCH ",
+			N: []string{"_ts"},
+		},
+	}
+
+	for _, test := range table {
+		stmt, names := test.B.ToCql()
+		if diff := cmp.Diff(test.S, stmt); diff != "" {
+			t.Error(diff)
+		}
+		if diff := cmp.Diff(test.N, names); diff != "" {
+			t.Error(diff)
+		}
+	}
+}
+
+func BenchmarkBatchBuilder(b *testing.B) {
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		Batch().Add("", mockBuilder{"INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) ", []string{"id", "user_uuid", "firstname"}}).ToCql()
+	}
+}