Support of keyspace field for BATCH message

Author: worryg0d

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