diff --git a/README.md b/README.md index 1ba1584..c55e1b1 100644 --- a/README.md +++ b/README.md @@ -91,8 +91,8 @@ type Person struct { i := qb.Insert("gocqlx_test.person").Columns("first_name", "last_name", "email") stmt, names := qb.Batch(). - Add("a.", i). - Add("b.", i). + AddWithPrefix("a", i). + AddWithPrefix("b", i). ToCql() batch := struct { diff --git a/example_test.go b/example_test.go index 39a82da..e55e417 100644 --- a/example_test.go +++ b/example_test.go @@ -90,8 +90,8 @@ func TestExample(t *testing.T) { i := qb.Insert("gocqlx_test.person").Columns("first_name", "last_name", "email") stmt, names := qb.Batch(). - Add("a.", i). - Add("b.", i). + AddWithPrefix("a", i). + AddWithPrefix("b", i). ToCql() batch := struct { diff --git a/qb/batch.go b/qb/batch.go index bcfa7fc..d36409f 100644 --- a/qb/batch.go +++ b/qb/batch.go @@ -55,6 +55,30 @@ func (b *BatchBuilder) ToCql() (stmt string, names []string) { return } +// Add adds another statement to the batch. +func (b *BatchBuilder) Add(builder Builder) *BatchBuilder { + stmt, names := builder.ToCql() + + b.stmts = append(b.stmts, stmt) + b.names = append(b.names, names...) + return b +} + +// AddWithPrefix adds another statement to the batch. Names returned by the +// builder are prefixed with the prefix + ".". +func (b *BatchBuilder) AddWithPrefix(prefix string, builder Builder) *BatchBuilder { + stmt, names := builder.ToCql() + + b.stmts = append(b.stmts, stmt) + for _, name := range names { + if prefix != "" { + name = fmt.Sprint(prefix, ".", name) + } + b.names = append(b.names, name) + } + return b +} + // UnLogged sets a UNLOGGED BATCH clause on the query. func (b *BatchBuilder) UnLogged() *BatchBuilder { b.unlogged = true @@ -78,14 +102,3 @@ 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 index 1168a71..facaf45 100644 --- a/qb/batch_test.go +++ b/qb/batch_test.go @@ -25,15 +25,15 @@ func TestBatchBuilder(t *testing.T) { }{ // Basic test for Batch { - B: Batch().Add("a.", m), + B: Batch().Add(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"}, + N: []string{"id", "user_uuid", "firstname"}, }, // Add statement { B: Batch(). - Add("a.", m). - Add("b.", m), + AddWithPrefix("a", m). + AddWithPrefix("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"}, }, @@ -75,6 +75,6 @@ func TestBatchBuilder(t *testing.T) { 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() + Batch().Add(mockBuilder{"INSERT INTO cycling.cyclist_name (id,user_uuid,firstname) VALUES (?,?,?) ", []string{"id", "user_uuid", "firstname"}}).ToCql() } }