Ensure _id is always first in inserted documents

GODRIVER-660

Change-Id: Ieed0294e8447d50fd49b00551d064522994260f7

Author: skriptble

mongo/mongo.go

@@ -103,17 +103,30 @@ func transformAndEnsureID(registry *bsoncodec.Registry, val interface{}) (bsonx.
 
 	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
-		}
+	idx := d.IndexOf("_id")
+	var idElem bsonx.Elem
+	switch idx {
+	case -1:
+		idElem = bsonx.Elem{"_id", bsonx.ObjectID(primitive.NewObjectID())}
+		d = append(d, bsonx.Elem{})
+		copy(d[1:], d)
+		d[0] = idElem
 	default:
-		oid := primitive.NewObjectID()
-		d = append(d, bsonx.Elem{"_id", bsonx.ObjectID(oid)})
-		id = oid
+		idElem = d[idx]
+		copy(d[1:idx+1], d[0:idx])
+		d[0] = idElem
 	}
+
+	t, data, err := idElem.Value.MarshalAppendBSONValue(buf[:0])
+	if err != nil {
+		return nil, nil, err
+	}
+
+	err = bson.RawValue{Type: t, Value: data}.UnmarshalWithRegistry(registry, &id)
+	if err != nil {
+		return nil, nil, err
+	}
+
 	return d, id, nil
 }
 

mongo/mongo_test.go

@@ -81,6 +81,58 @@ func TestTransformDocument(t *testing.T) {
 	}
 }
 
+func TestTransformAndEnsureID(t *testing.T) {
+	t.Run("newly added _id should be first element", func(t *testing.T) {
+		doc := bson.D{{"foo", "bar"}, {"baz", "qux"}, {"hello", "world"}}
+		want := bsonx.Doc{
+			{"_id", bsonx.Null()}, {"foo", bsonx.String("bar")},
+			{"baz", bsonx.String("qux")}, {"hello", bsonx.String("world")},
+		}
+		got, id, err := transformAndEnsureID(bson.DefaultRegistry, doc)
+		noerr(t, err)
+		oid, ok := id.(primitive.ObjectID)
+		if !ok {
+			t.Fatalf("Expected returned id to be a %T, but was %T", primitive.ObjectID{}, id)
+		}
+		want[0] = bsonx.Elem{"_id", bsonx.ObjectID(oid)}
+		if diff := cmp.Diff(got, want, cmp.AllowUnexported(bsonx.Elem{}, bsonx.Val{})); diff != "" {
+			t.Errorf("Returned documents differ: (-got +want)\n%s", diff)
+		}
+	})
+	t.Run("existing _id should be first element", func(t *testing.T) {
+		doc := bson.D{{"foo", "bar"}, {"baz", "qux"}, {"_id", 3.14159}, {"hello", "world"}}
+		want := bsonx.Doc{
+			{"_id", bsonx.Double(3.14159)}, {"foo", bsonx.String("bar")},
+			{"baz", bsonx.String("qux")}, {"hello", bsonx.String("world")},
+		}
+		got, id, err := transformAndEnsureID(bson.DefaultRegistry, doc)
+		noerr(t, err)
+		_, ok := id.(float64)
+		if !ok {
+			t.Fatalf("Expected returned id to be a %T, but was %T", float64(0), id)
+		}
+		if diff := cmp.Diff(got, want, cmp.AllowUnexported(bsonx.Elem{}, bsonx.Val{})); diff != "" {
+			t.Errorf("Returned documents differ: (-got +want)\n%s", diff)
+		}
+	})
+	t.Run("existing _id as first element should remain first element", func(t *testing.T) {
+		doc := bson.D{{"_id", 3.14159}, {"foo", "bar"}, {"baz", "qux"}, {"hello", "world"}}
+		want := bsonx.Doc{
+			{"_id", bsonx.Double(3.14159)}, {"foo", bsonx.String("bar")},
+			{"baz", bsonx.String("qux")}, {"hello", bsonx.String("world")},
+		}
+		got, id, err := transformAndEnsureID(bson.DefaultRegistry, doc)
+		noerr(t, err)
+		_, ok := id.(float64)
+		if !ok {
+			t.Fatalf("Expected returned id to be a %T, but was %T", float64(0), id)
+		}
+		if diff := cmp.Diff(got, want, cmp.AllowUnexported(bsonx.Elem{}, bsonx.Val{})); diff != "" {
+			t.Errorf("Returned documents differ: (-got +want)\n%s", diff)
+		}
+	})
+}
+
 func TestTransformAggregatePipeline(t *testing.T) {
 	index, arr := bsoncore.AppendArrayStart(nil)
 	dindex, arr := bsoncore.AppendDocumentElementStart(arr, "0")

x/bsonx/document.go

@@ -81,15 +81,17 @@ func (d Doc) Prepend(key string, val Val) Doc {
 // does not have an element with that key, the element is appended to the
 // document instead.
 func (d Doc) Set(key string, val Val) Doc {
-	idx := d.indexOf(key)
+	idx := d.IndexOf(key)
 	if idx == -1 {
 		return append(d, Elem{Key: key, Value: val})
 	}
 	d[idx] = Elem{Key: key, Value: val}
 	return d
 }
 
-func (d Doc) indexOf(key string) int {
+// IndexOf returns the index of the first element with a key of key, or -1 if no element with a key
+// was found.
+func (d Doc) IndexOf(key string) int {
 	for i, e := range d {
 		if e.Key == key {
 			return i
@@ -100,7 +102,7 @@ func (d Doc) indexOf(key string) int {
 
 // Delete removes the element with key if it exists and returns the updated Doc.
 func (d Doc) Delete(key string) Doc {
-	idx := d.indexOf(key)
+	idx := d.IndexOf(key)
 	if idx == -1 {
 		return d
 	}
@@ -147,7 +149,7 @@ func (d Doc) LookupElementErr(key ...string) (Elem, error) {
 
 	var elem Elem
 	var err error
-	idx := d.indexOf(key[0])
+	idx := d.IndexOf(key[0])
 	if idx == -1 {
 		return Elem{}, KeyNotFound{Key: key}
 	}

x/bsonx/value.go

@@ -160,64 +160,69 @@ func (v Val) Interface() interface{} {
 
 // MarshalBSONValue implements the bsoncodec.ValueMarshaler interface.
 func (v Val) MarshalBSONValue() (bsontype.Type, []byte, error) {
+	return v.MarshalAppendBSONValue(nil)
+}
+
+// MarshalAppendBSONValue is similar to MarshalBSONValue, but allows the caller to specify a slice
+// to add the bytes to.
+func (v Val) MarshalAppendBSONValue(dst []byte) (bsontype.Type, []byte, error) {
 	t := v.Type()
-	var data []byte
 	switch v.Type() {
 	case bsontype.Double:
-		data = bsoncore.AppendDouble(data, v.Double())
+		dst = bsoncore.AppendDouble(dst, v.Double())
 	case bsontype.String:
-		data = bsoncore.AppendString(data, v.String())
+		dst = bsoncore.AppendString(dst, v.String())
 	case bsontype.EmbeddedDocument:
 		switch v.primitive.(type) {
 		case Doc:
-			t, data, _ = v.primitive.(Doc).MarshalBSONValue() // Doc.MarshalBSONValue never returns an error.
+			t, dst, _ = v.primitive.(Doc).MarshalBSONValue() // Doc.MarshalBSONValue never returns an error.
 		case MDoc:
-			t, data, _ = v.primitive.(MDoc).MarshalBSONValue() // MDoc.MarshalBSONValue never returns an error.
+			t, dst, _ = v.primitive.(MDoc).MarshalBSONValue() // MDoc.MarshalBSONValue never returns an error.
 		}
 	case bsontype.Array:
-		t, data, _ = v.Array().MarshalBSONValue() // Arr.MarshalBSON never returns an error.
+		t, dst, _ = v.Array().MarshalBSONValue() // Arr.MarshalBSON never returns an error.
 	case bsontype.Binary:
 		subtype, bindata := v.Binary()
-		data = bsoncore.AppendBinary(data, subtype, bindata)
+		dst = bsoncore.AppendBinary(dst, subtype, bindata)
 	case bsontype.Undefined:
 	case bsontype.ObjectID:
-		data = bsoncore.AppendObjectID(data, v.ObjectID())
+		dst = bsoncore.AppendObjectID(dst, v.ObjectID())
 	case bsontype.Boolean:
-		data = bsoncore.AppendBoolean(data, v.Boolean())
+		dst = bsoncore.AppendBoolean(dst, v.Boolean())
 	case bsontype.DateTime:
-		data = bsoncore.AppendDateTime(data, int64(v.DateTime()))
+		dst = bsoncore.AppendDateTime(dst, int64(v.DateTime()))
 	case bsontype.Null:
 	case bsontype.Regex:
 		pattern, options := v.Regex()
-		data = bsoncore.AppendRegex(data, pattern, options)
+		dst = bsoncore.AppendRegex(dst, pattern, options)
 	case bsontype.DBPointer:
 		ns, ptr := v.DBPointer()
-		data = bsoncore.AppendDBPointer(data, ns, ptr)
+		dst = bsoncore.AppendDBPointer(dst, ns, ptr)
 	case bsontype.JavaScript:
-		data = bsoncore.AppendJavaScript(data, string(v.JavaScript()))
+		dst = bsoncore.AppendJavaScript(dst, string(v.JavaScript()))
 	case bsontype.Symbol:
-		data = bsoncore.AppendSymbol(data, string(v.Symbol()))
+		dst = bsoncore.AppendSymbol(dst, string(v.Symbol()))
 	case bsontype.CodeWithScope:
 		code, doc := v.CodeWithScope()
 		var scope []byte
 		scope, _ = doc.MarshalBSON() // Doc.MarshalBSON never returns an error.
-		data = bsoncore.AppendCodeWithScope(data, code, scope)
+		dst = bsoncore.AppendCodeWithScope(dst, code, scope)
 	case bsontype.Int32:
-		data = bsoncore.AppendInt32(data, v.Int32())
+		dst = bsoncore.AppendInt32(dst, v.Int32())
 	case bsontype.Timestamp:
 		t, i := v.Timestamp()
-		data = bsoncore.AppendTimestamp(data, t, i)
+		dst = bsoncore.AppendTimestamp(dst, t, i)
 	case bsontype.Int64:
-		data = bsoncore.AppendInt64(data, v.Int64())
+		dst = bsoncore.AppendInt64(dst, v.Int64())
 	case bsontype.Decimal128:
-		data = bsoncore.AppendDecimal128(data, v.Decimal128())
+		dst = bsoncore.AppendDecimal128(dst, v.Decimal128())
 	case bsontype.MinKey:
 	case bsontype.MaxKey:
 	default:
 		panic(fmt.Errorf("invalid BSON type %v", t))
 	}
 
-	return t, data, nil
+	return t, dst, nil
 }
 
 // UnmarshalBSONValue implements the bsoncodec.ValueUnmarshaler interface.