def _save_update(self, doc, save_condition, write_concern): """Update an existing document. Helper method, should only be used inside save(). """ doc_alias = None if doc.get('properties', {}).get('db_alias', None): doc_alias = doc['properties'].pop('db_alias') if doc.get('properties', {}).get('conn_settings', {}): connection_properties = doc['properties'].pop('conn_settings') self.verify_connection_setting(doc_alias, connection_properties) if doc_alias: alias_db = get_db(doc_alias) collection = alias_db.get_collection(self._meta['collection']) else: collection = self._get_collection() object_id = doc["_id"] created = False select_dict = {} if save_condition is not None: select_dict = transform.query(self.__class__, **save_condition) select_dict["_id"] = object_id # Need to add shard key to query, or you get an error shard_key = self._meta.get("shard_key", tuple()) for k in shard_key: path = self._lookup_field(k.split(".")) actual_key = [p.db_field for p in path] val = doc for ak in actual_key: val = val[ak] select_dict[".".join(actual_key)] = val update_doc = self._get_update_doc() if update_doc: upsert = save_condition is None with set_write_concern(collection, write_concern) as wc_collection: last_error = wc_collection.update_one(select_dict, update_doc, upsert=upsert).raw_result if not upsert and last_error["n"] == 0: raise SaveConditionError( "Race condition preventing document update detected") if last_error is not None: updated_existing = last_error.get("updatedExisting") if updated_existing is False: created = True # !!! This is bad, means we accidentally created a new, # potentially corrupted document. See # https://github.com/MongoEngine/mongoengine/issues/564 return object_id, created
def store_data(data, name, overwrite=False): """Create a new BaseStudy or overwrite an existing one with the given data.""" existing = get_study(name) if existing and overwrite: archive = existing elif existing: raise SaveConditionError(f"Study '{name}' already exists, but 'overwrite' set to False.") else: archive = BaseStudy(name=name) archive.data = data archive.save() return get_data(name)
def _save_update(self, doc, save_condition, write_concern): """Update an existing document. Helper method, should only be used inside save(). """ collection = self._get_collection() object_id = doc['_id'] created = False select_dict = {} if save_condition is not None: select_dict = transform.query(self.__class__, **save_condition) select_dict['_id'] = object_id # Need to add shard key to query, or you get an error shard_key = self._meta.get('shard_key', tuple()) for k in shard_key: path = self._lookup_field(k.split('.')) actual_key = [p.db_field for p in path] val = doc for ak in actual_key: val = val[ak] select_dict['.'.join(actual_key)] = val updates, removals = self._delta() update_query = {} if updates: update_query['$set'] = updates if removals: update_query['$unset'] = removals if updates or removals: upsert = save_condition is None last_error = collection.update(select_dict, update_query, upsert=upsert, **write_concern) if not upsert and last_error['n'] == 0: raise SaveConditionError('Race condition preventing' ' document update detected') if last_error is not None: updated_existing = last_error.get('updatedExisting') if updated_existing is False: created = True # !!! This is bad, means we accidentally created a new, # potentially corrupted document. See # https://github.com/MongoEngine/mongoengine/issues/564 return object_id, created
def _save_update(self, doc, save_condition, write_concern): """Update an existing document. Helper method, should only be used inside save(). """ collection = self._get_collection() object_id = doc["_id"] created = False select_dict = {} if save_condition is not None: select_dict = transform.query(self.__class__, **save_condition) select_dict["_id"] = object_id select_dict = self._integrate_shard_key(doc, select_dict) update_doc = self._get_update_doc() if update_doc: upsert = save_condition is None with set_write_concern(collection, write_concern) as wc_collection: last_error = wc_collection.update_one( select_dict, update_doc, upsert=upsert, session=self._get_local_session(), ).raw_result if not upsert and last_error["n"] == 0: raise SaveConditionError( "Race condition preventing document update detected") if last_error is not None: updated_existing = last_error.get("updatedExisting") if updated_existing is False: created = True # !!! This is bad, means we accidentally created a new, # potentially corrupted document. See # https://github.com/MongoEngine/mongoengine/issues/564 return object_id, created
def save(self, force_insert=False, validate=True, clean=True, write_concern=None, cascade=None, cascade_kwargs=None, _refs=None, save_condition=None, signal_kwargs=None, **kwargs): """Save the :class:`~mongoengine.Document` to the database. If the document already exists, it will be updated, otherwise it will be created. :param force_insert: only try to create a new document, don't allow updates of existing documents :param validate: validates the document; set to ``False`` to skip. :param clean: call the document clean method, requires `validate` to be True. :param write_concern: Extra keyword arguments are passed down to :meth:`~pymongo.collection.Collection.save` OR :meth:`~pymongo.collection.Collection.insert` which will be used as options for the resultant ``getLastError`` command. For example, ``save(..., write_concern={w: 2, fsync: True}, ...)`` will wait until at least two servers have recorded the write and will force an fsync on the primary server. :param cascade: Sets the flag for cascading saves. You can set a default by setting "cascade" in the document __meta__ :param cascade_kwargs: (optional) kwargs dictionary to be passed throw to cascading saves. Implies ``cascade=True``. :param _refs: A list of processed references used in cascading saves :param save_condition: only perform save if matching record in db satisfies condition(s) (e.g. version number). Raises :class:`OperationError` if the conditions are not satisfied :parm signal_kwargs: (optional) kwargs dictionary to be passed to the signal calls. .. versionchanged:: 0.5 In existing documents it only saves changed fields using set / unset. Saves are cascaded and any :class:`~bson.dbref.DBRef` objects that have changes are saved as well. .. versionchanged:: 0.6 Added cascading saves .. versionchanged:: 0.8 Cascade saves are optional and default to False. If you want fine grain control then you can turn off using document meta['cascade'] = True. Also you can pass different kwargs to the cascade save using cascade_kwargs which overwrites the existing kwargs with custom values. .. versionchanged:: 0.8.5 Optional save_condition that only overwrites existing documents if the condition is satisfied in the current db record. .. versionchanged:: 0.10 :class:`OperationError` exception raised if save_condition fails. .. versionchanged:: 0.10.1 :class: save_condition failure now raises a `SaveConditionError` .. versionchanged:: 0.10.7 Add signal_kwargs argument """ signal_kwargs = signal_kwargs or {} signals.pre_save.send(self.__class__, document=self, **signal_kwargs) if validate: self.validate(clean=clean) if write_concern is None: write_concern = {'w': 1} doc = self.to_mongo() created = ('_id' not in doc or self._created or force_insert) signals.pre_save_post_validation.send(self.__class__, document=self, created=created, **signal_kwargs) try: collection = self._get_collection() if self._meta.get('auto_create_index', True): self.ensure_indexes() if created: if force_insert: object_id = collection.insert(doc, **write_concern) else: object_id = collection.save(doc, **write_concern) # In PyMongo 3.0, the save() call calls internally the _update() call # but they forget to return the _id value passed back, therefore getting it back here # Correct behaviour in 2.X and in 3.0.1+ versions if not object_id and pymongo.version_tuple == (3, 0): pk_as_mongo_obj = self._fields.get(self._meta['id_field']).to_mongo(self.pk) object_id = ( self._qs.filter(pk=pk_as_mongo_obj).first() and self._qs.filter(pk=pk_as_mongo_obj).first().pk ) # TODO doesn't this make 2 queries? else: object_id = doc['_id'] updates, removals = self._delta() # Need to add shard key to query, or you get an error if save_condition is not None: select_dict = transform.query(self.__class__, **save_condition) else: select_dict = {} select_dict['_id'] = object_id shard_key = self._meta.get('shard_key', tuple()) for k in shard_key: path = self._lookup_field(k.split('.')) actual_key = [p.db_field for p in path] val = doc for ak in actual_key: val = val[ak] select_dict['.'.join(actual_key)] = val def is_new_object(last_error): if last_error is not None: updated = last_error.get('updatedExisting') if updated is not None: return not updated return created update_query = {} if updates: update_query['$set'] = updates if removals: update_query['$unset'] = removals if updates or removals: upsert = save_condition is None last_error = collection.update(select_dict, update_query, upsert=upsert, **write_concern) if not upsert and last_error['n'] == 0: raise SaveConditionError('Race condition preventing' ' document update detected') created = is_new_object(last_error) if cascade is None: cascade = self._meta.get( 'cascade', False) or cascade_kwargs is not None if cascade: kwargs = { 'force_insert': force_insert, 'validate': validate, 'write_concern': write_concern, 'cascade': cascade } if cascade_kwargs: # Allow granular control over cascades kwargs.update(cascade_kwargs) kwargs['_refs'] = _refs self.cascade_save(**kwargs) except pymongo.errors.DuplicateKeyError as err: message = 'Tried to save duplicate unique keys (%s)' raise NotUniqueError(message % six.text_type(err)) except pymongo.errors.OperationFailure as err: message = 'Could not save document (%s)' if re.match('^E1100[01] duplicate key', six.text_type(err)): # E11000 - duplicate key error index # E11001 - duplicate key on update message = 'Tried to save duplicate unique keys (%s)' raise NotUniqueError(message % six.text_type(err)) raise OperationError(message % six.text_type(err)) id_field = self._meta['id_field'] if created or id_field not in self._meta.get('shard_key', []): self[id_field] = self._fields[id_field].to_python(object_id) signals.post_save.send(self.__class__, document=self, created=created, **signal_kwargs) self._clear_changed_fields() self._created = False return self