CASSGO-11 Support vector type

cassandra_test.go

@@ -2245,15 +2245,16 @@ func TestViewMetadata(t *testing.T) {
 		textType = TypeVarchar
 	}
 
+	protoVer := byte(session.cfg.ProtoVersion)
 	expectedView := ViewMetadata{
 		Keyspace:   "gocql_test",
 		Name:       "basicview",
 		FieldNames: []string{"birthday", "nationality", "weight", "height"},
 		FieldTypes: []TypeInfo{
-			NativeType{typ: TypeTimestamp},
-			NativeType{typ: textType},
-			NativeType{typ: textType},
-			NativeType{typ: textType},
+			NativeType{typ: TypeTimestamp, proto: protoVer},
+			NativeType{typ: textType, proto: protoVer},
+			NativeType{typ: textType, proto: protoVer},
+			NativeType{typ: textType, proto: protoVer},
 		},
 	}
 
@@ -2351,18 +2352,19 @@ func TestAggregateMetadata(t *testing.T) {
 		t.Fatal("expected two aggregates")
 	}
 
+	protoVer := byte(session.cfg.ProtoVersion)
 	expectedAggregrate := AggregateMetadata{
 		Keyspace:      "gocql_test",
 		Name:          "average",
-		ArgumentTypes: []TypeInfo{NativeType{typ: TypeInt}},
+		ArgumentTypes: []TypeInfo{NativeType{typ: TypeInt, proto: protoVer}},
 		InitCond:      "(0, 0)",
-		ReturnType:    NativeType{typ: TypeDouble},
+		ReturnType:    NativeType{typ: TypeDouble, proto: protoVer},
 		StateType: TupleTypeInfo{
-			NativeType: NativeType{typ: TypeTuple},
+			NativeType: NativeType{typ: TypeTuple, proto: protoVer},
 
 			Elems: []TypeInfo{
-				NativeType{typ: TypeInt},
-				NativeType{typ: TypeBigInt},
+				NativeType{typ: TypeInt, proto: protoVer},
+				NativeType{typ: TypeBigInt, proto: protoVer},
 			},
 		},
 		stateFunc: "avgstate",
@@ -2401,28 +2403,29 @@ func TestFunctionMetadata(t *testing.T) {
 	avgState := functions[1]
 	avgFinal := functions[0]
 
+	protoVer := byte(session.cfg.ProtoVersion)
 	avgStateBody := "if (val !=null) {state.setInt(0, state.getInt(0)+1); state.setLong(1, state.getLong(1)+val.intValue());}return state;"
 	expectedAvgState := FunctionMetadata{
 		Keyspace: "gocql_test",
 		Name:     "avgstate",
 		ArgumentTypes: []TypeInfo{
 			TupleTypeInfo{
-				NativeType: NativeType{typ: TypeTuple},
+				NativeType: NativeType{typ: TypeTuple, proto: protoVer},
 
 				Elems: []TypeInfo{
-					NativeType{typ: TypeInt},
-					NativeType{typ: TypeBigInt},
+					NativeType{typ: TypeInt, proto: protoVer},
+					NativeType{typ: TypeBigInt, proto: protoVer},
 				},
 			},
-			NativeType{typ: TypeInt},
+			NativeType{typ: TypeInt, proto: protoVer},
 		},
 		ArgumentNames: []string{"state", "val"},
 		ReturnType: TupleTypeInfo{
-			NativeType: NativeType{typ: TypeTuple},
+			NativeType: NativeType{typ: TypeTuple, proto: protoVer},
 
 			Elems: []TypeInfo{
-				NativeType{typ: TypeInt},
-				NativeType{typ: TypeBigInt},
+				NativeType{typ: TypeInt, proto: protoVer},
+				NativeType{typ: TypeBigInt, proto: protoVer},
 			},
 		},
 		CalledOnNullInput: true,
@@ -2439,16 +2442,16 @@ func TestFunctionMetadata(t *testing.T) {
 		Name:     "avgfinal",
 		ArgumentTypes: []TypeInfo{
 			TupleTypeInfo{
-				NativeType: NativeType{typ: TypeTuple},
+				NativeType: NativeType{typ: TypeTuple, proto: protoVer},
 
 				Elems: []TypeInfo{
-					NativeType{typ: TypeInt},
-					NativeType{typ: TypeBigInt},
+					NativeType{typ: TypeInt, proto: protoVer},
+					NativeType{typ: TypeBigInt, proto: protoVer},
 				},
 			},
 		},
 		ArgumentNames:     []string{"state"},
-		ReturnType:        NativeType{typ: TypeDouble},
+		ReturnType:        NativeType{typ: TypeDouble, proto: protoVer},
 		CalledOnNullInput: true,
 		Language:          "java",
 		Body:              finalStateBody,
@@ -2557,15 +2560,16 @@ func TestKeyspaceMetadata(t *testing.T) {
 	if flagCassVersion.Before(3, 0, 0) {
 		textType = TypeVarchar
 	}
+	protoVer := byte(session.cfg.ProtoVersion)
 	expectedType := UserTypeMetadata{
 		Keyspace:   "gocql_test",
 		Name:       "basicview",
 		FieldNames: []string{"birthday", "nationality", "weight", "height"},
 		FieldTypes: []TypeInfo{
-			NativeType{typ: TypeTimestamp},
-			NativeType{typ: textType},
-			NativeType{typ: textType},
-			NativeType{typ: textType},
+			NativeType{typ: TypeTimestamp, proto: protoVer},
+			NativeType{typ: textType, proto: protoVer},
+			NativeType{typ: textType, proto: protoVer},
+			NativeType{typ: textType, proto: protoVer},
 		},
 	}
 	if !reflect.DeepEqual(*keyspaceMetadata.UserTypes["basicview"], expectedType) {

common_test.go

@@ -28,6 +28,7 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"math/rand"
 	"net"
 	"reflect"
 	"strings"
@@ -52,6 +53,10 @@ var (
 	flagCassVersion cassVersion
 )
 
+var seededRand *rand.Rand = rand.New(rand.NewSource(time.Now().UnixNano()))
+
+const randCharset = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+
 func init() {
 	flag.Var(&flagCassVersion, "gocql.cversion", "the cassandra version being tested against")
 
@@ -277,6 +282,14 @@ func assertTrue(t *testing.T, description string, value bool) {
 	}
 }
 
+func randomText(size int) string {
+	result := make([]byte, size)
+	for i := range result {
+		result[i] = randCharset[rand.Intn(len(randCharset))]
+	}
+	return string(result)
+}
+
 func assertEqual(t *testing.T, description string, expected, actual interface{}) {
 	t.Helper()
 	if expected != actual {

frame.go

@@ -32,6 +32,7 @@ import (
 	"io/ioutil"
 	"net"
 	"runtime"
+	"strconv"
 	"strings"
 	"time"
 )
@@ -928,6 +929,22 @@ func (f *framer) readTypeInfo() TypeInfo {
 		collection.Elem = f.readTypeInfo()
 
 		return collection
+	case TypeCustom:
+		if strings.HasPrefix(simple.custom, VECTOR_TYPE) {
+			spec := strings.TrimPrefix(simple.custom, VECTOR_TYPE)
+			spec = spec[1 : len(spec)-1] // remove parenthesis
+			idx := strings.LastIndex(spec, ",")
+			typeStr := spec[:idx]
+			dimStr := spec[idx+1:]
+			subType := getTypeInfo(strings.TrimSpace(typeStr), f.proto, nopLogger{})
+			dim, _ := strconv.Atoi(strings.TrimSpace(dimStr))
+			vector := VectorType{
+				NativeType: simple,
+				SubType:    subType,
+				Dimensions: dim,
+			}
+			return vector
+		}
 	}
 
 	return simple

helpers.go

@@ -25,10 +25,12 @@
 package gocql
 
 import (
+	"encoding/hex"
 	"fmt"
 	"math/big"
 	"net"
 	"reflect"
+	"strconv"
 	"strings"
 	"time"
 
@@ -162,47 +164,83 @@ func getCassandraBaseType(name string) Type {
 	}
 }
 
-func getCassandraType(name string, logger StdLogger) TypeInfo {
+func getCassandraType(name string, protoVer byte, logger StdLogger) TypeInfo {
 	if strings.HasPrefix(name, "frozen<") {
-		return getCassandraType(strings.TrimPrefix(name[:len(name)-1], "frozen<"), logger)
+		return getCassandraType(strings.TrimPrefix(name[:len(name)-1], "frozen<"), protoVer, logger)
 	} else if strings.HasPrefix(name, "set<") {
 		return CollectionType{
-			NativeType: NativeType{typ: TypeSet},
-			Elem:       getCassandraType(strings.TrimPrefix(name[:len(name)-1], "set<"), logger),
+			NativeType: NativeType{typ: TypeSet, proto: protoVer},
+			Elem:       getCassandraType(strings.TrimPrefix(name[:len(name)-1], "set<"), protoVer, logger),
 		}
 	} else if strings.HasPrefix(name, "list<") {
 		return CollectionType{
-			NativeType: NativeType{typ: TypeList},
-			Elem:       getCassandraType(strings.TrimPrefix(name[:len(name)-1], "list<"), logger),
+			NativeType: NativeType{typ: TypeList, proto: protoVer},
+			Elem:       getCassandraType(strings.TrimPrefix(name[:len(name)-1], "list<"), protoVer, logger),
 		}
 	} else if strings.HasPrefix(name, "map<") {
 		names := splitCompositeTypes(strings.TrimPrefix(name[:len(name)-1], "map<"))
 		if len(names) != 2 {
 			logger.Printf("Error parsing map type, it has %d subelements, expecting 2\n", len(names))
 			return NativeType{
-				typ: TypeCustom,
+				proto: protoVer,
+				typ:   TypeCustom,
 			}
 		}
 		return CollectionType{
-			NativeType: NativeType{typ: TypeMap},
-			Key:        getCassandraType(names[0], logger),
-			Elem:       getCassandraType(names[1], logger),
+			NativeType: NativeType{typ: TypeMap, proto: protoVer},
+			Key:        getCassandraType(names[0], protoVer, logger),
+			Elem:       getCassandraType(names[1], protoVer, logger),
 		}
 	} else if strings.HasPrefix(name, "tuple<") {
 		names := splitCompositeTypes(strings.TrimPrefix(name[:len(name)-1], "tuple<"))
 		types := make([]TypeInfo, len(names))
 
 		for i, name := range names {
-			types[i] = getCassandraType(name, logger)
+			types[i] = getCassandraType(name, protoVer, logger)
 		}
 
 		return TupleTypeInfo{
-			NativeType: NativeType{typ: TypeTuple},
+			NativeType: NativeType{typ: TypeTuple, proto: protoVer},
 			Elems:      types,
 		}
+	} else if strings.HasPrefix(name, "udt<") {
+		names := splitCompositeTypes(strings.TrimPrefix(name[:len(name)-1], "udt<"))
+		fields := make([]UDTField, len(names)-2)
+
+		for i := 2; i < len(names); i++ {
+			spec := strings.Split(names[i], ":")
+			fieldName, _ := hex.DecodeString(spec[0])
+			fields[i-2] = UDTField{
+				Name: string(fieldName),
+				Type: getTypeInfo(spec[1], protoVer, logger),
+			}
+		}
+
+		udtName, _ := hex.DecodeString(names[1])
+		return UDTTypeInfo{
+			NativeType: NativeType{typ: TypeUDT, proto: protoVer},
+			KeySpace:   names[0],
+			Name:       string(udtName),
+			Elements:   fields,
+		}
+	} else if strings.HasPrefix(name, "vector<") {
+		names := splitCompositeTypes(strings.TrimPrefix(name[:len(name)-1], "vector<"))
+		subType := getCassandraType(strings.TrimSpace(names[0]), protoVer, logger)
+		dim, _ := strconv.Atoi(strings.TrimSpace(names[1]))
+
+		return VectorType{
+			NativeType: NativeType{
+				proto:  protoVer,
+				typ:    TypeCustom,
+				custom: VECTOR_TYPE,
+			},
+			SubType:    subType,
+			Dimensions: dim,
+		}
 	} else {
 		return NativeType{
-			typ: getCassandraBaseType(name),
+			proto: protoVer,
+			typ:   getCassandraBaseType(name),
 		}
 	}
 }
@@ -236,19 +274,48 @@ func splitCompositeTypes(name string) []string {
 }
 
 func apacheToCassandraType(t string) string {
-	t = strings.Replace(t, apacheCassandraTypePrefix, "", -1)
 	t = strings.Replace(t, "(", "<", -1)
 	t = strings.Replace(t, ")", ">", -1)
 	types := strings.FieldsFunc(t, func(r rune) bool {
 		return r == '<' || r == '>' || r == ','
 	})
-	for _, typ := range types {
-		t = strings.Replace(t, typ, getApacheCassandraType(typ).String(), -1)
+	skip := 0
+	for _, class := range types {
+		class = strings.TrimSpace(class)
+		if !isDigitsOnly(class) {
+			// vector types include dimension (digits) as second type parameter
+			// UDT fields are represented in format {field id}:{class}, example 66697273745f6e616d65:org.apache.cassandra.db.marshal.UTF8Type
+			if skip > 0 {
+				skip -= 1
+				continue
+			}
+			idx := strings.Index(class, ":")
+			class = class[idx+1:]
+			act := getApacheCassandraType(class)
+			val := act.String()
+			switch act {
+			case TypeUDT:
+				val = "udt"
+				skip = 2 // skip next two parameters (keyspace and type ID), do not attempt to resolve their type
+			case TypeCustom:
+				val = getApacheCassandraCustomSubType(class)
+			}
+			t = strings.Replace(t, class, val, -1)
+		}
 	}
 	// This is done so it exactly matches what Cassandra returns
 	return strings.Replace(t, ",", ", ", -1)
 }
 
+func isDigitsOnly(s string) bool {
+	for _, c := range s {
+		if c < '0' || c > '9' {
+			return false
+		}
+	}
+	return true
+}
+
 func getApacheCassandraType(class string) Type {
 	switch strings.TrimPrefix(class, apacheCassandraTypePrefix) {
 	case "AsciiType":
@@ -297,11 +364,23 @@ func getApacheCassandraType(class string) Type {
 		return TypeTuple
 	case "DurationType":
 		return TypeDuration
+	case "SimpleDateType":
+		return TypeDate
+	case "UserType":
+		return TypeUDT
 	default:
 		return TypeCustom
 	}
 }
 
+func getApacheCassandraCustomSubType(class string) string {
+	switch strings.TrimPrefix(class, apacheCassandraTypePrefix) {
+	case "VectorType":
+		return "vector"
+	}
+	return "custom"
+}
+
 func (r *RowData) rowMap(m map[string]interface{}) {
 	for i, column := range r.Columns {
 		val := dereference(r.Values[i])