Skip to content

Commit 47a3ae9

Browse files
worryg0dworryg0d
committed
Support for Native Protocol 5 release version
1 parent 974fa12 commit 47a3ae9

16 files changed

+1424
-46
lines changed

‎batch_test.go

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
package gocql
2929

3030
import (
31+
"github.com/stretchr/testify/require"
3132
"testing"
3233
"time"
3334
)
@@ -84,3 +85,84 @@ func TestBatch_WithTimestamp(t *testing.T) {
8485
t.Errorf("got ts %d, expected %d", storedTs, micros)
8586
}
8687
}
88+
89+
func TestBatch_WithNowInSeconds(t *testing.T) {
90+
session := createSession(t)
91+
defer session.Close()
92+
93+
if session.cfg.ProtoVersion < protoVersion5 {
94+
t.Skip("Batch now in seconds are only available on protocol >= 5")
95+
}
96+
97+
if err := createTable(session, `CREATE TABLE batch_now_in_seconds (id int primary key, val text)`); err != nil {
98+
t.Fatal(err)
99+
}
100+
101+
b := session.NewBatch(LoggedBatch)
102+
b.WithNowInSeconds(0)
103+
b.Query("INSERT INTO batch_now_in_seconds (id, val) VALUES (?, ?) USING TTL 20", 1, "val")
104+
if err := session.ExecuteBatch(b); err != nil {
105+
t.Fatal(err)
106+
}
107+
108+
var remainingTTL int
109+
err := session.Query(`SELECT TTL(val) FROM batch_now_in_seconds WHERE id = ?`, 1).
110+
WithNowInSeconds(10).
111+
Scan(&remainingTTL)
112+
if err != nil {
113+
t.Fatal(err)
114+
}
115+
116+
require.Equal(t, remainingTTL, 10)
117+
}
118+
119+
func TestBatch_SetKeyspace(t *testing.T) {
120+
session := createSession(t)
121+
defer session.Close()
122+
123+
if session.cfg.ProtoVersion < protoVersion5 {
124+
t.Skip("keyspace for BATCH message is not supported in protocol < 5")
125+
}
126+
127+
const keyspaceStmt = `
128+
CREATE KEYSPACE IF NOT EXISTS gocql_keyspace_override_test
129+
WITH replication = {
130+
'class': 'SimpleStrategy',
131+
'replication_factor': '1'
132+
};
133+
`
134+
135+
err := session.Query(keyspaceStmt).Exec()
136+
if err != nil {
137+
t.Fatal(err)
138+
}
139+
140+
err = createTable(session, "CREATE TABLE IF NOT EXISTS gocql_keyspace_override_test.batch_keyspace(id int, value text, PRIMARY KEY (id))")
141+
if err != nil {
142+
t.Fatal(err)
143+
}
144+
145+
ids := []int{1, 2}
146+
texts := []string{"val1", "val2"}
147+
148+
b := session.NewBatch(LoggedBatch).SetKeyspace("gocql_keyspace_override_test")
149+
b.Query("INSERT INTO batch_keyspace(id, value) VALUES (?, ?)", ids[0], texts[0])
150+
b.Query("INSERT INTO batch_keyspace(id, value) VALUES (?, ?)", ids[1], texts[1])
151+
err = session.ExecuteBatch(b)
152+
if err != nil {
153+
t.Fatal(err)
154+
}
155+
156+
var (
157+
id int
158+
text string
159+
)
160+
161+
iter := session.Query("SELECT * FROM gocql_keyspace_override_test.batch_keyspace").Iter()
162+
defer iter.Close()
163+
164+
for i := 0; iter.Scan(&id, &text); i++ {
165+
require.Equal(t, id, ids[i])
166+
require.Equal(t, text, texts[i])
167+
}
168+
}

‎cassandra_test.go

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ import (
3232
"context"
3333
"errors"
3434
"fmt"
35+
"github.com/stretchr/testify/require"
3536
"io"
3637
"math"
3738
"math/big"
@@ -3288,3 +3289,150 @@ func TestQuery_NamedValues(t *testing.T) {
32883289
t.Fatal(err)
32893290
}
32903291
}
3292+
3293+
func TestQuery_WithNowInSeconds(t *testing.T) {
3294+
session := createSession(t)
3295+
defer session.Close()
3296+
3297+
if session.cfg.ProtoVersion < protoVersion5 {
3298+
t.Skip("Query now in seconds are only available on protocol >= 5")
3299+
}
3300+
3301+
if err := createTable(session, `CREATE TABLE query_now_in_seconds (id int primary key, val text)`); err != nil {
3302+
t.Fatal(err)
3303+
}
3304+
3305+
err := session.Query("INSERT INTO query_now_in_seconds (id, val) VALUES (?, ?) USING TTL 20", 1, "val").
3306+
WithNowInSeconds(int(0)).
3307+
Exec()
3308+
if err != nil {
3309+
t.Fatal(err)
3310+
}
3311+
3312+
var remainingTTL int
3313+
err = session.Query(`SELECT TTL(val) FROM query_now_in_seconds WHERE id = ?`, 1).
3314+
WithNowInSeconds(10).
3315+
Scan(&remainingTTL)
3316+
if err != nil {
3317+
t.Fatal(err)
3318+
}
3319+
3320+
require.Equal(t, remainingTTL, 10)
3321+
}
3322+
3323+
func TestQuery_SetKeyspace(t *testing.T) {
3324+
session := createSession(t)
3325+
defer session.Close()
3326+
3327+
if session.cfg.ProtoVersion < protoVersion5 {
3328+
t.Skip("keyspace for QUERY message is not supported in protocol < 5")
3329+
}
3330+
3331+
const keyspaceStmt = `
3332+
CREATE KEYSPACE IF NOT EXISTS gocql_query_keyspace_override_test
3333+
WITH replication = {
3334+
'class': 'SimpleStrategy',
3335+
'replication_factor': '1'
3336+
};
3337+
`
3338+
3339+
err := session.Query(keyspaceStmt).Exec()
3340+
if err != nil {
3341+
t.Fatal(err)
3342+
}
3343+
3344+
err = createTable(session, "CREATE TABLE IF NOT EXISTS gocql_query_keyspace_override_test.query_keyspace(id int, value text, PRIMARY KEY (id))")
3345+
if err != nil {
3346+
t.Fatal(err)
3347+
}
3348+
3349+
expectedID := 1
3350+
expectedText := "text"
3351+
3352+
// Testing PREPARE message
3353+
err = session.Query("INSERT INTO gocql_query_keyspace_override_test.query_keyspace (id, value) VALUES (?, ?)", expectedID, expectedText).Exec()
3354+
if err != nil {
3355+
t.Fatal(err)
3356+
}
3357+
3358+
var (
3359+
id int
3360+
text string
3361+
)
3362+
3363+
q := session.Query("SELECT * FROM gocql_query_keyspace_override_test.query_keyspace").
3364+
SetKeyspace("gocql_query_keyspace_override_test")
3365+
err = q.Scan(&id, &text)
3366+
if err != nil {
3367+
t.Fatal(err)
3368+
}
3369+
3370+
require.Equal(t, expectedID, id)
3371+
require.Equal(t, expectedText, text)
3372+
3373+
// Testing QUERY message
3374+
id = 0
3375+
text = ""
3376+
3377+
q = session.Query("SELECT * FROM gocql_query_keyspace_override_test.query_keyspace").
3378+
SetKeyspace("gocql_query_keyspace_override_test")
3379+
q.skipPrepare = true
3380+
err = q.Scan(&id, &text)
3381+
if err != nil {
3382+
t.Fatal(err)
3383+
}
3384+
3385+
require.Equal(t, expectedID, id)
3386+
require.Equal(t, expectedText, text)
3387+
}
3388+
3389+
func TestLargeSizeQuery(t *testing.T) {
3390+
session := createSession(t)
3391+
defer session.Close()
3392+
3393+
if err := createTable(session, "CREATE TABLE gocql_test.large_size_query(id int, text_col text, PRIMARY KEY (id))"); err != nil {
3394+
t.Fatal(err)
3395+
}
3396+
3397+
defer session.Close()
3398+
3399+
longString := strings.Repeat("a", 500_000)
3400+
3401+
err := session.Query("INSERT INTO gocql_test.large_size_query (id, text_col) VALUES (?, ?)", "1", longString).Exec()
3402+
if err != nil {
3403+
t.Fatal(err)
3404+
}
3405+
3406+
var result string
3407+
err = session.Query("SELECT text_col FROM gocql_test.large_size_query").Scan(&result)
3408+
if err != nil {
3409+
t.Fatal(err)
3410+
}
3411+
3412+
require.Equal(t, longString, result)
3413+
}
3414+
3415+
func TestQueryCompressionNotWorthIt(t *testing.T) {
3416+
session := createSession(t)
3417+
defer session.Close()
3418+
3419+
if err := createTable(session, "CREATE TABLE gocql_test.compression_now_worth_it(id int, text_col text, PRIMARY KEY (id))"); err != nil {
3420+
t.Fatal(err)
3421+
}
3422+
3423+
defer session.Close()
3424+
3425+
str := "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890!@#$%^&*()_+"
3426+
err := session.Query("INSERT INTO gocql_test.large_size_query (id, text_col) VALUES (?, ?)", "1", str).Exec()
3427+
if err != nil {
3428+
t.Fatal(err)
3429+
}
3430+
3431+
var result string
3432+
err = session.Query("SELECT text_col FROM gocql_test.large_size_query").Scan(&result)
3433+
if err != nil {
3434+
t.Fatal(err)
3435+
}
3436+
3437+
require.Equal(t, str, result)
3438+
}

‎common_test.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,8 @@ func createCluster(opts ...func(*ClusterConfig)) *ClusterConfig {
111111
switch *flagCompressTest {
112112
case "snappy":
113113
cluster.Compressor = &SnappyCompressor{}
114+
case "lz4":
115+
cluster.Compressor = &LZ4Compressor{}
114116
case "":
115117
default:
116118
panic("invalid compressor: " + *flagCompressTest)

‎compressor.go

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,13 +25,17 @@
2525
package gocql
2626

2727
import (
28+
"encoding/binary"
29+
"fmt"
2830
"github.com/golang/snappy"
31+
"github.com/pierrec/lz4/v4"
2932
)
3033

3134
type Compressor interface {
3235
Name() string
3336
Encode(data []byte) ([]byte, error)
3437
Decode(data []byte) ([]byte, error)
38+
DecodeSized(data []byte, size uint32) ([]byte, error)
3539
}
3640

3741
// SnappyCompressor implements the Compressor interface and can be used to
@@ -50,3 +54,51 @@ func (s SnappyCompressor) Encode(data []byte) ([]byte, error) {
5054
func (s SnappyCompressor) Decode(data []byte) ([]byte, error) {
5155
return snappy.Decode(nil, data)
5256
}
57+
58+
func (s SnappyCompressor) DecodeSized(data []byte, size uint32) ([]byte, error) {
59+
buf := make([]byte, size)
60+
return snappy.Decode(buf, data)
61+
}
62+
63+
type LZ4Compressor struct{}
64+
65+
func (s LZ4Compressor) Name() string {
66+
return "lz4"
67+
}
68+
69+
func (s LZ4Compressor) Encode(data []byte) ([]byte, error) {
70+
buf := make([]byte, lz4.CompressBlockBound(len(data)+4))
71+
var compressor lz4.Compressor
72+
n, err := compressor.CompressBlock(data, buf[4:])
73+
// According to lz4.CompressBlock doc, it doesn't fail as long as the dst
74+
// buffer length is at least lz4.CompressBlockBound(len(data))) bytes, but
75+
// we check for error anyway just to be thorough.
76+
if err != nil {
77+
return nil, err
78+
}
79+
binary.BigEndian.PutUint32(buf, uint32(len(data)))
80+
return buf[:n+4], nil
81+
}
82+
83+
func (s LZ4Compressor) Decode(data []byte) ([]byte, error) {
84+
if len(data) < 4 {
85+
return nil, fmt.Errorf("cassandra lz4 block size should be >4, got=%d", len(data))
86+
}
87+
uncompressedLength := binary.BigEndian.Uint32(data)
88+
if uncompressedLength == 0 {
89+
return nil, nil
90+
}
91+
buf := make([]byte, uncompressedLength)
92+
n, err := lz4.UncompressBlock(data[4:], buf)
93+
return buf[:n], err
94+
}
95+
96+
func (s LZ4Compressor) DecodeSized(data []byte, size uint32) ([]byte, error) {
97+
buf := make([]byte, size)
98+
_, err := lz4.UncompressBlock(data, buf)
99+
if err != nil {
100+
return nil, err
101+
}
102+
103+
return buf, nil
104+
}

0 commit comments

Comments
 (0)