change session save api (foundation) to detect before_save changes, and don't break on inserts. tests provided

Author: CastixGitHub

ming/odm/mapper.py

@@ -3,7 +3,7 @@
 from copy import copy, deepcopy
 
 from ming.base import Object, NoDefault
-from ming.utils import wordwrap, doc_to_set
+from ming.utils import wordwrap
 
 from .base import ObjectState, state, _with_hooks
 from .property import FieldProperty
@@ -84,15 +84,8 @@ def insert(self, obj, state, session, **kwargs):
 
     @_with_hooks('update')
     def update(self, obj, state, session, **kwargs):
-        fields = state.options.get('fields', None)
-        if fields is None:
-            # here we do a symmetric difference to see what fields did change
-            fields = tuple(set((k for k, v in
-                                doc_to_set(state.original_document)
-                                ^ doc_to_set(state.document))))
-
         doc = self.collection(state.document, skip_from_bson=True)
-        ret = session.impl.save(doc, *fields, validate=False)
+        ret = session.impl.save(doc, validate=False, state=state)
         state.status = state.clean
         return ret
 

ming/session.py

@@ -9,7 +9,7 @@
 import six
 
 from .base import Cursor, Object
-from .utils import fixup_index, fix_write_concern
+from .utils import fixup_index, fix_write_concern, doc_to_set
 from . import exc
 
 log = logging.getLogger(__name__)
@@ -139,7 +139,8 @@ def find_and_modify(self, cls, query=None, sort=None, new=False, **kw):
 
     def _prep_save(self, doc, validate):
         hook = doc.m.before_save
-        if hook: hook(doc)
+        if hook:
+            hook(doc)
         if validate:
             if doc.m.schema is None:
                 data = dict(doc)
@@ -151,9 +152,15 @@ def _prep_save(self, doc, validate):
         return data
 
     @annotate_doc_failure
-    def save(self, doc, *args, **kwargs):
+    def save(self, doc, state=None, **kwargs):
+        # args was meant to be the list of changed fields
+        # but actually we ended up checking the differences here
         data = self._prep_save(doc, kwargs.pop('validate', True))
-        if args:
+        if state is not None and state.original_document:
+            if state is not None:
+                args = tuple(set((k for k, v in
+                                  doc_to_set(state.original_document)
+                                  ^ doc_to_set(data))))
             values = dict((arg, data[arg]) for arg in args)
             result = self._impl(doc).update(
                 dict(_id=doc._id), {'$set': values}, **fix_write_concern(kwargs))

ming/tests/odm/test_declarative.py

@@ -1,6 +1,7 @@
 import sys
 from collections import defaultdict
 from unittest import TestCase, SkipTest
+from unittest.mock import MagicMock
 
 from ming import schema as S
 from ming import create_datastore
@@ -811,3 +812,60 @@ def test_hook_base(self):
                          [
                              {'_id': doc._id, 'a': doc.a}
                          ])
+
+
+class TestReplacingSession(TestCase):
+
+    def setUp(self):
+        Mapper._mapper_by_classname.clear()
+        self.datastore = create_datastore('mim:///test_db')
+        self.session = ODMSession(bind=self.datastore)
+        class Basic(MappedClass):
+            class __mongometa__:
+                name = 'hook'
+                session = self.session
+            _id = FieldProperty(S.ObjectId)
+            a = FieldProperty(int)
+        Mapper.compile_all()
+        self.Basic = Basic
+        self.session.remove(self.Basic)
+
+    def test_hook_base(self):
+        assert id(self.Basic.query.session) == id(self.session)
+        session2 = MagicMock()
+        new_session = ODMSession(bind=session2)
+        Mapper.replace_session(new_session)
+        assert id(self.Basic.query.session) == id(new_session)
+        assert id(self.session) != id(new_session)
+
+class TestBeforeSave(TestCase):
+
+    def setUp(self):
+        Mapper._mapper_by_classname.clear()
+        self.datastore = create_datastore('mim:///test_db')
+        self.session = ODMSession(bind=self.datastore)
+        class Basic(MappedClass):
+            class __mongometa__:
+                name = 'hook'
+                session = self.session
+                def before_save(instance):
+                    instance.a = 9
+
+            _id = FieldProperty(S.ObjectId)
+            a = FieldProperty(int)
+        Mapper.compile_all()
+        self.Basic = Basic
+        self.session.remove(self.Basic)
+
+    def test_hook_base(self):
+        doc = self.Basic()
+        doc.a = 5
+        self.session.flush()  # first insert
+        self.session.close()
+        doc = self.Basic.query.get(doc._id)
+        assert doc.a == 9, doc.a
+        doc.a = 6 
+        self.session.flush()  # then save
+        self.session.close()
+        doc = self.Basic.query.get(doc._id)
+        assert doc.a == 9, doc.a

ming/tests/test_session.py

@@ -81,9 +81,9 @@ def test_base_session(self):
         doc = self.TestDocNoSchema({'_id':5, 'a':5})
         sess.save(doc)
         impl.save.assert_called_with(dict(_id=5, a=5))
-        doc = self.TestDocNoSchema({'_id':5, 'a':5})
-        sess.save(doc, 'a')
-        impl.update.assert_called_with(dict(_id=5), {'$set':dict(a=5)})
+        # doc = self.TestDocNoSchema({'_id':5, 'a':5})
+        # sess.save(doc, 'a')
+        # impl.update.assert_called_with(dict(_id=5), {'$set':dict(a=5)})
         doc = self.TestDocNoSchema({'_id':5, 'a':5})
         impl.insert.return_value = bson.ObjectId()
         sess.insert(doc)