GODRIVER-1416 Add bson.MarshalValue (#241)

Author: Divjot Arora

bson/bsonrw/value_writer.go

@@ -56,6 +56,13 @@ func (bvwp *BSONValueWriterPool) Get(w io.Writer) ValueWriter {
 	return vw
 }
 
+// GetAtModeElement retrieves a ValueWriterFlusher from the pool and resets it to use w as the destination.
+func (bvwp *BSONValueWriterPool) GetAtModeElement(w io.Writer) ValueWriterFlusher {
+	vw := bvwp.Get(w).(*valueWriter)
+	vw.push(mElement)
+	return vw
+}
+
 // Put inserts a ValueWriter into the pool. If the ValueWriter is not a BSON ValueWriter, nothing
 // happens and ok will be false.
 func (bvwp *BSONValueWriterPool) Put(vw ValueWriter) (ok bool) {
@@ -512,17 +519,8 @@ func (vw *valueWriter) WriteDocumentEnd() error {
 	}
 
 	if vw.stack[vw.frame].mode == mTopLevel {
-		if vw.w != nil {
-			if sw, ok := vw.w.(*SliceWriter); ok {
-				*sw = vw.buf
-			} else {
-				_, err = vw.w.Write(vw.buf)
-				if err != nil {
-					return err
-				}
-				// reset buffer
-				vw.buf = vw.buf[:0]
-			}
+		if err = vw.Flush(); err != nil {
+			return err
 		}
 	}
 
@@ -537,6 +535,23 @@ func (vw *valueWriter) WriteDocumentEnd() error {
 	return nil
 }
 
+func (vw *valueWriter) Flush() error {
+	if vw.w == nil {
+		return nil
+	}
+
+	if sw, ok := vw.w.(*SliceWriter); ok {
+		*sw = vw.buf
+		return nil
+	}
+	if _, err := vw.w.Write(vw.buf); err != nil {
+		return err
+	}
+	// reset buffer
+	vw.buf = vw.buf[:0]
+	return nil
+}
+
 func (vw *valueWriter) WriteArrayElement() (ValueWriter, error) {
 	if vw.stack[vw.frame].mode != mArray {
 		return nil, vw.invalidTransitionError(mValue, "WriteArrayElement", []mode{mArray})

bson/bsonrw/writer.go

@@ -55,6 +55,12 @@ type ValueWriter interface {
 	WriteUndefined() error
 }
 
+// ValueWriterFlusher is a superset of ValueWriter that exposes functionality to flush to the underlying buffer.
+type ValueWriterFlusher interface {
+	ValueWriter
+	Flush() error
+}
+
 // BytesWriter is the interface used to write BSON bytes to a ValueWriter.
 // This interface is meant to be a superset of ValueWriter, so that types that
 // implement ValueWriter may also implement this interface.

bson/marshal.go

@@ -32,7 +32,8 @@ type ValueMarshaler interface {
 	MarshalBSONValue() (bsontype.Type, []byte, error)
 }
 
-// Marshal returns the BSON encoding of val.
+// Marshal returns the BSON encoding of val as a BSON document. If val is not a type that can be transformed into a
+// document, MarshalValue should be used instead.
 //
 // Marshal will use the default registry created by NewRegistry to recursively
 // marshal val into a []byte. Marshal will inspect struct tags and alter the
@@ -41,34 +42,37 @@ func Marshal(val interface{}) ([]byte, error) {
 	return MarshalWithRegistry(DefaultRegistry, val)
 }
 
-// MarshalAppend will append the BSON encoding of val to dst. If dst is not
-// large enough to hold the BSON encoding of val, dst will be grown.
+// MarshalAppend will encode val as a BSON document and append the bytes to dst. If dst is not large enough to hold the
+// bytes, it will be grown. If val is not a type that can be transformed into a document, MarshalValueAppend should be
+// used instead.
 func MarshalAppend(dst []byte, val interface{}) ([]byte, error) {
 	return MarshalAppendWithRegistry(DefaultRegistry, dst, val)
 }
 
-// MarshalWithRegistry returns the BSON encoding of val using Registry r.
+// MarshalWithRegistry returns the BSON encoding of val as a BSON document. If val is not a type that can be transformed
+// into a document, MarshalValueWithRegistry should be used instead.
 func MarshalWithRegistry(r *bsoncodec.Registry, val interface{}) ([]byte, error) {
 	dst := make([]byte, 0, 256) // TODO: make the default cap a constant
 	return MarshalAppendWithRegistry(r, dst, val)
 }
 
-// MarshalWithContext returns the BSON encoding of val using EncodeContext ec.
+// MarshalWithContext returns the BSON encoding of val as a BSON document using EncodeContext ec. If val is not a type
+// that can be transformed into a document, MarshalValueWithContext should be used instead.
 func MarshalWithContext(ec bsoncodec.EncodeContext, val interface{}) ([]byte, error) {
 	dst := make([]byte, 0, 256) // TODO: make the default cap a constant
 	return MarshalAppendWithContext(ec, dst, val)
 }
 
-// MarshalAppendWithRegistry will append the BSON encoding of val to dst using
-// Registry r. If dst is not large enough to hold the BSON encoding of val, dst
-// will be grown.
+// MarshalAppendWithRegistry will encode val as a BSON document using Registry r and append the bytes to dst. If dst is
+// not large enough to hold the bytes, it will be grown. If val is not a type that can be transformed into a document,
+// MarshalValueAppendWithRegistry should be used instead.
 func MarshalAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{}) ([]byte, error) {
 	return MarshalAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val)
 }
 
-// MarshalAppendWithContext will append the BSON encoding of val to dst using
-// EncodeContext ec. If dst is not large enough to hold the BSON encoding of val, dst
-// will be grown.
+// MarshalAppendWithContext will encode val as a BSON document using Registry r and EncodeContext ec and append the
+// bytes to dst. If dst is not large enough to hold the bytes, it will be grown. If val is not a type that can be
+// transformed into a document, MarshalValueAppendWithContext should be used instead.
 func MarshalAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interface{}) ([]byte, error) {
 	sw := new(bsonrw.SliceWriter)
 	*sw = dst
@@ -95,6 +99,69 @@ func MarshalAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interf
 	return *sw, nil
 }
 
+// MarshalValue returns the BSON encoding of val.
+//
+// MarshalValue will use bson.DefaultRegistry to transform val into a BSON value. If val is a struct, this function will
+// inspect struct tags and alter the marshalling process accordingly.
+func MarshalValue(val interface{}) (bsontype.Type, []byte, error) {
+	return MarshalValueWithRegistry(DefaultRegistry, val)
+}
+
+// MarshalValueAppend will append the BSON encoding of val to dst. If dst is not large enough to hold the BSON encoding
+// of val, dst will be grown.
+func MarshalValueAppend(dst []byte, val interface{}) (bsontype.Type, []byte, error) {
+	return MarshalValueAppendWithRegistry(DefaultRegistry, dst, val)
+}
+
+// MarshalValueWithRegistry returns the BSON encoding of val using Registry r.
+func MarshalValueWithRegistry(r *bsoncodec.Registry, val interface{}) (bsontype.Type, []byte, error) {
+	dst := make([]byte, 0, defaultDstCap)
+	return MarshalValueAppendWithRegistry(r, dst, val)
+}
+
+// MarshalValueWithContext returns the BSON encoding of val using EncodeContext ec.
+func MarshalValueWithContext(ec bsoncodec.EncodeContext, val interface{}) (bsontype.Type, []byte, error) {
+	dst := make([]byte, 0, defaultDstCap)
+	return MarshalValueAppendWithContext(ec, dst, val)
+}
+
+// MarshalValueAppendWithRegistry will append the BSON encoding of val to dst using Registry r. If dst is not large
+// enough to hold the BSON encoding of val, dst will be grown.
+func MarshalValueAppendWithRegistry(r *bsoncodec.Registry, dst []byte, val interface{}) (bsontype.Type, []byte, error) {
+	return MarshalValueAppendWithContext(bsoncodec.EncodeContext{Registry: r}, dst, val)
+}
+
+// MarshalValueAppendWithContext will append the BSON encoding of val to dst using EncodeContext ec. If dst is not large
+// enough to hold the BSON encoding of val, dst will be grown.
+func MarshalValueAppendWithContext(ec bsoncodec.EncodeContext, dst []byte, val interface{}) (bsontype.Type, []byte, error) {
+	// get a ValueWriter configured to write to dst
+	sw := new(bsonrw.SliceWriter)
+	*sw = dst
+	vwFlusher := bvwPool.GetAtModeElement(sw)
+
+	// get an Encoder and encode the value
+	enc := encPool.Get().(*Encoder)
+	defer encPool.Put(enc)
+	if err := enc.Reset(vwFlusher); err != nil {
+		return 0, nil, err
+	}
+	if err := enc.SetContext(ec); err != nil {
+		return 0, nil, err
+	}
+	if err := enc.Encode(val); err != nil {
+		return 0, nil, err
+	}
+
+	// flush the bytes written because we cannot guarantee that a full document has been written
+	// after the flush, *sw will be in the format
+	// [value type, 0 (null byte to indicate end of empty element name), value bytes..]
+	if err := vwFlusher.Flush(); err != nil {
+		return 0, nil, err
+	}
+	buffer := *sw
+	return bsontype.Type(buffer[0]), buffer[2:], nil
+}
+
 // MarshalExtJSON returns the extended JSON encoding of val.
 func MarshalExtJSON(val interface{}, canonical, escapeHTML bool) ([]byte, error) {
 	return MarshalExtJSONWithRegistry(DefaultRegistry, val, canonical, escapeHTML)

bson/marshal_value_test.go

@@ -0,0 +1,168 @@
+// Copyright (C) MongoDB, Inc. 2017-present.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+package bson
+
+import (
+	"io"
+	"testing"
+
+	"go.mongodb.org/mongo-driver/bson/bsoncodec"
+	"go.mongodb.org/mongo-driver/bson/bsontype"
+	"go.mongodb.org/mongo-driver/bson/primitive"
+	"go.mongodb.org/mongo-driver/internal/testutil/assert"
+	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
+)
+
+// helper type for testing MarshalValue that implements io.Reader
+type marshalValueInterfaceInner struct {
+	Foo int
+}
+
+var _ io.Reader = marshalValueInterfaceInner{}
+
+func (marshalValueInterfaceInner) Read([]byte) (int, error) {
+	return 0, nil
+}
+
+// helper type for testing MarshalValue that contains an interface
+type marshalValueInterfaceOuter struct {
+	Reader io.Reader
+}
+
+// helper type for testing MarshalValue that implements ValueMarshaler
+type marshalValueMarshaler struct {
+	Foo int
+}
+
+var _ ValueMarshaler = marshalValueMarshaler{}
+
+func (mvi marshalValueMarshaler) MarshalBSONValue() (bsontype.Type, []byte, error) {
+	return bsontype.Int32, bsoncore.AppendInt32(nil, int32(mvi.Foo)), nil
+}
+
+type marshalValueStruct struct {
+	Foo int
+}
+
+type marshalValueTestCase struct {
+	name          string
+	val           interface{}
+	expectedType  bsontype.Type
+	expectedBytes []byte
+}
+
+func TestMarshalValue(t *testing.T) {
+	oid := primitive.NewObjectID()
+	regex := primitive.Regex{Pattern: "pattern", Options: "imx"}
+	dbPointer := primitive.DBPointer{DB: "db", Pointer: primitive.NewObjectID()}
+	codeWithScope := primitive.CodeWithScope{Code: "code", Scope: D{{"a", "b"}}}
+	idx, scopeCore := bsoncore.AppendDocumentStart(nil)
+	scopeCore = bsoncore.AppendStringElement(scopeCore, "a", "b")
+	scopeCore, _ = bsoncore.AppendDocumentEnd(scopeCore, idx)
+	decimal128 := primitive.NewDecimal128(5, 10)
+	interfaceTest := marshalValueInterfaceOuter{
+		Reader: marshalValueInterfaceInner{
+			Foo: 10,
+		},
+	}
+	interfaceCore, err := Marshal(interfaceTest)
+	assert.Nil(t, err, "Marshal error: %v", err)
+	structTest := marshalValueStruct{Foo: 10}
+	structCore, err := Marshal(structTest)
+	assert.Nil(t, err, "Marshal error: %v", err)
+
+	marshalValueTestCases := []marshalValueTestCase{
+		{"double", 3.14, bsontype.Double, bsoncore.AppendDouble(nil, 3.14)},
+		{"string", "hello world", bsontype.String, bsoncore.AppendString(nil, "hello world")},
+		{"binary", primitive.Binary{1, []byte{1, 2}}, bsontype.Binary, bsoncore.AppendBinary(nil, 1, []byte{1, 2})},
+		{"undefined", primitive.Undefined{}, bsontype.Undefined, []byte{}},
+		{"object id", oid, bsontype.ObjectID, bsoncore.AppendObjectID(nil, oid)},
+		{"boolean", true, bsontype.Boolean, bsoncore.AppendBoolean(nil, true)},
+		{"datetime", primitive.DateTime(5), bsontype.DateTime, bsoncore.AppendDateTime(nil, 5)},
+		{"null", primitive.Null{}, bsontype.Null, []byte{}},
+		{"regex", regex, bsontype.Regex, bsoncore.AppendRegex(nil, regex.Pattern, regex.Options)},
+		{"dbpointer", dbPointer, bsontype.DBPointer, bsoncore.AppendDBPointer(nil, dbPointer.DB, dbPointer.Pointer)},
+		{"javascript", primitive.JavaScript("js"), bsontype.JavaScript, bsoncore.AppendJavaScript(nil, "js")},
+		{"symbol", primitive.Symbol("symbol"), bsontype.Symbol, bsoncore.AppendSymbol(nil, "symbol")},
+		{"code with scope", codeWithScope, bsontype.CodeWithScope, bsoncore.AppendCodeWithScope(nil, "code", scopeCore)},
+		{"int32", 5, bsontype.Int32, bsoncore.AppendInt32(nil, 5)},
+		{"int64", int64(5), bsontype.Int64, bsoncore.AppendInt64(nil, 5)},
+		{"timestamp", primitive.Timestamp{T: 1, I: 5}, bsontype.Timestamp, bsoncore.AppendTimestamp(nil, 1, 5)},
+		{"decimal128", decimal128, bsontype.Decimal128, bsoncore.AppendDecimal128(nil, decimal128)},
+		{"min key", primitive.MinKey{}, bsontype.MinKey, []byte{}},
+		{"max key", primitive.MaxKey{}, bsontype.MaxKey, []byte{}},
+		{"struct", structTest, bsontype.EmbeddedDocument, structCore},
+		{"interface", interfaceTest, bsontype.EmbeddedDocument, interfaceCore},
+		{"D", D{{"foo", 10}}, bsontype.EmbeddedDocument, structCore},
+		{"M", M{"foo": 10}, bsontype.EmbeddedDocument, structCore},
+		{"ValueMarshaler", marshalValueMarshaler{Foo: 10}, bsontype.Int32, bsoncore.AppendInt32(nil, 10)},
+	}
+
+	t.Run("MarshalValue", func(t *testing.T) {
+		for _, tc := range marshalValueTestCases {
+			t.Run(tc.name, func(t *testing.T) {
+				valueType, valueBytes, err := MarshalValue(tc.val)
+				assert.Nil(t, err, "MarshalValue error: %v", err)
+				compareMarshalValueResults(t, tc, valueType, valueBytes)
+			})
+		}
+	})
+	t.Run("MarshalValueAppend", func(t *testing.T) {
+		for _, tc := range marshalValueTestCases {
+			t.Run(tc.name, func(t *testing.T) {
+				valueType, valueBytes, err := MarshalValueAppend(nil, tc.val)
+				assert.Nil(t, err, "MarshalValueAppend error: %v", err)
+				compareMarshalValueResults(t, tc, valueType, valueBytes)
+			})
+		}
+	})
+	t.Run("MarshalValueWithRegistry", func(t *testing.T) {
+		for _, tc := range marshalValueTestCases {
+			t.Run(tc.name, func(t *testing.T) {
+				valueType, valueBytes, err := MarshalValueWithRegistry(DefaultRegistry, tc.val)
+				assert.Nil(t, err, "MarshalValueWithRegistry error: %v", err)
+				compareMarshalValueResults(t, tc, valueType, valueBytes)
+			})
+		}
+	})
+	t.Run("MarshalValueWithContext", func(t *testing.T) {
+		ec := bsoncodec.EncodeContext{Registry: DefaultRegistry}
+		for _, tc := range marshalValueTestCases {
+			t.Run(tc.name, func(t *testing.T) {
+				valueType, valueBytes, err := MarshalValueWithContext(ec, tc.val)
+				assert.Nil(t, err, "MarshalValueWithContext error: %v", err)
+				compareMarshalValueResults(t, tc, valueType, valueBytes)
+			})
+		}
+	})
+	t.Run("MarshalValueAppendWithRegistry", func(t *testing.T) {
+		for _, tc := range marshalValueTestCases {
+			t.Run(tc.name, func(t *testing.T) {
+				valueType, valueBytes, err := MarshalValueAppendWithRegistry(DefaultRegistry, nil, tc.val)
+				assert.Nil(t, err, "MarshalValueAppendWithRegistry error: %v", err)
+				compareMarshalValueResults(t, tc, valueType, valueBytes)
+			})
+		}
+	})
+	t.Run("MarshalValueAppendWithContext", func(t *testing.T) {
+		ec := bsoncodec.EncodeContext{Registry: DefaultRegistry}
+		for _, tc := range marshalValueTestCases {
+			t.Run(tc.name, func(t *testing.T) {
+				valueType, valueBytes, err := MarshalValueAppendWithContext(ec, nil, tc.val)
+				assert.Nil(t, err, "MarshalValueWithContext error: %v", err)
+				compareMarshalValueResults(t, tc, valueType, valueBytes)
+			})
+		}
+	})
+}
+
+func compareMarshalValueResults(t *testing.T, tc marshalValueTestCase, gotType bsontype.Type, gotBytes []byte) {
+	t.Helper()
+	expectedValue := RawValue{Type: tc.expectedType, Value: tc.expectedBytes}
+	gotValue := RawValue{Type: gotType, Value: gotBytes}
+	assert.Equal(t, expectedValue, gotValue, "value mismatch; expected %s, got %s", expectedValue, gotValue)
+}