CSHARP-810: Save should work with a polymorphic _id (or any custom serialization of an _id).

Author: rstam

MongoDB.Driver/MongoCollection.cs

@@ -21,6 +21,7 @@
 using MongoDB.Bson;
 using MongoDB.Bson.IO;
 using MongoDB.Bson.Serialization;
+using MongoDB.Bson.Serialization.Options;
 using MongoDB.Bson.Serialization.Serializers;
 using MongoDB.Driver.Builders;
 using MongoDB.Driver.Internal;
@@ -1393,71 +1394,65 @@ public virtual WriteConcernResult Save(Type nominalType, object document)
         /// <returns>A WriteConcernResult (or null if WriteConcern is disabled).</returns>
         public virtual WriteConcernResult Save(Type nominalType, object document, MongoInsertOptions options)
         {
+            if (nominalType == null)
+            {
+                throw new ArgumentNullException("nominalType");
+            }
             if (document == null)
             {
                 throw new ArgumentNullException("document");
             }
+
             var serializer = BsonSerializer.LookupSerializer(document.GetType());
+
+            // if we can determine for sure that it is a new document and we can generate an Id for it then insert it
             var idProvider = serializer as IBsonIdProvider;
-            object id;
-            Type idNominalType;
-            IIdGenerator idGenerator;
-            if (idProvider != null && idProvider.GetDocumentId(document, out id, out idNominalType, out idGenerator))
+            if (idProvider != null)
             {
-                if (id == null && idGenerator == null)
+                object id;
+                Type idNominalType;
+                IIdGenerator idGenerator;
+                var hasId = idProvider.GetDocumentId(document, out id, out idNominalType, out idGenerator);
+
+                if (idGenerator == null && (!hasId || id == null))
                 {
                     throw new InvalidOperationException("No IdGenerator found.");
                 }
 
-                if (idGenerator != null && idGenerator.IsEmpty(id))
+                if (idGenerator != null && (!hasId || idGenerator.IsEmpty(id)))
                 {
                     id = idGenerator.GenerateId(this, document);
                     idProvider.SetDocumentId(document, id);
                     return Insert(nominalType, document, options);
                 }
-                else
-                {
-                    BsonValue idBsonValue;
-                    var documentType = document.GetType();
-                    if (BsonClassMap.IsClassMapRegistered(documentType))
-                    {
-                        var classMap = BsonClassMap.LookupClassMap(documentType);
-                        var idMemberMap = classMap.IdMemberMap;
-                        var idSerializer = idMemberMap.GetSerializer(id.GetType());
-                        // we only care about the serialized _id value but we need a dummy document to serialize it into
-                        var bsonDocument = new BsonDocument();
-                        var bsonDocumentWriterSettings = new BsonDocumentWriterSettings
-                        {
-                            GuidRepresentation = _settings.GuidRepresentation
-                        };
-                        var bsonWriter = new BsonDocumentWriter(bsonDocument, bsonDocumentWriterSettings);
-                        bsonWriter.WriteStartDocument();
-                        bsonWriter.WriteName("_id");
-                        idSerializer.Serialize(bsonWriter, id.GetType(), id, idMemberMap.SerializationOptions);
-                        bsonWriter.WriteEndDocument();
-                        idBsonValue = bsonDocument[0]; // extract the _id value from the dummy document
-                    } else {
-                        if (!BsonTypeMapper.TryMapToBsonValue(id, out idBsonValue))
-                        {
-                            idBsonValue = BsonDocumentWrapper.Create(idNominalType, id);
-                        }
-                    }
+            }
 
-                    var query = Query.EQ("_id", idBsonValue);
-                    var update = Builders.Update.Replace(nominalType, document);
-                    var updateOptions = new MongoUpdateOptions
-                    {
-                        CheckElementNames = options.CheckElementNames,
-                        Flags = UpdateFlags.Upsert,
-                        WriteConcern = options.WriteConcern
-                    };
-                    return Update(query, update, updateOptions);
-                }
+            // since we can't determine for sure whether it's a new document or not upsert it
+            // the only safe way to get the serialized _id value needed for the query is to serialize the entire document
+
+            var bsonDocument = new BsonDocument();
+            var writerSettings = new BsonDocumentWriterSettings { GuidRepresentation = _settings.GuidRepresentation };
+            using (var bsonWriter = new BsonDocumentWriter(bsonDocument, writerSettings))
+            {
+                serializer.Serialize(bsonWriter, nominalType, document, null);
             }
-            else
+
+            BsonValue idBsonValue;
+            if (!bsonDocument.TryGetValue("_id", out idBsonValue))
             {
                 throw new InvalidOperationException("Save can only be used with documents that have an Id.");
             }
+
+            var query = Query.EQ("_id", idBsonValue);
+            var update = Builders.Update.Replace(bsonDocument);
+            var updateOptions = new MongoUpdateOptions
+            {
+                CheckElementNames = options.CheckElementNames,
+                Flags = UpdateFlags.Upsert,
+                WriteConcern = options.WriteConcern
+            };
+
+            return Update(query, update, updateOptions);
         }
 
         /// <summary>