diff --git a/example_test.go b/example_test.go index 4eba8aa..627d7d9 100644 --- a/example_test.go +++ b/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() diff --git a/qb/batch.go b/qb/batch.go new file mode 100644 index 0000000..bcfa7fc --- /dev/null +++ b/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 +} diff --git a/qb/batch_test.go b/qb/batch_test.go new file mode 100644 index 0000000..1168a71 --- /dev/null +++ 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() + } +}