Support of keyspace field for BATCH message

Author: worryg0d

cassandra_test.go

@@ -3288,3 +3288,46 @@ func TestQuery_NamedValues(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestBatchKeyspaceField(t *testing.T) {
+	session := createSession(t)
+	defer session.Close()
+
+	if session.cfg.ProtoVersion < protoVersion5 {
+		t.Skip("keyspace for BATCH message is not supported in protocol < 5")
+	}
+
+	err := createTable(session, "CREATE TABLE batch_keyspace(id int, value text, PRIMARY KEY (id))")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ids := []int{1, 2}
+	texts := []string{"val1", "val2"}
+
+	b := session.NewBatch(LoggedBatch)
+	b.Query("INSERT INTO batch_keyspace(id, value) VALUES (?, ?)", ids[0], texts[0])
+	b.Query("INSERT INTO batch_keyspace(id, value) VALUES (?, ?)", ids[1], texts[1])
+	err = session.ExecuteBatch(b)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var (
+		id   int
+		text string
+	)
+
+	iter := session.Query("SELECT * FROM batch_keyspace").Iter()
+	defer iter.Close()
+
+	for i := 0; iter.Scan(&id, &text); i++ {
+		if id != ids[i] {
+			t.Fatalf("expected id %v, got %v", ids[i], id)
+		}
+
+		if text != texts[i] {
+			t.Fatalf("expected text %v, got %v", texts[i], text)
+		}
+	}
+}

conn.go

@@ -1554,6 +1554,10 @@ func (c *Conn) executeBatch(ctx context.Context, batch *Batch) *Iter {
 		customPayload:         batch.CustomPayload,
 	}
 
+	if c.version > protoVersion4 {
+		req.keyspace = c.currentKeyspace
+	}
+
 	stmts := make(map[string]string, len(batch.Entries))
 
 	for i := 0; i < n; i++ {

frame.go

@@ -1659,6 +1659,9 @@ type writeBatchFrame struct {
 
 	//v4+
 	customPayload map[string][]byte
+
+	//v5+
+	keyspace string
 }
 
 func (w *writeBatchFrame) buildFrame(framer *framer, streamID int) error {
@@ -1718,6 +1721,13 @@ func (f *framer) writeBatchFrame(streamID int, w *writeBatchFrame, customPayload
 			flags |= flagDefaultTimestamp
 		}
 
+		if w.keyspace != "" {
+			if f.proto < protoVersion5 {
+				panic(fmt.Errorf("the keyspace can only be set with protocol 5 or higher"))
+			}
+			flags |= flagWithKeyspace
+		}
+
 		if f.proto > protoVersion4 {
 			f.writeUint(uint32(flags))
 		} else {
@@ -1737,6 +1747,10 @@ func (f *framer) writeBatchFrame(streamID int, w *writeBatchFrame, customPayload
 			}
 			f.writeLong(ts)
 		}
+
+		if w.keyspace != "" {
+			f.writeString(w.keyspace)
+		}
 	}
 
 	return f.finish()