GODRIVER-2407 Add DocumentDecodeType to DecodeContext (#951)

Co-authored-by: Benjamin Rewis <[email protected]>

Co-authored-by: Benjamin Rewis <[email protected]>

Author: prestonvasquez

bson/bsoncodec/bsoncodec.go

@@ -13,6 +13,7 @@ import (
 
 	"go.mongodb.org/mongo-driver/bson/bsonrw"
 	"go.mongodb.org/mongo-driver/bson/bsontype"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 )
 
 var (
@@ -118,11 +119,32 @@ type EncodeContext struct {
 type DecodeContext struct {
 	*Registry
 	Truncate bool
+
 	// Ancestor is the type of a containing document. This is mainly used to determine what type
 	// should be used when decoding an embedded document into an empty interface. For example, if
 	// Ancestor is a bson.M, BSON embedded document values being decoded into an empty interface
 	// will be decoded into a bson.M.
+	//
+	// Deprecated: Use DefaultDocumentM or DefaultDocumentD instead.
 	Ancestor reflect.Type
+
+	// defaultDocumentType specifies the Go type to decode top-level and nested BSON documents into. In particular, the
+	// usage for this field is restricted to data typed as "interface{}" or "map[string]interface{}". If DocumentType is
+	// set to a type that a BSON document cannot be unmarshaled into (e.g. "string"), unmarshalling will result in an
+	// error. DocumentType overrides the Ancestor field.
+	defaultDocumentType reflect.Type
+}
+
+// DefaultDocumentM will decode empty documents using the primitive.M type. This behavior is restricted to data typed as
+// "interface{}" or "map[string]interface{}".
+func (dc *DecodeContext) DefaultDocumentM() {
+	dc.defaultDocumentType = reflect.TypeOf(primitive.M{})
+}
+
+// DefaultDocumentD will decode empty documents using the primitive.D type. This behavior is restricted to data typed as
+// "interface{}" or "map[string]interface{}".
+func (dc *DecodeContext) DefaultDocumentD() {
+	dc.defaultDocumentType = reflect.TypeOf(primitive.D{})
 }
 
 // ValueCodec is the interface that groups the methods to encode and decode

bson/bsoncodec/empty_interface_codec.go

@@ -57,11 +57,18 @@ func (eic EmptyInterfaceCodec) EncodeValue(ec EncodeContext, vw bsonrw.ValueWrit
 
 func (eic EmptyInterfaceCodec) getEmptyInterfaceDecodeType(dc DecodeContext, valueType bsontype.Type) (reflect.Type, error) {
 	isDocument := valueType == bsontype.Type(0) || valueType == bsontype.EmbeddedDocument
-	if isDocument && dc.Ancestor != nil {
-		// Using ancestor information rather than looking up the type map entry forces consistent decoding.
-		// If we're decoding into a bson.D, subdocuments should also be decoded as bson.D, even if a type map entry
-		// has been registered.
-		return dc.Ancestor, nil
+	if isDocument {
+		if dc.defaultDocumentType != nil {
+			// If the bsontype is an embedded document and the DocumentType is set on the DecodeContext, then return
+			// that type.
+			return dc.defaultDocumentType, nil
+		}
+		if dc.Ancestor != nil {
+			// Using ancestor information rather than looking up the type map entry forces consistent decoding.
+			// If we're decoding into a bson.D, subdocuments should also be decoded as bson.D, even if a type map entry
+			// has been registered.
+			return dc.Ancestor, nil
+		}
 	}
 
 	rtype, err := dc.LookupTypeMapEntry(valueType)

bson/decoder.go

@@ -33,6 +33,11 @@ var decPool = sync.Pool{
 type Decoder struct {
 	dc bsoncodec.DecodeContext
 	vr bsonrw.ValueReader
+
+	// We persist defaultDocumentM and defaultDocumentD on the Decoder to prevent overwriting from
+	// (*Decoder).SetContext.
+	defaultDocumentM bool
+	defaultDocumentD bool
 }
 
 // NewDecoder returns a new decoder that uses the DefaultRegistry to read from vr.
@@ -95,6 +100,12 @@ func (d *Decoder) Decode(val interface{}) error {
 	if err != nil {
 		return err
 	}
+	if d.defaultDocumentM {
+		d.dc.DefaultDocumentM()
+	}
+	if d.defaultDocumentD {
+		d.dc.DefaultDocumentD()
+	}
 	return decoder.DecodeValue(d.dc, d.vr, rval)
 }
 
@@ -116,3 +127,15 @@ func (d *Decoder) SetContext(dc bsoncodec.DecodeContext) error {
 	d.dc = dc
 	return nil
 }
+
+// DefaultDocumentM will decode empty documents using the primitive.M type. This behavior is restricted to data typed as
+// "interface{}" or "map[string]interface{}".
+func (d *Decoder) DefaultDocumentM() {
+	d.defaultDocumentM = true
+}
+
+// DefaultDocumentD will decode empty documents using the primitive.D type. This behavior is restricted to data typed as
+// "interface{}" or "map[string]interface{}".
+func (d *Decoder) DefaultDocumentD() {
+	d.defaultDocumentD = true
+}

bson/decoder_test.go

@@ -18,6 +18,7 @@ import (
 	"go.mongodb.org/mongo-driver/bson/bsonrw"
 	"go.mongodb.org/mongo-driver/bson/bsonrw/bsonrwtest"
 	"go.mongodb.org/mongo-driver/bson/bsontype"
+	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
 )
 
@@ -199,12 +200,12 @@ func TestDecoderv2(t *testing.T) {
 		dc2 := bsoncodec.DecodeContext{Registry: NewRegistryBuilder().Build()}
 		dec, err := NewDecoderWithContext(dc1, bsonrw.NewBSONDocumentReader([]byte{}))
 		noerr(t, err)
-		if dec.dc != dc1 {
+		if !reflect.DeepEqual(dec.dc, dc1) {
 			t.Errorf("Decoder should use the Registry provided. got %v; want %v", dec.dc, dc1)
 		}
 		err = dec.SetContext(dc2)
 		noerr(t, err)
-		if dec.dc != dc2 {
+		if !reflect.DeepEqual(dec.dc, dc2) {
 			t.Errorf("Decoder should use the Registry provided. got %v; want %v", dec.dc, dc2)
 		}
 	})
@@ -214,12 +215,12 @@ func TestDecoderv2(t *testing.T) {
 		dc2 := bsoncodec.DecodeContext{Registry: r2}
 		dec, err := NewDecoder(bsonrw.NewBSONDocumentReader([]byte{}))
 		noerr(t, err)
-		if dec.dc != dc1 {
+		if !reflect.DeepEqual(dec.dc, dc1) {
 			t.Errorf("Decoder should use the Registry provided. got %v; want %v", dec.dc, dc1)
 		}
 		err = dec.SetRegistry(r2)
 		noerr(t, err)
-		if dec.dc != dc2 {
+		if !reflect.DeepEqual(dec.dc, dc2) {
 			t.Errorf("Decoder should use the Registry provided. got %v; want %v", dec.dc, dc2)
 		}
 	})
@@ -235,6 +236,85 @@ func TestDecoderv2(t *testing.T) {
 			t.Fatalf("Decode error mismatch; expected %v, got %v", ErrDecodeToNil, err)
 		}
 	})
+	t.Run("SetDocumentType embedded map as empty interface", func(t *testing.T) {
+		type someMap map[string]interface{}
+
+		in := make(someMap)
+		in["foo"] = map[string]interface{}{"bar": "baz"}
+
+		bytes, err := Marshal(in)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		var bsonOut someMap
+		dec, err := NewDecoder(bsonrw.NewBSONDocumentReader(bytes))
+		if err != nil {
+			t.Fatal(err)
+		}
+		dec.DefaultDocumentM()
+		if err := dec.Decode(&bsonOut); err != nil {
+			t.Fatal(err)
+		}
+
+		// Ensure that interface{}-typed top-level data is converted to the document type.
+		bsonOutType := reflect.TypeOf(bsonOut)
+		inType := reflect.TypeOf(in)
+		assert.Equal(t, inType, bsonOutType, "expected %v to equal %v", inType.String(), bsonOutType.String())
+
+		// Ensure that the embedded type is a primitive map.
+		mType := reflect.TypeOf(primitive.M{})
+		bsonFooOutType := reflect.TypeOf(bsonOut["foo"])
+		assert.Equal(t, mType, bsonFooOutType, "expected %v to equal %v", mType.String(), bsonFooOutType.String())
+	})
+	t.Run("SetDocumentType for decoding into interface{} alias", func(t *testing.T) {
+		var in interface{} = map[string]interface{}{"bar": "baz"}
+
+		bytes, err := Marshal(in)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		var bsonOut interface{}
+		dec, err := NewDecoder(bsonrw.NewBSONDocumentReader(bytes))
+		if err != nil {
+			t.Fatal(err)
+		}
+		dec.DefaultDocumentD()
+		if err := dec.Decode(&bsonOut); err != nil {
+			t.Fatal(err)
+		}
+
+		// Ensure that interface{}-typed top-level data is converted to the document type.
+		dType := reflect.TypeOf(primitive.D{})
+		bsonOutType := reflect.TypeOf(bsonOut)
+		assert.Equal(t, dType, bsonOutType,
+			"expected %v to equal %v", dType.String(), bsonOutType.String())
+	})
+	t.Run("SetDocumentType for decoding into non-interface{} alias", func(t *testing.T) {
+		var in interface{} = map[string]interface{}{"bar": "baz"}
+
+		bytes, err := Marshal(in)
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		var bsonOut struct{}
+		dec, err := NewDecoder(bsonrw.NewBSONDocumentReader(bytes))
+		if err != nil {
+			t.Fatal(err)
+		}
+		dec.DefaultDocumentD()
+		if err := dec.Decode(&bsonOut); err != nil {
+			t.Fatal(err)
+		}
+
+		// Ensure that typed top-level data is not converted to the document type.
+		dType := reflect.TypeOf(primitive.D{})
+		bsonOutType := reflect.TypeOf(bsonOut)
+		assert.NotEqual(t, dType, bsonOutType,
+			"expected %v to not equal %v", dType.String(), bsonOutType.String())
+	})
 }
 
 type testUnmarshaler struct {