Native Protocol v5 support

Author: worryg0d

.github/workflows/main.yml

@@ -37,7 +37,7 @@ jobs:
         go: [ '1.19', '1.20' ]
         cassandra_version: [ '4.0.8', '4.1.1' ]
         auth: [ "false" ]
-        compressor: [ "snappy" ]
+        compressor: [ "lz4" ]
         tags: [ "cassandra", "integration", "ccm" ]
     steps:
       - uses: actions/checkout@v2
@@ -101,7 +101,7 @@ jobs:
           ccm status
           ccm node1 nodetool status
 
-          args="-gocql.timeout=60s -runssl -proto=4 -rf=3 -clusterSize=3 -autowait=2000ms -compressor=${{ matrix.compressor }} -gocql.cversion=$VERSION -cluster=$(ccm liveset) ./..."
+          args="-gocql.timeout=60s -runssl -proto=5 -rf=3 -clusterSize=3 -autowait=2000ms -compressor=${{ matrix.compressor }} -gocql.cversion=$VERSION -cluster=$(ccm liveset) ./..."
 
           echo "args=$args" >> $GITHUB_ENV
           echo "JVM_EXTRA_OPTS=$JVM_EXTRA_OPTS" >> $GITHUB_ENV
@@ -127,7 +127,7 @@ jobs:
       matrix:
         go: [ '1.19', '1.20' ]
         cassandra_version: [ '4.0.8' ]
-        compressor: [ "snappy" ]
+        compressor: [ "lz4" ]
         tags: [ "integration" ]
 
     steps:
@@ -190,7 +190,7 @@ jobs:
           ccm status
           ccm node1 nodetool status
 
-          args="-gocql.timeout=60s -runssl -proto=4 -rf=3 -clusterSize=1 -autowait=2000ms -compressor=${{ matrix.compressor }} -gocql.cversion=$VERSION -cluster=$(ccm liveset) ./..."
+          args="-gocql.timeout=60s -runssl -proto=5 -rf=3 -clusterSize=1 -autowait=2000ms -compressor=${{ matrix.compressor }} -gocql.cversion=$VERSION -cluster=$(ccm liveset) ./..."
 
           echo "args=$args" >> $GITHUB_ENV
           echo "JVM_EXTRA_OPTS=$JVM_EXTRA_OPTS" >> $GITHUB_ENV

cassandra_test.go

@@ -35,6 +35,7 @@ import (
 	"io"
 	"math"
 	"math/big"
+	"math/rand"
 	"net"
 	"reflect"
 	"strconv"
@@ -3288,3 +3289,40 @@ func TestQuery_NamedValues(t *testing.T) {
 		t.Fatal(err)
 	}
 }
+
+func TestLargeSizeQuery(t *testing.T) {
+	session := createSession(t)
+	defer session.Close()
+
+	if err := createTable(session, "CREATE TABLE gocql_test.large_size_query(id int, text_col text, PRIMARY KEY (id))"); err != nil {
+		t.Fatal(err)
+	}
+
+	defer session.Close()
+
+	longString := randomString(2_000_000)
+
+	err := session.Query("INSERT INTO gocql_test.large_size_query (id, text_col) VALUES (?, ?)", "1", longString).Exec()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var result string
+	err = session.Query("SELECT text_col FROM gocql_test.large_size_query").Scan(&result)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	assertEqual(t, "result should equal inserted longString", longString, result)
+}
+
+func randomString(n int) string {
+	source := rand.NewSource(time.Now().UnixMilli())
+	r := rand.New(source)
+	var aplhabet = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+	buf := make([]byte, n)
+	for i := 0; i < n; i++ {
+		buf[i] = aplhabet[r.Intn(len(aplhabet))]
+	}
+	return string(buf)
+}

common_test.go

@@ -111,6 +111,8 @@ func createCluster(opts ...func(*ClusterConfig)) *ClusterConfig {
 	switch *flagCompressTest {
 	case "snappy":
 		cluster.Compressor = &SnappyCompressor{}
+	case "lz4":
+		cluster.Compressor = &LZ4Compressor{}
 	case "":
 	default:
 		panic("invalid compressor: " + *flagCompressTest)

compressor.go

@@ -25,13 +25,17 @@
 package gocql
 
 import (
+	"encoding/binary"
+	"fmt"
 	"github.com/golang/snappy"
+	"github.com/pierrec/lz4/v4"
 )
 
 type Compressor interface {
 	Name() string
 	Encode(data []byte) ([]byte, error)
 	Decode(data []byte) ([]byte, error)
+	DecodeSized(data []byte, size uint32) ([]byte, error)
 }
 
 // SnappyCompressor implements the Compressor interface and can be used to
@@ -50,3 +54,51 @@ func (s SnappyCompressor) Encode(data []byte) ([]byte, error) {
 func (s SnappyCompressor) Decode(data []byte) ([]byte, error) {
 	return snappy.Decode(nil, data)
 }
+
+func (s SnappyCompressor) DecodeSized(data []byte, size uint32) ([]byte, error) {
+	buf := make([]byte, size)
+	return snappy.Decode(buf, data)
+}
+
+type LZ4Compressor struct{}
+
+func (s LZ4Compressor) Name() string {
+	return "lz4"
+}
+
+func (s LZ4Compressor) Encode(data []byte) ([]byte, error) {
+	buf := make([]byte, lz4.CompressBlockBound(len(data)+4))
+	var compressor lz4.Compressor
+	n, err := compressor.CompressBlock(data, buf[4:])
+	// According to lz4.CompressBlock doc, it doesn't fail as long as the dst
+	// buffer length is at least lz4.CompressBlockBound(len(data))) bytes, but
+	// we check for error anyway just to be thorough.
+	if err != nil {
+		return nil, err
+	}
+	binary.BigEndian.PutUint32(buf, uint32(len(data)))
+	return buf[:n+4], nil
+}
+
+func (s LZ4Compressor) Decode(data []byte) ([]byte, error) {
+	if len(data) < 4 {
+		return nil, fmt.Errorf("cassandra lz4 block size should be >4, got=%d", len(data))
+	}
+	uncompressedLength := binary.BigEndian.Uint32(data)
+	if uncompressedLength == 0 {
+		return nil, nil
+	}
+	buf := make([]byte, uncompressedLength)
+	n, err := lz4.UncompressBlock(data[4:], buf)
+	return buf[:n], err
+}
+
+func (s LZ4Compressor) DecodeSized(data []byte, size uint32) ([]byte, error) {
+	buf := make([]byte, size)
+	_, err := lz4.UncompressBlock(data, buf)
+	if err != nil {
+		return nil, err
+	}
+
+	return buf, nil
+}