GODRIVER-650 Add MarshalBSON and UnmarshalBSON to DefaultValueEncoders and DefaultValueDecoders

Change-Id: Ie989f65078a48062a8109ee118d389e1a70c490a

Author: Isabella Siu

bson/bsoncodec/default_value_decoders.go

@@ -48,6 +48,7 @@ func (dvd DefaultValueDecoders) RegisterDefaultDecoders(rb *RegistryBuilder) {
 		RegisterDecoder(reflect.PtrTo(tJSONNumber), ValueDecoderFunc(dvd.JSONNumberDecodeValue)).
 		RegisterDecoder(reflect.PtrTo(tURL), ValueDecoderFunc(dvd.URLDecodeValue)).
 		RegisterDecoder(tValueUnmarshaler, ValueDecoderFunc(dvd.ValueUnmarshalerDecodeValue)).
+		RegisterDecoder(tUnmarshaler, ValueDecoderFunc(dvd.UnmarshalerDecodeValue)).
 		RegisterDefaultDecoder(reflect.Bool, ValueDecoderFunc(dvd.BooleanDecodeValue)).
 		RegisterDefaultDecoder(reflect.Int, ValueDecoderFunc(dvd.IntDecodeValue)).
 		RegisterDefaultDecoder(reflect.Int8, ValueDecoderFunc(dvd.IntDecodeValue)).
@@ -735,6 +736,32 @@ func (dvd DefaultValueDecoders) ValueUnmarshalerDecodeValue(dc DecodeContext, vr
 	return valueUnmarshaler.UnmarshalBSONValue(t, src)
 }
 
+// UnmarshalerDecodeValue is the ValueDecoderFunc for Unmarshaler implementations.
+func (dvd DefaultValueDecoders) UnmarshalerDecodeValue(dc DecodeContext, vr bsonrw.ValueReader, i interface{}) error {
+	val := reflect.ValueOf(i)
+	var unmarshaler Unmarshaler
+	if val.Kind() == reflect.Ptr && val.IsNil() {
+		return fmt.Errorf("UnmarshalerDecodeValue can only unmarshal into non-nil Unmarshaler values, got %T", i)
+	}
+	if val.Type().Implements(tUnmarshaler) {
+		unmarshaler = val.Interface().(Unmarshaler)
+	} else if val.Type().Kind() == reflect.Ptr && val.Elem().Type().Implements(tUnmarshaler) {
+		if val.Elem().Kind() == reflect.Ptr && val.Elem().IsNil() {
+			val.Elem().Set(reflect.New(val.Type().Elem().Elem()))
+		}
+		unmarshaler = val.Elem().Interface().(Unmarshaler)
+	} else {
+		return fmt.Errorf("UnmarshalerDecodeValue can only handle types or pointers to types that are a Unmarshaler, got %T", i)
+	}
+
+	_, src, err := bsonrw.Copier{}.CopyValueToBytes(vr)
+	if err != nil {
+		return err
+	}
+
+	return unmarshaler.UnmarshalBSON(src)
+}
+
 // EmptyInterfaceDecodeValue is the ValueDecoderFunc for interface{}.
 func (dvd DefaultValueDecoders) EmptyInterfaceDecodeValue(dc DecodeContext, vr bsonrw.ValueReader, i interface{}) error {
 	target, ok := i.(*interface{})

bson/bsoncodec/default_value_decoders_test.go

@@ -52,6 +52,7 @@ func TestDefaultValueDecoders(t *testing.T) {
 	now := time.Now().Truncate(time.Millisecond)
 	d128 := decimal.NewDecimal128(12345, 67890)
 	var ptrPtrValueUnmarshaler **testValueUnmarshaler
+	var ptrPtrUnmarshaler **testUnmarshaler
 
 	type subtest struct {
 		name   string
@@ -1099,6 +1100,36 @@ func TestDefaultValueDecoders(t *testing.T) {
 				},
 			},
 		},
+		{
+			"UnmarshalerDecodeValue",
+			ValueDecoderFunc(dvd.UnmarshalerDecodeValue),
+			[]subtest{
+				{
+					"wrong type",
+					wrong,
+					nil,
+					nil,
+					bsonrwtest.Nothing,
+					fmt.Errorf("UnmarshalerDecodeValue can only handle types or pointers to types that are a Unmarshaler, got %T", &wrong),
+				},
+				{
+					"Unmarshaler",
+					testUnmarshaler{Val: bsoncore.AppendDouble(nil, 3.14159)},
+					nil,
+					&bsonrwtest.ValueReaderWriter{BSONType: bsontype.Double, Return: float64(3.14159)},
+					bsonrwtest.ReadDouble,
+					nil,
+				},
+				{
+					"nil pointer to Unmarshaler",
+					ptrPtrUnmarshaler,
+					nil,
+					&bsonrwtest.ValueReaderWriter{BSONType: bsontype.Double, Return: float64(3.14159)},
+					bsonrwtest.Nothing,
+					fmt.Errorf("UnmarshalerDecodeValue can only unmarshal into non-nil Unmarshaler values, got %T", ptrPtrUnmarshaler),
+				},
+			},
+		},
 	}
 
 	for _, tc := range testCases {
@@ -1698,6 +1729,17 @@ func (tvu *testValueUnmarshaler) UnmarshalBSONValue(t bsontype.Type, val []byte)
 	tvu.t, tvu.val = t, val
 	return tvu.err
 }
+
+type testUnmarshaler struct {
+	Val []byte
+	Err error
+}
+
+func (tvu *testUnmarshaler) UnmarshalBSON(val []byte) error {
+	tvu.Val = val
+	return tvu.Err
+}
+
 func (tvu testValueUnmarshaler) Equal(tvu2 testValueUnmarshaler) bool {
 	return tvu.t == tvu2.t && bytes.Equal(tvu.val, tvu2.val)
 }

bson/bsoncodec/default_value_encoders.go

@@ -16,6 +16,7 @@ import (
 	"time"
 
 	"github.com/mongodb/mongo-go-driver/bson/bsonrw"
+	"github.com/mongodb/mongo-go-driver/bson/bsontype"
 	"github.com/mongodb/mongo-go-driver/bson/decimal"
 	"github.com/mongodb/mongo-go-driver/bson/objectid"
 )
@@ -41,6 +42,7 @@ func (dve DefaultValueEncoders) RegisterDefaultEncoders(rb *RegistryBuilder) {
 		RegisterEncoder(tJSONNumber, ValueEncoderFunc(dve.JSONNumberEncodeValue)).
 		RegisterEncoder(tURL, ValueEncoderFunc(dve.URLEncodeValue)).
 		RegisterEncoder(tValueMarshaler, ValueEncoderFunc(dve.ValueMarshalerEncodeValue)).
+		RegisterEncoder(tMarshaler, ValueEncoderFunc(dve.MarshalerEncodeValue)).
 		RegisterEncoder(tProxy, ValueEncoderFunc(dve.ProxyEncodeValue)).
 		RegisterDefaultEncoder(reflect.Bool, ValueEncoderFunc(dve.BooleanEncodeValue)).
 		RegisterDefaultEncoder(reflect.Int, ValueEncoderFunc(dve.IntEncodeValue)).
@@ -468,6 +470,24 @@ func (dve DefaultValueEncoders) ValueMarshalerEncodeValue(ec EncodeContext, vw b
 	return bsonrw.Copier{}.CopyValueFromBytes(vw, t, val)
 }
 
+// MarshalerEncodeValue is the ValueEncoderFunc for Marshaler implementations.
+func (dve DefaultValueEncoders) MarshalerEncodeValue(ec EncodeContext, vw bsonrw.ValueWriter, i interface{}) error {
+	vm, ok := i.(Marshaler)
+	if !ok {
+		return ValueEncoderError{
+			Name:     "MarshalerEncodeValue",
+			Types:    []interface{}{(Marshaler)(nil)},
+			Received: i,
+		}
+	}
+
+	val, err := vm.MarshalBSON()
+	if err != nil {
+		return err
+	}
+	return bsonrw.Copier{}.CopyValueFromBytes(vw, bsontype.EmbeddedDocument, val)
+}
+
 // ProxyEncodeValue is the ValueEncoderFunc for Proxy implementations.
 func (dve DefaultValueEncoders) ProxyEncodeValue(ec EncodeContext, vw bsonrw.ValueWriter, i interface{}) error {
 	proxy, ok := i.(Proxy)

bson/bsoncodec/default_value_encoders_test.go

@@ -472,6 +472,40 @@ func TestDefaultValueEncoders(t *testing.T) {
 				},
 			},
 		},
+		{
+			"MarshalerEncodeValue",
+			ValueEncoderFunc(dve.MarshalerEncodeValue),
+			[]subtest{
+				{
+					"wrong type",
+					wrong,
+					nil,
+					nil,
+					bsonrwtest.Nothing,
+					ValueEncoderError{
+						Name:     "MarshalerEncodeValue",
+						Types:    []interface{}{(ValueMarshaler)(nil)},
+						Received: wrong,
+					},
+				},
+				{
+					"MarshalBSON error",
+					testMarshaler{err: errors.New("mbson error")},
+					nil,
+					nil,
+					bsonrwtest.Nothing,
+					errors.New("mbson error"),
+				},
+				{
+					"success",
+					testMarshaler{buf: bsoncore.BuildDocument(nil, bsoncore.AppendDoubleElement(nil, "pi", 3.14159))},
+					nil,
+					nil,
+					bsonrwtest.WriteDocumentEnd,
+					nil,
+				},
+			},
+		},
 		{
 			"ProxyEncodeValue",
 			ValueEncoderFunc(dve.ProxyEncodeValue),
@@ -998,6 +1032,15 @@ func (tvm testValueMarshaler) MarshalBSONValue() (bsontype.Type, []byte, error)
 	return tvm.t, tvm.buf, tvm.err
 }
 
+type testMarshaler struct {
+	buf []byte
+	err error
+}
+
+func (tvm testMarshaler) MarshalBSON() ([]byte, error) {
+	return tvm.buf, tvm.err
+}
+
 type testProxy struct {
 	ret interface{}
 	err error

bson/bsoncodec/types.go

@@ -57,4 +57,6 @@ var tJSONNumber = reflect.TypeOf(json.Number(""))
 
 var tValueMarshaler = reflect.TypeOf((*ValueMarshaler)(nil)).Elem()
 var tValueUnmarshaler = reflect.TypeOf((*ValueUnmarshaler)(nil)).Elem()
+var tMarshaler = reflect.TypeOf((*Marshaler)(nil)).Elem()
+var tUnmarshaler = reflect.TypeOf((*Unmarshaler)(nil)).Elem()
 var tProxy = reflect.TypeOf((*Proxy)(nil)).Elem()