Ensure primitives are returned as insertedIDs

This change updates InsertOne and InsertMany to return a Go or BSON
primitive type instead of a bsonx type.

GODRIVER-651

Change-Id: Iccaae5ee169cbf72cce09d996e3e55f46659faf8

Author: skriptble

mongo/collection.go

@@ -213,13 +213,11 @@ func (coll *Collection) InsertOne(ctx context.Context, document interface{},
 		ctx = context.Background()
 	}
 
-	doc, err := transformDocument(coll.registry, document)
+	doc, insertedID, err := transformAndEnsureID(coll.registry, document)
 	if err != nil {
 		return nil, err
 	}
 
-	doc, insertedID := ensureID(doc)
-
 	sess := sessionFromContext(ctx)
 
 	err = coll.client.ValidSession(sess)
@@ -277,11 +275,10 @@ func (coll *Collection) InsertMany(ctx context.Context, documents []interface{},
 	docs := make([]bsonx.Doc, len(documents))
 
 	for i, doc := range documents {
-		bdoc, err := transformDocument(coll.registry, doc)
+		bdoc, insertedID, err := transformAndEnsureID(coll.registry, doc)
 		if err != nil {
 			return nil, err
 		}
-		bdoc, insertedID := ensureID(bdoc)
 
 		docs[i] = bdoc
 		result[i] = insertedID

mongo/collection_internal_test.go

@@ -283,8 +283,8 @@ func TestCollection_InsertOne(t *testing.T) {
 	t.Parallel()
 
 	id := primitive.NewObjectID()
-	want := bsonx.Elem{"_id", bsonx.ObjectID(id)}
-	doc := bsonx.Doc{want, {"x", bsonx.Int32(1)}}
+	want := id
+	doc := bsonx.Doc{bsonx.Elem{"_id", bsonx.ObjectID(id)}, {"x", bsonx.Int32(1)}}
 	coll := createTestCollection(t, nil, nil)
 
 	result, err := coll.InsertOne(context.Background(), doc)
@@ -358,12 +358,12 @@ func TestCollection_InsertMany(t *testing.T) {
 
 	t.Parallel()
 
-	want1 := bsonx.Elem{"_id", bsonx.Int32(11)}
-	want2 := bsonx.Elem{"_id", bsonx.Int32(12)}
+	want1 := int32(11)
+	want2 := int32(12)
 	docs := []interface{}{
-		bsonx.Doc{want1},
+		bsonx.Doc{bsonx.Elem{"_id", bsonx.Int32(11)}},
 		bsonx.Doc{{"x", bsonx.Int32(6)}},
-		bsonx.Doc{want2},
+		bsonx.Doc{bsonx.Elem{"_id", bsonx.Int32(12)}},
 	}
 	coll := createTestCollection(t, nil, nil)
 

mongo/crud_util_test.go

@@ -591,7 +591,7 @@ func verifyInsertOneResult(t *testing.T, res *InsertOneResult, result json.RawMe
 
 	if expectedID != nil {
 		require.NotNil(t, res)
-		require.Equal(t, expectedID, res.InsertedID.(bsonx.Elem).Value.Interface())
+		require.Equal(t, expectedID, res.InsertedID)
 	}
 }
 
@@ -606,10 +606,6 @@ func verifyInsertManyResult(t *testing.T, res *InsertManyResult, result json.Raw
 	if expected.InsertedIds != nil {
 		replaceFloatsWithInts(expected.InsertedIds)
 
-		for i, elem := range res.InsertedIDs {
-			res.InsertedIDs[i] = elem.(bsonx.Elem).Value.Interface()
-		}
-
 		for _, val := range expected.InsertedIds {
 			require.Contains(t, res.InsertedIDs, val)
 		}

mongo/mongo.go

@@ -69,6 +69,53 @@ func (me MarshalError) Error() string {
 //
 type Pipeline []bson.D
 
+// transformAndEnsureID is a hack that makes it easy to get a RawValue as the _id value. This will
+// be removed when we switch from using bsonx to bsoncore for the driver package.
+func transformAndEnsureID(registry *bsoncodec.Registry, val interface{}) (bsonx.Doc, interface{}, error) {
+	// TODO: performance is going to be pretty bad for bsonx.Doc here since we turn it into a []byte
+	// only to turn it back into a bsonx.Doc. We can fix this post beta1 when we refactor the driver
+	// package to use bsoncore.Document instead of bsonx.Doc.
+	if registry == nil {
+		registry = bson.NewRegistryBuilder().Build()
+	}
+	switch tt := val.(type) {
+	case nil:
+		val = bsonx.Doc{}
+	case bsonx.Doc:
+		val = tt.Copy()
+	case []byte:
+		// Slight optimization so we'll just use MarshalBSON and not go through the codec machinery.
+		val = bson.Raw(tt)
+	}
+
+	// TODO(skriptble): Use a pool of these instead.
+	buf := make([]byte, 0, 256)
+	b, err := bson.MarshalAppendWithRegistry(registry, buf, val)
+	if err != nil {
+		return nil, nil, MarshalError{Value: val, Err: err}
+	}
+
+	d, err := bsonx.ReadDoc(b)
+	if err != nil {
+		return nil, nil, err
+	}
+
+	var id interface{}
+
+	v, err := bson.Raw(b).LookupErr("_id")
+	switch err.(type) {
+	case nil:
+		if err := v.Unmarshal(&id); err != nil {
+			return nil, nil, err
+		}
+	default:
+		oid := primitive.NewObjectID()
+		d = append(d, bsonx.Elem{"_id", bsonx.ObjectID(oid)})
+		id = oid
+	}
+	return d, id, nil
+}
+
 func transformDocument(registry *bsoncodec.Registry, val interface{}) (bsonx.Doc, error) {
 	if registry == nil {
 		registry = bson.NewRegistryBuilder().Build()