Merge pull request #14 from scylladb/mmt/batch_fix

batch: extracted AddWithPrefix

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 {

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 {

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
-}

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()
 	}
 }