Repository: MongoEngine/mongoengine

Issue:
Un-pickling a document removes new fields from the global document Class
Hello,
I do apologize if this was an existing issue. There were many pages to comb through :).
## TL;DR

A `setattr` is performed on `type(document)` when `__setstate__` is invoked which wipes the available fields from the global Document class instance. Instead, perhaps this should be a soft merge on the current document only. [Here is how I fixed it](https://github.com/YoApp/mongoengine/commit/d3451a8ca85f3146605c79c4dff13342a2dc1bdb).
## The symptom

When using cache or pickled object instances you may notice that new fields introduced post-cache are not available after the first de-reference of cached instance. For example:
- Cache a `User(Document)` object whose definition contains the fields: `name, password`.
- Add a new field to the definition named `email`.
- When referencing `User._fields_ordered` the value is now `name, password, email`.
- Uncache the existing user object.
- When referencing `User._fields_ordered` the value is now `name, password`.
## A semi-long explanation

The private key `_fields_ordered` is ONLY populated during a set for dynamic Documents, and in the `__new__` method of the specified metaclass. That being said, it **should** represent a single source of truth for what attributes are available and defined on the Document. This is also the field used for overriding `getattr`. That is, if you attempt to access a attribute, for example, `myuser.email`, code is executed to the extent of this:

```
if attr not in type(self)._fields_ordered:
    raise AttributeError
```

That being said, mutation of this field on a global scale doesn't seem like a good idea. Ever (from the outside looking in anyway).
## The quick solution

Above is a link to my simple commit that fixed this behavior. Basically, instead of doing a set on the global Document class it copies all attributes from the global `type(self)` and simply adds anything that was pickled but not present. That is, a soft merge, so that no fields are lost.
## Why this is a issue and not a PR

I am not 100% sure of the consequences my quick fix could have. For 1, the fields "added" from the unpickled object are not given a proper ordering. Also, There may be a possibility that writing this object to the db would re-introduce previously removed fields.
I hope that the good folks at MongoEngine as well as the contributors can sort out a proper solution.



Discussion:

Resolving patch:
diff --git a/mongoengine/base/document.py b/mongoengine/base/document.py
--- a/mongoengine/base/document.py
+++ b/mongoengine/base/document.py
@@ -206,7 +206,12 @@ def __setstate__(self, data):
             if k in data:
                 setattr(self, k, data[k])
         if '_fields_ordered' in data:
-            setattr(type(self), '_fields_ordered', data['_fields_ordered'])
+            if self._dynamic:
+                setattr(self, '_fields_ordered', data['_fields_ordered'])
+            else:
+                _super_fields_ordered = type(self)._fields_ordered
+                setattr(self, '_fields_ordered', _super_fields_ordered)
+
         dynamic_fields = data.get('_dynamic_fields') or SON()
         for k in dynamic_fields.keys():
             setattr(self, k, data["_data"].get(k))
diff --git a/mongoengine/dereference.py b/mongoengine/dereference.py
--- a/mongoengine/dereference.py
+++ b/mongoengine/dereference.py
@@ -102,24 +102,24 @@ def _find_references(self, items, depth=0):
                 for field_name, field in item._fields.iteritems():
                     v = item._data.get(field_name, None)
                     if isinstance(v, (DBRef)):
-                        reference_map.setdefault(field.document_type, []).append(v.id)
+                        reference_map.setdefault(field.document_type, set()).add(v.id)
                     elif isinstance(v, (dict, SON)) and '_ref' in v:
-                        reference_map.setdefault(get_document(v['_cls']), []).append(v['_ref'].id)
+                        reference_map.setdefault(get_document(v['_cls']), set()).add(v['_ref'].id)
                     elif isinstance(v, (dict, list, tuple)) and depth <= self.max_depth:
                         field_cls = getattr(getattr(field, 'field', None), 'document_type', None)
                         references = self._find_references(v, depth)
                         for key, refs in references.iteritems():
                             if isinstance(field_cls, (Document, TopLevelDocumentMetaclass)):
                                 key = field_cls
-                            reference_map.setdefault(key, []).extend(refs)
+                            reference_map.setdefault(key, set()).update(refs)
             elif isinstance(item, (DBRef)):
-                reference_map.setdefault(item.collection, []).append(item.id)
+                reference_map.setdefault(item.collection, set()).add(item.id)
             elif isinstance(item, (dict, SON)) and '_ref' in item:
-                reference_map.setdefault(get_document(item['_cls']), []).append(item['_ref'].id)
+                reference_map.setdefault(get_document(item['_cls']), set()).add(item['_ref'].id)
             elif isinstance(item, (dict, list, tuple)) and depth - 1 <= self.max_depth:
                 references = self._find_references(item, depth - 1)
                 for key, refs in references.iteritems():
-                    reference_map.setdefault(key, []).extend(refs)
+                    reference_map.setdefault(key, set()).update(refs)
 
         return reference_map
 
@@ -128,8 +128,8 @@ def _fetch_objects(self, doc_type=None):
         """
         object_map = {}
         for collection, dbrefs in self.reference_map.iteritems():
-            keys = object_map.keys()
-            refs = list(set([dbref for dbref in dbrefs if unicode(dbref).encode('utf-8') not in keys]))
+            refs = [dbref for dbref in dbrefs
+                    if unicode(dbref).encode('utf-8') not in object_map]
             if hasattr(collection, 'objects'):  # We have a document class for the refs
                 references = collection.objects.in_bulk(refs)
                 for key, doc in references.iteritems():


Test patch:
diff --git a/tests/document/instance.py b/tests/document/instance.py
--- a/tests/document/instance.py
+++ b/tests/document/instance.py
@@ -10,6 +10,7 @@
 
 from datetime import datetime
 from bson import DBRef, ObjectId
+from tests import fixtures
 from tests.fixtures import (PickleEmbedded, PickleTest, PickleSignalsTest,
                             PickleDyanmicEmbedded, PickleDynamicTest)
 
@@ -2085,6 +2086,29 @@ def test_picklable(self):
         self.assertEqual(pickle_doc.string, "Two")
         self.assertEqual(pickle_doc.lists, ["1", "2", "3"])
 
+    def test_regular_document_pickle(self):
+
+        pickle_doc = PickleTest(number=1, string="One", lists=['1', '2'])
+        pickled_doc = pickle.dumps(pickle_doc)  # make sure pickling works even before the doc is saved
+        pickle_doc.save()
+
+        pickled_doc = pickle.dumps(pickle_doc)
+
+        # Test that when a document's definition changes the new
+        # definition is used
+        fixtures.PickleTest = fixtures.NewDocumentPickleTest
+
+        resurrected = pickle.loads(pickled_doc)
+        self.assertEqual(resurrected.__class__,
+                         fixtures.NewDocumentPickleTest)
+        self.assertEqual(resurrected._fields_ordered,
+                         fixtures.NewDocumentPickleTest._fields_ordered)
+        self.assertNotEqual(resurrected._fields_ordered,
+                            pickle_doc._fields_ordered)
+
+        # The local PickleTest is still a ref to the original
+        fixtures.PickleTest = PickleTest
+
     def test_dynamic_document_pickle(self):
 
         pickle_doc = PickleDynamicTest(
diff --git a/tests/fixtures.py b/tests/fixtures.py
--- a/tests/fixtures.py
+++ b/tests/fixtures.py
@@ -17,6 +17,15 @@ class PickleTest(Document):
     photo = FileField()
 
 
+class NewDocumentPickleTest(Document):
+    number = IntField()
+    string = StringField(choices=(('One', '1'), ('Two', '2')))
+    embedded = EmbeddedDocumentField(PickleEmbedded)
+    lists = ListField(StringField())
+    photo = FileField()
+    new_field = StringField()
+
+
 class PickleDyanmicEmbedded(DynamicEmbeddedDocument):
     date = DateTimeField(default=datetime.now)