Check keys in Collection.save PYTHON-245.

behackett · behackett · commit fad1c38f1bb0 · 2011-11-03T18:28:59.000-07:00
diff --git a/pymongo/_cmessagemodule.c b/pymongo/_cmessagemodule.c
@@ -222,18 +222,19 @@ static PyObject* _cbson_update_message(PyObject* self, PyObject* args) {
     unsigned char multi;
     unsigned char upsert;
     unsigned char safe;
+    unsigned char check_keys;
     PyObject* last_error_args;
     int options;
     buffer_t buffer;
     int length_location, message_length;
     PyObject* result;
 
-    if (!PyArg_ParseTuple(args, "et#bbOObO",
+    if (!PyArg_ParseTuple(args, "et#bbOObbO",
                           "utf-8",
                           &collection_name,
                           &collection_name_length,
                           &upsert, &multi, &spec, &doc, &safe,
-                          &last_error_args)) {
+                          &check_keys, &last_error_args)) {
         return NULL;
     }
 
@@ -282,7 +283,7 @@ static PyObject* _cbson_update_message(PyObject* self, PyObject* args) {
     max_size = buffer_get_position(buffer) - before;
 
     before = buffer_get_position(buffer);
-    if (!write_dict(buffer, doc, 0, 1)) {
+    if (!write_dict(buffer, doc, check_keys, 1)) {
         buffer_free(buffer);
         PyMem_Free(collection_name);
         return NULL;
diff --git a/pymongo/collection.py b/pymongo/collection.py
@@ -218,7 +218,7 @@ def save(self, to_save, manipulate=True, safe=False, **kwargs):
             return self.insert(to_save, manipulate, safe, **kwargs)
         else:
             self.update({"_id": to_save["_id"]}, to_save, True,
-                        manipulate, safe, **kwargs)
+                        manipulate, safe, _check_keys=True, **kwargs)
             return to_save.get("_id", None)
 
     def insert(self, doc_or_docs, manipulate=True,
@@ -298,7 +298,7 @@ def insert(self, doc_or_docs, manipulate=True,
         return return_one and ids[0] or ids
 
     def update(self, spec, document, upsert=False, manipulate=False,
-               safe=False, multi=False, **kwargs):
+               safe=False, multi=False, _check_keys=False, **kwargs):
         """Update a document(s) in this collection.
 
         Raises :class:`TypeError` if either `spec` or `document` is
@@ -385,9 +385,12 @@ def update(self, spec, document, upsert=False, manipulate=False,
             if not kwargs:
                 kwargs.update(self.get_lasterror_options())
 
+        # _check_keys is used by save() so we don't upsert pre-existing
+        # documents after adding an invalid key like 'a.b'. It can't really
+        # be used for any other update operations.
         return self.__database.connection._send_message(
             message.update(self.__full_name, upsert, multi,
-                           spec, document, safe, kwargs), safe)
+                           spec, document, safe, _check_keys, kwargs), safe)
 
     def drop(self):
         """Alias for :meth:`~pymongo.database.Database.drop_collection`.
diff --git a/pymongo/message.py b/pymongo/message.py
@@ -88,7 +88,8 @@ def insert(collection_name, docs, check_keys,
     insert = _cmessage._insert_message
 
 
-def update(collection_name, upsert, multi, spec, doc, safe, last_error_args):
+def update(collection_name, upsert, multi,
+           spec, doc, safe, check_keys, last_error_args):
     """Get an **update** message.
     """
     options = 0
@@ -101,7 +102,7 @@ def update(collection_name, upsert, multi, spec, doc, safe, last_error_args):
     data += bson._make_c_string(collection_name)
     data += struct.pack("<i", options)
     data += bson.BSON.encode(spec)
-    encoded = bson.BSON.encode(doc)
+    encoded = bson.BSON.encode(doc, check_keys)
     data += encoded
     if safe:
         (_, update_message) = __pack_message(2001, data)
diff --git a/test/test_collection.py b/test/test_collection.py
@@ -504,6 +504,13 @@ def test_save(self):
         self.assertEqual(self.db.test.find_one()["_id"], id)
         self.assert_(isinstance(id, ObjectId))
 
+    def test_save_with_invalid_key(self):
+        self.db.drop_collection("test")
+        self.assert_(self.db.test.insert({"hello": "world"}))
+        doc = self.db.test.find_one()
+        doc['a.b'] = 'c'
+        self.assertRaises(InvalidDocument, self.db.test.save, doc)
+
     def test_unique_index(self):
         db = self.db
 
