Ver Fonte

Merge pull request #489 from smira/tuple-marshal

Add support for tuple marshalling
Chris Bannister há 10 anos atrás
pai
commit
6e86253e1c
3 ficheiros alterados com 71 adições e 0 exclusões
  1. 1 0
      AUTHORS
  2. 33 0
      marshal.go
  3. 37 0
      marshal_test.go

+ 1 - 0
AUTHORS

@@ -54,3 +54,4 @@ Jamie Cuthill <jamie.cuthill@gmail.com>
 Adrian Casajus <adriancasajus@gmail.com>
 John Weldon <johnweldon4@gmail.com>
 Adrien Bustany <adrien@bustany.org>
+Andrey Smirnov <smirnov.andrey@gmail.com>

+ 33 - 0
marshal.go

@@ -91,6 +91,8 @@ func Marshal(info TypeInfo, value interface{}) ([]byte, error) {
 		return marshalVarint(info, value)
 	case TypeInet:
 		return marshalInet(info, value)
+	case TypeTuple:
+		return marshalTuple(info, value)
 	case TypeUDT:
 		return marshalUDT(info, value)
 	}
@@ -1194,6 +1196,37 @@ func unmarshalInet(info TypeInfo, data []byte, value interface{}) error {
 	return unmarshalErrorf("cannot unmarshal %s into %T", info, value)
 }
 
+func marshalTuple(info TypeInfo, value interface{}) ([]byte, error) {
+	tuple := info.(TupleTypeInfo)
+	switch v := value.(type) {
+	case []interface{}:
+		var buf []byte
+
+		if len(v) != len(tuple.Elems) {
+			return nil, unmarshalErrorf("cannont marshal tuple: wrong number of elements")
+		}
+
+		for i, elem := range v {
+			data, err := Marshal(tuple.Elems[i], elem)
+			if err != nil {
+				return nil, err
+			}
+
+			n := len(data)
+			buf = append(buf, byte(n>>24),
+				byte(n>>16),
+				byte(n>>8),
+				byte(n))
+
+			buf = append(buf, data...)
+		}
+
+		return buf, nil
+	}
+
+	return nil, unmarshalErrorf("cannot marshal %T into %s", value, tuple)
+}
+
 // currently only support unmarshal into a list of values, this makes it possible
 // to support tuples without changing the query API. In the future this can be extend
 // to allow unmarshalling into custom tuple types.

+ 37 - 0
marshal_test.go

@@ -860,3 +860,40 @@ func TestMarshalTimestamp(t *testing.T) {
 		}
 	}
 }
+
+func TestMarshalTuple(t *testing.T) {
+	info := TupleTypeInfo{
+		NativeType: NativeType{proto: 3, typ: TypeTuple},
+		Elems: []TypeInfo{
+			NativeType{proto: 3, typ: TypeVarchar},
+			NativeType{proto: 3, typ: TypeVarchar},
+		},
+	}
+
+	expectedData := []byte("\x00\x00\x00\x03foo\x00\x00\x00\x03bar")
+	value := []interface{}{"foo", "bar"}
+
+	data, err := Marshal(info, value)
+	if err != nil {
+		t.Errorf("marshalTest: %v", err)
+		return
+	}
+
+	if !bytes.Equal(data, expectedData) {
+		t.Errorf("marshalTest: expected %x (%v), got %x (%v)",
+			expectedData, decBigInt(expectedData), data, decBigInt(data))
+		return
+	}
+
+	var s1, s2 string
+	val := []interface{}{&s1, &s2}
+	err = Unmarshal(info, expectedData, val)
+	if err != nil {
+		t.Errorf("unmarshalTest: %v", err)
+		return
+	}
+
+	if s1 != "foo" || s2 != "bar" {
+		t.Errorf("unmarshalTest: expected [foo, bar], got [%s, %s]", s1, s2)
+	}
+}