From 0bac74ba88e107a6d711d267131b2cad6bf5e591 Mon Sep 17 00:00:00 2001
From: mbdavid <mauricio.david@gmail.com>
Date: Mon, 13 Apr 2020 16:44:33 -0300
Subject: [PATCH] Fixing  and _id when include occurs before filter

---
 LiteDB.Tests/Database/DbRef_Include_Tests.cs |  2 +-
 LiteDB/Client/Mapper/BsonMapper.cs           | 45 ++++++++++++++------
 LiteDB/Engine/Query/Pipeline/BasePipe.cs     |  8 +++-
 LiteDB/Utils/Constants.cs                    |  6 +--
 4 files changed, 42 insertions(+), 19 deletions(-)

diff --git a/LiteDB.Tests/Database/DbRef_Include_Tests.cs b/LiteDB.Tests/Database/DbRef_Include_Tests.cs
index 4d00c0aa2..a4fd56c2d 100644
--- a/LiteDB.Tests/Database/DbRef_Include_Tests.cs
+++ b/LiteDB.Tests/Database/DbRef_Include_Tests.cs
@@ -114,7 +114,7 @@ public void DbRef_Include()
 
                 orders.Insert(order);
 
-                orders.EnsureIndex(x => x.Customer.Id);
+                //orders.EnsureIndex(x => x.Customer.Id);
 
                 // query orders using customer $id
                 // include customer data and customer address
diff --git a/LiteDB/Client/Mapper/BsonMapper.cs b/LiteDB/Client/Mapper/BsonMapper.cs
index 1d57e2c65..ce8b79a75 100644
--- a/LiteDB/Client/Mapper/BsonMapper.cs
+++ b/LiteDB/Client/Mapper/BsonMapper.cs
@@ -460,17 +460,30 @@ private static void RegisterDbRefItem(BsonMapper mapper, MemberMapper member, IT
 
             member.Deserialize = (bson, m) =>
             {
-                var idRef = bson.IsDocument ? bson["$id"] : BsonValue.Null;
-                var missing = bson.IsDocument ? (bson["$missing"] == true) : false;
-                
+                // if not a document (maybe BsonValue.null) returns null
+                if (bson == null || bson.IsDocument == false) return null;
+
+                var doc = bson.AsDocument;
+                var idRef = doc["$id"];
+                var missing = doc["$missing"] == true;
+                var included = doc.ContainsKey("$ref") == false;
+
                 if (missing) return null;
 
-                return m.Deserialize(entity.ForType,
-                    idRef.IsNull ?
-                    bson : // if has no $id object was full loaded (via Include) - so deserialize using normal function
-                    bson.AsDocument.ContainsKey("$type") ?
-                        new BsonDocument { ["_id"] = idRef, ["_type"] = bson["$type"] } :
-                        new BsonDocument { ["_id"] = idRef }); // if has $id, deserialize object using only _id object
+                if (included)
+                {
+                    doc["_id"] = idRef;
+
+                    return m.Deserialize(entity.ForType, doc);
+                }
+                else
+                {
+                    return m.Deserialize(entity.ForType,
+                        doc.ContainsKey("$type") ?
+                            new BsonDocument { ["_id"] = idRef, ["_type"] = bson["$type"] } :
+                            new BsonDocument { ["_id"] = idRef }); // if has $id, deserialize object using only _id object
+                }
+
             };
         }
 
@@ -526,20 +539,26 @@ private static void RegisterDbRefList(BsonMapper mapper, MemberMapper member, IT
 
                 foreach (var item in array)
                 {
-                    var refId = item["$id"];
-                    var missing = item["$missing"] == true;
+                    if (item.IsDocument == false) continue;
+
+                    var doc = item.AsDocument;
+                    var idRef = doc["$id"];
+                    var missing = doc["$missing"] == true;
+                    var included = doc.ContainsKey("$ref") == false;
                     
                     // if referece document are missing, do not inlcude on output list
                     if (missing) continue;
 
                     // if refId is null was included by "include" query, so "item" is full filled document
-                    if (refId.IsNull)
+                    if (included)
                     {
+                        item["_id"] = idRef;
+
                         result.Add(item);
                     }
                     else
                     {
-                        var bsonDocument = new BsonDocument { ["_id"] = refId };
+                        var bsonDocument = new BsonDocument { ["_id"] = idRef };
 
                         if (item.AsDocument.ContainsKey("$type"))
                         {
diff --git a/LiteDB/Engine/Query/Pipeline/BasePipe.cs b/LiteDB/Engine/Query/Pipeline/BasePipe.cs
index 9426bc8d9..d3b49738f 100644
--- a/LiteDB/Engine/Query/Pipeline/BasePipe.cs
+++ b/LiteDB/Engine/Query/Pipeline/BasePipe.cs
@@ -114,10 +114,14 @@ void DoInclude(BsonDocument value)
                         // load document based on dataBlock position
                         var refDoc = lookup.Load(node);
 
-                        value.Remove("$id");
+                        //do not remove $id
                         value.Remove("$ref");
 
-                        refDoc.CopyTo(value);
+                        // copy values from refDocument into current documet (except _id - will keep $id)
+                        foreach (var element in refDoc.Where(x => x.Key != "_id"))
+                        {
+                            value[element.Key] = element.Value;
+                        }
                     }
                     else
                     {
diff --git a/LiteDB/Utils/Constants.cs b/LiteDB/Utils/Constants.cs
index a8d404245..2005b3c54 100644
--- a/LiteDB/Utils/Constants.cs
+++ b/LiteDB/Utils/Constants.cs
@@ -102,9 +102,9 @@ internal class Constants
         [Conditional("DEBUG")]
         public static void LOG(string message, string category)
         {
-            //var threadID = Environment.CurrentManagedThreadId;
-            //
-            //Debug.WriteLine(message, threadID + "|" + category);
+            var threadID = Environment.CurrentManagedThreadId;
+            
+            Debug.WriteLine(message, threadID + "|" + category);
         }
 
         /// <summary>