base.py 70 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912
  1. from __future__ import absolute_import
  2. import copy
  3. import itertools
  4. import pprint
  5. import re
  6. import warnings
  7. from bson import SON, json_util
  8. from bson.code import Code
  9. import pymongo
  10. import pymongo.errors
  11. from pymongo.common import validate_read_preference
  12. import six
  13. from mongoengine import signals
  14. from mongoengine.base import get_document
  15. from mongoengine.common import _import_class
  16. from mongoengine.connection import get_db
  17. from mongoengine.context_managers import set_write_concern, switch_db
  18. from mongoengine.errors import (InvalidQueryError, LookUpError,
  19. NotUniqueError, OperationError)
  20. from mongoengine.python_support import IS_PYMONGO_3
  21. from mongoengine.queryset import transform
  22. from mongoengine.queryset.field_list import QueryFieldList
  23. from mongoengine.queryset.visitor import Q, QNode
  24. if IS_PYMONGO_3:
  25. from pymongo.collection import ReturnDocument
  26. __all__ = ('BaseQuerySet', 'DO_NOTHING', 'NULLIFY', 'CASCADE', 'DENY', 'PULL')
  27. # Delete rules
  28. DO_NOTHING = 0
  29. NULLIFY = 1
  30. CASCADE = 2
  31. DENY = 3
  32. PULL = 4
  33. class BaseQuerySet(object):
  34. """A set of results returned from a query. Wraps a MongoDB cursor,
  35. providing :class:`~mongoengine.Document` objects as the results.
  36. """
  37. __dereference = False
  38. _auto_dereference = True
  39. def __init__(self, document, collection):
  40. self._document = document
  41. self._collection_obj = collection
  42. self._mongo_query = None
  43. self._query_obj = Q()
  44. self._initial_query = {}
  45. self._where_clause = None
  46. self._loaded_fields = QueryFieldList()
  47. self._ordering = None
  48. self._snapshot = False
  49. self._timeout = True
  50. self._class_check = True
  51. self._slave_okay = False
  52. self._read_preference = None
  53. self._iter = False
  54. self._scalar = []
  55. self._none = False
  56. self._as_pymongo = False
  57. self._search_text = None
  58. # If inheritance is allowed, only return instances and instances of
  59. # subclasses of the class being used
  60. if document._meta.get('allow_inheritance') is True:
  61. if len(self._document._subclasses) == 1:
  62. self._initial_query = {'_cls': self._document._subclasses[0]}
  63. else:
  64. self._initial_query = {
  65. '_cls': {'$in': self._document._subclasses}}
  66. self._loaded_fields = QueryFieldList(always_include=['_cls'])
  67. self._cursor_obj = None
  68. self._limit = None
  69. self._skip = None
  70. self._hint = -1 # Using -1 as None is a valid value for hint
  71. self._batch_size = None
  72. self.only_fields = []
  73. self._max_time_ms = None
  74. self._comment = None
  75. def __call__(self, q_obj=None, class_check=True, read_preference=None,
  76. **query):
  77. """Filter the selected documents by calling the
  78. :class:`~mongoengine.queryset.QuerySet` with a query.
  79. :param q_obj: a :class:`~mongoengine.queryset.Q` object to be used in
  80. the query; the :class:`~mongoengine.queryset.QuerySet` is filtered
  81. multiple times with different :class:`~mongoengine.queryset.Q`
  82. objects, only the last one will be used
  83. :param class_check: If set to False bypass class name check when
  84. querying collection
  85. :param read_preference: if set, overrides connection-level
  86. read_preference from `ReplicaSetConnection`.
  87. :param query: Django-style query keyword arguments
  88. """
  89. query = Q(**query)
  90. if q_obj:
  91. # make sure proper query object is passed
  92. if not isinstance(q_obj, QNode):
  93. msg = ('Not a query object: %s. '
  94. 'Did you intend to use key=value?' % q_obj)
  95. raise InvalidQueryError(msg)
  96. query &= q_obj
  97. if read_preference is None:
  98. queryset = self.clone()
  99. else:
  100. # Use the clone provided when setting read_preference
  101. queryset = self.read_preference(read_preference)
  102. queryset._query_obj &= query
  103. queryset._mongo_query = None
  104. queryset._cursor_obj = None
  105. queryset._class_check = class_check
  106. return queryset
  107. def __getstate__(self):
  108. """
  109. Need for pickling queryset
  110. See https://github.com/MongoEngine/mongoengine/issues/442
  111. """
  112. obj_dict = self.__dict__.copy()
  113. # don't picke collection, instead pickle collection params
  114. obj_dict.pop('_collection_obj')
  115. # don't pickle cursor
  116. obj_dict['_cursor_obj'] = None
  117. return obj_dict
  118. def __setstate__(self, obj_dict):
  119. """
  120. Need for pickling queryset
  121. See https://github.com/MongoEngine/mongoengine/issues/442
  122. """
  123. obj_dict['_collection_obj'] = obj_dict['_document']._get_collection()
  124. # update attributes
  125. self.__dict__.update(obj_dict)
  126. # forse load cursor
  127. # self._cursor
  128. def __getitem__(self, key):
  129. """Return a document instance corresponding to a given index if
  130. the key is an integer. If the key is a slice, translate its
  131. bounds into a skip and a limit, and return a cloned queryset
  132. with that skip/limit applied. For example:
  133. >>> User.objects[0]
  134. <User: User object>
  135. >>> User.objects[1:3]
  136. [<User: User object>, <User: User object>]
  137. """
  138. queryset = self.clone()
  139. # Handle a slice
  140. if isinstance(key, slice):
  141. queryset._cursor_obj = queryset._cursor[key]
  142. queryset._skip, queryset._limit = key.start, key.stop
  143. if key.start and key.stop:
  144. queryset._limit = key.stop - key.start
  145. # Allow further QuerySet modifications to be performed
  146. return queryset
  147. # Handle an index
  148. elif isinstance(key, int):
  149. if queryset._scalar:
  150. return queryset._get_scalar(
  151. queryset._document._from_son(
  152. queryset._cursor[key],
  153. _auto_dereference=self._auto_dereference,
  154. only_fields=self.only_fields
  155. )
  156. )
  157. if queryset._as_pymongo:
  158. return queryset._get_as_pymongo(queryset._cursor[key])
  159. return queryset._document._from_son(
  160. queryset._cursor[key],
  161. _auto_dereference=self._auto_dereference,
  162. only_fields=self.only_fields
  163. )
  164. raise AttributeError('Provide a slice or an integer index')
  165. def __iter__(self):
  166. raise NotImplementedError
  167. def _has_data(self):
  168. """Return True if cursor has any data."""
  169. queryset = self.order_by()
  170. return False if queryset.first() is None else True
  171. def __bool__(self):
  172. """Avoid to open all records in an if stmt in Py3."""
  173. return self._has_data()
  174. __nonzero__ = __bool__ # For Py2 support
  175. # Core functions
  176. def all(self):
  177. """Returns all documents."""
  178. return self.__call__()
  179. def filter(self, *q_objs, **query):
  180. """An alias of :meth:`~mongoengine.queryset.QuerySet.__call__`
  181. """
  182. return self.__call__(*q_objs, **query)
  183. def search_text(self, text, language=None):
  184. """
  185. Start a text search, using text indexes.
  186. Require: MongoDB server version 2.6+.
  187. :param language: The language that determines the list of stop words
  188. for the search and the rules for the stemmer and tokenizer.
  189. If not specified, the search uses the default language of the index.
  190. For supported languages, see
  191. `Text Search Languages <http://docs.mongodb.org/manual/reference/text-search-languages/#text-search-languages>`.
  192. """
  193. queryset = self.clone()
  194. if queryset._search_text:
  195. raise OperationError(
  196. 'It is not possible to use search_text two times.')
  197. query_kwargs = SON({'$search': text})
  198. if language:
  199. query_kwargs['$language'] = language
  200. queryset._query_obj &= Q(__raw__={'$text': query_kwargs})
  201. queryset._mongo_query = None
  202. queryset._cursor_obj = None
  203. queryset._search_text = text
  204. return queryset
  205. def get(self, *q_objs, **query):
  206. """Retrieve the the matching object raising
  207. :class:`~mongoengine.queryset.MultipleObjectsReturned` or
  208. `DocumentName.MultipleObjectsReturned` exception if multiple results
  209. and :class:`~mongoengine.queryset.DoesNotExist` or
  210. `DocumentName.DoesNotExist` if no results are found.
  211. .. versionadded:: 0.3
  212. """
  213. queryset = self.clone()
  214. queryset = queryset.order_by().limit(2)
  215. queryset = queryset.filter(*q_objs, **query)
  216. try:
  217. result = six.next(queryset)
  218. except StopIteration:
  219. msg = ('%s matching query does not exist.'
  220. % queryset._document._class_name)
  221. raise queryset._document.DoesNotExist(msg)
  222. try:
  223. six.next(queryset)
  224. except StopIteration:
  225. return result
  226. # If we were able to retrieve the 2nd doc, rewind the cursor and
  227. # raise the MultipleObjectsReturned exception.
  228. queryset.rewind()
  229. message = u'%d items returned, instead of 1' % queryset.count()
  230. raise queryset._document.MultipleObjectsReturned(message)
  231. def create(self, **kwargs):
  232. """Create new object. Returns the saved object instance.
  233. .. versionadded:: 0.4
  234. """
  235. return self._document(**kwargs).save(force_insert=True)
  236. def first(self):
  237. """Retrieve the first object matching the query."""
  238. queryset = self.clone()
  239. try:
  240. result = queryset[0]
  241. except IndexError:
  242. result = None
  243. return result
  244. def insert(self, doc_or_docs, load_bulk=True,
  245. write_concern=None, signal_kwargs=None):
  246. """bulk insert documents
  247. :param doc_or_docs: a document or list of documents to be inserted
  248. :param load_bulk (optional): If True returns the list of document
  249. instances
  250. :param write_concern: Extra keyword arguments are passed down to
  251. :meth:`~pymongo.collection.Collection.insert`
  252. which will be used as options for the resultant
  253. ``getLastError`` command. For example,
  254. ``insert(..., {w: 2, fsync: True})`` will wait until at least
  255. two servers have recorded the write and will force an fsync on
  256. each server being written to.
  257. :parm signal_kwargs: (optional) kwargs dictionary to be passed to
  258. the signal calls.
  259. By default returns document instances, set ``load_bulk`` to False to
  260. return just ``ObjectIds``
  261. .. versionadded:: 0.5
  262. .. versionchanged:: 0.10.7
  263. Add signal_kwargs argument
  264. """
  265. Document = _import_class('Document')
  266. if write_concern is None:
  267. write_concern = {}
  268. docs = doc_or_docs
  269. return_one = False
  270. if isinstance(docs, Document) or issubclass(docs.__class__, Document):
  271. return_one = True
  272. docs = [docs]
  273. for doc in docs:
  274. if not isinstance(doc, self._document):
  275. msg = ("Some documents inserted aren't instances of %s"
  276. % str(self._document))
  277. raise OperationError(msg)
  278. if doc.pk and not doc._created:
  279. msg = 'Some documents have ObjectIds use doc.update() instead'
  280. raise OperationError(msg)
  281. signal_kwargs = signal_kwargs or {}
  282. signals.pre_bulk_insert.send(self._document,
  283. documents=docs, **signal_kwargs)
  284. raw = [doc.to_mongo() for doc in docs]
  285. with set_write_concern(self._collection, write_concern) as collection:
  286. insert_func = collection.insert_many
  287. if return_one:
  288. raw = raw[0]
  289. insert_func = collection.insert_one
  290. try:
  291. inserted_result = insert_func(raw)
  292. ids = [inserted_result.inserted_id] if return_one else inserted_result.inserted_ids
  293. except pymongo.errors.DuplicateKeyError as err:
  294. message = 'Could not save document (%s)'
  295. raise NotUniqueError(message % six.text_type(err))
  296. except pymongo.errors.BulkWriteError as err:
  297. # inserting documents that already have an _id field will
  298. # give huge performance debt or raise
  299. message = u'Document must not have _id value before bulk write (%s)'
  300. raise NotUniqueError(message % six.text_type(err))
  301. except pymongo.errors.OperationFailure as err:
  302. message = 'Could not save document (%s)'
  303. if re.match('^E1100[01] duplicate key', six.text_type(err)):
  304. # E11000 - duplicate key error index
  305. # E11001 - duplicate key on update
  306. message = u'Tried to save duplicate unique keys (%s)'
  307. raise NotUniqueError(message % six.text_type(err))
  308. raise OperationError(message % six.text_type(err))
  309. # Apply inserted_ids to documents
  310. for doc, doc_id in zip(docs, ids):
  311. doc.pk = doc_id
  312. if not load_bulk:
  313. signals.post_bulk_insert.send(
  314. self._document, documents=docs, loaded=False, **signal_kwargs)
  315. return ids[0] if return_one else ids
  316. documents = self.in_bulk(ids)
  317. results = [documents.get(obj_id) for obj_id in ids]
  318. signals.post_bulk_insert.send(
  319. self._document, documents=results, loaded=True, **signal_kwargs)
  320. return results[0] if return_one else results
  321. def count(self, with_limit_and_skip=False):
  322. """Count the selected elements in the query.
  323. :param with_limit_and_skip (optional): take any :meth:`limit` or
  324. :meth:`skip` that has been applied to this cursor into account when
  325. getting the count
  326. """
  327. if self._limit == 0 and with_limit_and_skip is False or self._none:
  328. return 0
  329. return self._cursor.count(with_limit_and_skip=with_limit_and_skip)
  330. def delete(self, write_concern=None, _from_doc_delete=False,
  331. cascade_refs=None):
  332. """Delete the documents matched by the query.
  333. :param write_concern: Extra keyword arguments are passed down which
  334. will be used as options for the resultant
  335. ``getLastError`` command. For example,
  336. ``save(..., write_concern={w: 2, fsync: True}, ...)`` will
  337. wait until at least two servers have recorded the write and
  338. will force an fsync on the primary server.
  339. :param _from_doc_delete: True when called from document delete therefore
  340. signals will have been triggered so don't loop.
  341. :returns number of deleted documents
  342. """
  343. queryset = self.clone()
  344. doc = queryset._document
  345. if write_concern is None:
  346. write_concern = {}
  347. # Handle deletes where skips or limits have been applied or
  348. # there is an untriggered delete signal
  349. has_delete_signal = signals.signals_available and (
  350. signals.pre_delete.has_receivers_for(doc) or
  351. signals.post_delete.has_receivers_for(doc)
  352. )
  353. call_document_delete = (queryset._skip or queryset._limit or
  354. has_delete_signal) and not _from_doc_delete
  355. if call_document_delete:
  356. cnt = 0
  357. for doc in queryset:
  358. doc.delete(**write_concern)
  359. cnt += 1
  360. return cnt
  361. delete_rules = doc._meta.get('delete_rules') or {}
  362. delete_rules = list(delete_rules.items())
  363. # Check for DENY rules before actually deleting/nullifying any other
  364. # references
  365. for rule_entry, rule in delete_rules:
  366. document_cls, field_name = rule_entry
  367. if document_cls._meta.get('abstract'):
  368. continue
  369. if rule == DENY:
  370. refs = document_cls.objects(**{field_name + '__in': self})
  371. if refs.limit(1).count() > 0:
  372. raise OperationError(
  373. 'Could not delete document (%s.%s refers to it)'
  374. % (document_cls.__name__, field_name)
  375. )
  376. # Check all the other rules
  377. for rule_entry, rule in delete_rules:
  378. document_cls, field_name = rule_entry
  379. if document_cls._meta.get('abstract'):
  380. continue
  381. if rule == CASCADE:
  382. cascade_refs = set() if cascade_refs is None else cascade_refs
  383. # Handle recursive reference
  384. if doc._collection == document_cls._collection:
  385. for ref in queryset:
  386. cascade_refs.add(ref.id)
  387. refs = document_cls.objects(**{field_name + '__in': self,
  388. 'pk__nin': cascade_refs})
  389. if refs.count() > 0:
  390. refs.delete(write_concern=write_concern,
  391. cascade_refs=cascade_refs)
  392. elif rule == NULLIFY:
  393. document_cls.objects(**{field_name + '__in': self}).update(
  394. write_concern=write_concern,
  395. **{'unset__%s' % field_name: 1})
  396. elif rule == PULL:
  397. document_cls.objects(**{field_name + '__in': self}).update(
  398. write_concern=write_concern,
  399. **{'pull_all__%s' % field_name: self})
  400. result = queryset._collection.remove(queryset._query, **write_concern)
  401. if result:
  402. return result.get('n')
  403. def update(self, upsert=False, multi=True, write_concern=None,
  404. full_result=False, **update):
  405. """Perform an atomic update on the fields matched by the query.
  406. :param upsert: insert if document doesn't exist (default ``False``)
  407. :param multi: Update multiple documents.
  408. :param write_concern: Extra keyword arguments are passed down which
  409. will be used as options for the resultant
  410. ``getLastError`` command. For example,
  411. ``save(..., write_concern={w: 2, fsync: True}, ...)`` will
  412. wait until at least two servers have recorded the write and
  413. will force an fsync on the primary server.
  414. :param full_result: Return the full result dictionary rather than just the number
  415. updated, e.g. return
  416. ``{'n': 2, 'nModified': 2, 'ok': 1.0, 'updatedExisting': True}``.
  417. :param update: Django-style update keyword arguments
  418. .. versionadded:: 0.2
  419. """
  420. if not update and not upsert:
  421. raise OperationError('No update parameters, would remove data')
  422. if write_concern is None:
  423. write_concern = {}
  424. queryset = self.clone()
  425. query = queryset._query
  426. update = transform.update(queryset._document, **update)
  427. # If doing an atomic upsert on an inheritable class
  428. # then ensure we add _cls to the update operation
  429. if upsert and '_cls' in query:
  430. if '$set' in update:
  431. update['$set']['_cls'] = queryset._document._class_name
  432. else:
  433. update['$set'] = {'_cls': queryset._document._class_name}
  434. try:
  435. with set_write_concern(queryset._collection, write_concern) as collection:
  436. update_func = collection.update_one
  437. if multi:
  438. update_func = collection.update_many
  439. result = update_func(query, update, upsert=upsert)
  440. if full_result:
  441. return result
  442. elif result.raw_result:
  443. return result.raw_result['n']
  444. except pymongo.errors.DuplicateKeyError as err:
  445. raise NotUniqueError(u'Update failed (%s)' % six.text_type(err))
  446. except pymongo.errors.OperationFailure as err:
  447. if six.text_type(err) == u'multi not coded yet':
  448. message = u'update() method requires MongoDB 1.1.3+'
  449. raise OperationError(message)
  450. raise OperationError(u'Update failed (%s)' % six.text_type(err))
  451. def upsert_one(self, write_concern=None, **update):
  452. """Overwrite or add the first document matched by the query.
  453. :param write_concern: Extra keyword arguments are passed down which
  454. will be used as options for the resultant
  455. ``getLastError`` command. For example,
  456. ``save(..., write_concern={w: 2, fsync: True}, ...)`` will
  457. wait until at least two servers have recorded the write and
  458. will force an fsync on the primary server.
  459. :param update: Django-style update keyword arguments
  460. :returns the new or overwritten document
  461. .. versionadded:: 0.10.2
  462. """
  463. atomic_update = self.update(multi=False, upsert=True,
  464. write_concern=write_concern,
  465. full_result=True, **update)
  466. if atomic_update.raw_result['updatedExisting']:
  467. document = self.get()
  468. else:
  469. document = self._document.objects.with_id(atomic_update.upserted_id)
  470. return document
  471. def update_one(self, upsert=False, write_concern=None, **update):
  472. """Perform an atomic update on the fields of the first document
  473. matched by the query.
  474. :param upsert: insert if document doesn't exist (default ``False``)
  475. :param write_concern: Extra keyword arguments are passed down which
  476. will be used as options for the resultant
  477. ``getLastError`` command. For example,
  478. ``save(..., write_concern={w: 2, fsync: True}, ...)`` will
  479. wait until at least two servers have recorded the write and
  480. will force an fsync on the primary server.
  481. :param update: Django-style update keyword arguments
  482. .. versionadded:: 0.2
  483. """
  484. return self.update(
  485. upsert=upsert, multi=False, write_concern=write_concern, **update)
  486. def modify(self, upsert=False, full_response=False, remove=False, new=False, **update):
  487. """Update and return the updated document.
  488. Returns either the document before or after modification based on `new`
  489. parameter. If no documents match the query and `upsert` is false,
  490. returns ``None``. If upserting and `new` is false, returns ``None``.
  491. If the full_response parameter is ``True``, the return value will be
  492. the entire response object from the server, including the 'ok' and
  493. 'lastErrorObject' fields, rather than just the modified document.
  494. This is useful mainly because the 'lastErrorObject' document holds
  495. information about the command's execution.
  496. :param upsert: insert if document doesn't exist (default ``False``)
  497. :param full_response: return the entire response object from the
  498. server (default ``False``, not available for PyMongo 3+)
  499. :param remove: remove rather than updating (default ``False``)
  500. :param new: return updated rather than original document
  501. (default ``False``)
  502. :param update: Django-style update keyword arguments
  503. .. versionadded:: 0.9
  504. """
  505. if remove and new:
  506. raise OperationError('Conflicting parameters: remove and new')
  507. if not update and not upsert and not remove:
  508. raise OperationError(
  509. 'No update parameters, must either update or remove')
  510. queryset = self.clone()
  511. query = queryset._query
  512. if not IS_PYMONGO_3 or not remove:
  513. update = transform.update(queryset._document, **update)
  514. sort = queryset._ordering
  515. try:
  516. if IS_PYMONGO_3:
  517. if full_response:
  518. msg = 'With PyMongo 3+, it is not possible anymore to get the full response.'
  519. warnings.warn(msg, DeprecationWarning)
  520. if remove:
  521. result = queryset._collection.find_one_and_delete(
  522. query, sort=sort, **self._cursor_args)
  523. else:
  524. if new:
  525. return_doc = ReturnDocument.AFTER
  526. else:
  527. return_doc = ReturnDocument.BEFORE
  528. result = queryset._collection.find_one_and_update(
  529. query, update, upsert=upsert, sort=sort, return_document=return_doc,
  530. **self._cursor_args)
  531. else:
  532. result = queryset._collection.find_and_modify(
  533. query, update, upsert=upsert, sort=sort, remove=remove, new=new,
  534. full_response=full_response, **self._cursor_args)
  535. except pymongo.errors.DuplicateKeyError as err:
  536. raise NotUniqueError(u'Update failed (%s)' % err)
  537. except pymongo.errors.OperationFailure as err:
  538. raise OperationError(u'Update failed (%s)' % err)
  539. if full_response:
  540. if result['value'] is not None:
  541. result['value'] = self._document._from_son(result['value'], only_fields=self.only_fields)
  542. else:
  543. if result is not None:
  544. result = self._document._from_son(result, only_fields=self.only_fields)
  545. return result
  546. def with_id(self, object_id):
  547. """Retrieve the object matching the id provided. Uses `object_id` only
  548. and raises InvalidQueryError if a filter has been applied. Returns
  549. `None` if no document exists with that id.
  550. :param object_id: the value for the id of the document to look up
  551. .. versionchanged:: 0.6 Raises InvalidQueryError if filter has been set
  552. """
  553. queryset = self.clone()
  554. if not queryset._query_obj.empty:
  555. msg = 'Cannot use a filter whilst using `with_id`'
  556. raise InvalidQueryError(msg)
  557. return queryset.filter(pk=object_id).first()
  558. def in_bulk(self, object_ids):
  559. """Retrieve a set of documents by their ids.
  560. :param object_ids: a list or tuple of ``ObjectId``\ s
  561. :rtype: dict of ObjectIds as keys and collection-specific
  562. Document subclasses as values.
  563. .. versionadded:: 0.3
  564. """
  565. doc_map = {}
  566. docs = self._collection.find({'_id': {'$in': object_ids}},
  567. **self._cursor_args)
  568. if self._scalar:
  569. for doc in docs:
  570. doc_map[doc['_id']] = self._get_scalar(
  571. self._document._from_son(doc, only_fields=self.only_fields))
  572. elif self._as_pymongo:
  573. for doc in docs:
  574. doc_map[doc['_id']] = self._get_as_pymongo(doc)
  575. else:
  576. for doc in docs:
  577. doc_map[doc['_id']] = self._document._from_son(
  578. doc,
  579. only_fields=self.only_fields,
  580. _auto_dereference=self._auto_dereference)
  581. return doc_map
  582. def none(self):
  583. """Helper that just returns a list"""
  584. queryset = self.clone()
  585. queryset._none = True
  586. return queryset
  587. def no_sub_classes(self):
  588. """
  589. Only return instances of this document and not any inherited documents
  590. """
  591. if self._document._meta.get('allow_inheritance') is True:
  592. self._initial_query = {'_cls': self._document._class_name}
  593. return self
  594. def using(self, alias):
  595. """This method is for controlling which database the QuerySet will be
  596. evaluated against if you are using more than one database.
  597. :param alias: The database alias
  598. .. versionadded:: 0.9
  599. """
  600. with switch_db(self._document, alias) as cls:
  601. collection = cls._get_collection()
  602. return self._clone_into(self.__class__(self._document, collection))
  603. def clone(self):
  604. """Create a copy of the current queryset."""
  605. return self._clone_into(self.__class__(self._document, self._collection_obj))
  606. def _clone_into(self, new_qs):
  607. """Copy all of the relevant properties of this queryset to
  608. a new queryset (which has to be an instance of
  609. :class:`~mongoengine.queryset.base.BaseQuerySet`).
  610. """
  611. if not isinstance(new_qs, BaseQuerySet):
  612. raise OperationError(
  613. '%s is not a subclass of BaseQuerySet' % new_qs.__name__)
  614. copy_props = ('_mongo_query', '_initial_query', '_none', '_query_obj',
  615. '_where_clause', '_loaded_fields', '_ordering',
  616. '_snapshot', '_timeout', '_class_check', '_slave_okay',
  617. '_read_preference', '_iter', '_scalar', '_as_pymongo',
  618. '_limit', '_skip', '_hint', '_auto_dereference',
  619. '_search_text', 'only_fields', '_max_time_ms',
  620. '_comment')
  621. for prop in copy_props:
  622. val = getattr(self, prop)
  623. setattr(new_qs, prop, copy.copy(val))
  624. if self._cursor_obj:
  625. new_qs._cursor_obj = self._cursor_obj.clone()
  626. return new_qs
  627. def select_related(self, max_depth=1):
  628. """Handles dereferencing of :class:`~bson.dbref.DBRef` objects or
  629. :class:`~bson.object_id.ObjectId` a maximum depth in order to cut down
  630. the number queries to mongodb.
  631. .. versionadded:: 0.5
  632. """
  633. # Make select related work the same for querysets
  634. max_depth += 1
  635. queryset = self.clone()
  636. return queryset._dereference(queryset, max_depth=max_depth)
  637. def limit(self, n):
  638. """Limit the number of returned documents to `n`. This may also be
  639. achieved using array-slicing syntax (e.g. ``User.objects[:5]``).
  640. :param n: the maximum number of objects to return if n is greater than 0.
  641. When 0 is passed, returns all the documents in the cursor
  642. """
  643. queryset = self.clone()
  644. queryset._limit = n
  645. # If a cursor object has already been created, apply the limit to it.
  646. if queryset._cursor_obj:
  647. queryset._cursor_obj.limit(queryset._limit)
  648. return queryset
  649. def skip(self, n):
  650. """Skip `n` documents before returning the results. This may also be
  651. achieved using array-slicing syntax (e.g. ``User.objects[5:]``).
  652. :param n: the number of objects to skip before returning results
  653. """
  654. queryset = self.clone()
  655. queryset._skip = n
  656. # If a cursor object has already been created, apply the skip to it.
  657. if queryset._cursor_obj:
  658. queryset._cursor_obj.skip(queryset._skip)
  659. return queryset
  660. def hint(self, index=None):
  661. """Added 'hint' support, telling Mongo the proper index to use for the
  662. query.
  663. Judicious use of hints can greatly improve query performance. When
  664. doing a query on multiple fields (at least one of which is indexed)
  665. pass the indexed field as a hint to the query.
  666. Hinting will not do anything if the corresponding index does not exist.
  667. The last hint applied to this cursor takes precedence over all others.
  668. .. versionadded:: 0.5
  669. """
  670. queryset = self.clone()
  671. queryset._hint = index
  672. # If a cursor object has already been created, apply the hint to it.
  673. if queryset._cursor_obj:
  674. queryset._cursor_obj.hint(queryset._hint)
  675. return queryset
  676. def batch_size(self, size):
  677. """Limit the number of documents returned in a single batch (each
  678. batch requires a round trip to the server).
  679. See http://api.mongodb.com/python/current/api/pymongo/cursor.html#pymongo.cursor.Cursor.batch_size
  680. for details.
  681. :param size: desired size of each batch.
  682. """
  683. queryset = self.clone()
  684. queryset._batch_size = size
  685. # If a cursor object has already been created, apply the batch size to it.
  686. if queryset._cursor_obj:
  687. queryset._cursor_obj.batch_size(queryset._batch_size)
  688. return queryset
  689. def distinct(self, field):
  690. """Return a list of distinct values for a given field.
  691. :param field: the field to select distinct values from
  692. .. note:: This is a command and won't take ordering or limit into
  693. account.
  694. .. versionadded:: 0.4
  695. .. versionchanged:: 0.5 - Fixed handling references
  696. .. versionchanged:: 0.6 - Improved db_field refrence handling
  697. """
  698. queryset = self.clone()
  699. try:
  700. field = self._fields_to_dbfields([field]).pop()
  701. except LookUpError:
  702. pass
  703. distinct = self._dereference(queryset._cursor.distinct(field), 1,
  704. name=field, instance=self._document)
  705. doc_field = self._document._fields.get(field.split('.', 1)[0])
  706. instance = None
  707. # We may need to cast to the correct type eg. ListField(EmbeddedDocumentField)
  708. EmbeddedDocumentField = _import_class('EmbeddedDocumentField')
  709. ListField = _import_class('ListField')
  710. GenericEmbeddedDocumentField = _import_class('GenericEmbeddedDocumentField')
  711. if isinstance(doc_field, ListField):
  712. doc_field = getattr(doc_field, 'field', doc_field)
  713. if isinstance(doc_field, (EmbeddedDocumentField, GenericEmbeddedDocumentField)):
  714. instance = getattr(doc_field, 'document_type', None)
  715. # handle distinct on subdocuments
  716. if '.' in field:
  717. for field_part in field.split('.')[1:]:
  718. # if looping on embedded document, get the document type instance
  719. if instance and isinstance(doc_field, (EmbeddedDocumentField, GenericEmbeddedDocumentField)):
  720. doc_field = instance
  721. # now get the subdocument
  722. doc_field = getattr(doc_field, field_part, doc_field)
  723. # We may need to cast to the correct type eg. ListField(EmbeddedDocumentField)
  724. if isinstance(doc_field, ListField):
  725. doc_field = getattr(doc_field, 'field', doc_field)
  726. if isinstance(doc_field, (EmbeddedDocumentField, GenericEmbeddedDocumentField)):
  727. instance = getattr(doc_field, 'document_type', None)
  728. if instance and isinstance(doc_field, (EmbeddedDocumentField,
  729. GenericEmbeddedDocumentField)):
  730. distinct = [instance(**doc) for doc in distinct]
  731. return distinct
  732. def only(self, *fields):
  733. """Load only a subset of this document's fields. ::
  734. post = BlogPost.objects(...).only('title', 'author.name')
  735. .. note :: `only()` is chainable and will perform a union ::
  736. So with the following it will fetch both: `title` and `author.name`::
  737. post = BlogPost.objects.only('title').only('author.name')
  738. :func:`~mongoengine.queryset.QuerySet.all_fields` will reset any
  739. field filters.
  740. :param fields: fields to include
  741. .. versionadded:: 0.3
  742. .. versionchanged:: 0.5 - Added subfield support
  743. """
  744. fields = {f: QueryFieldList.ONLY for f in fields}
  745. self.only_fields = fields.keys()
  746. return self.fields(True, **fields)
  747. def exclude(self, *fields):
  748. """Opposite to .only(), exclude some document's fields. ::
  749. post = BlogPost.objects(...).exclude('comments')
  750. .. note :: `exclude()` is chainable and will perform a union ::
  751. So with the following it will exclude both: `title` and `author.name`::
  752. post = BlogPost.objects.exclude('title').exclude('author.name')
  753. :func:`~mongoengine.queryset.QuerySet.all_fields` will reset any
  754. field filters.
  755. :param fields: fields to exclude
  756. .. versionadded:: 0.5
  757. """
  758. fields = {f: QueryFieldList.EXCLUDE for f in fields}
  759. return self.fields(**fields)
  760. def fields(self, _only_called=False, **kwargs):
  761. """Manipulate how you load this document's fields. Used by `.only()`
  762. and `.exclude()` to manipulate which fields to retrieve. If called
  763. directly, use a set of kwargs similar to the MongoDB projection
  764. document. For example:
  765. Include only a subset of fields:
  766. posts = BlogPost.objects(...).fields(author=1, title=1)
  767. Exclude a specific field:
  768. posts = BlogPost.objects(...).fields(comments=0)
  769. To retrieve a subrange of array elements:
  770. posts = BlogPost.objects(...).fields(slice__comments=5)
  771. :param kwargs: A set of keyword arguments identifying what to
  772. include, exclude, or slice.
  773. .. versionadded:: 0.5
  774. """
  775. # Check for an operator and transform to mongo-style if there is
  776. operators = ['slice']
  777. cleaned_fields = []
  778. for key, value in kwargs.items():
  779. parts = key.split('__')
  780. if parts[0] in operators:
  781. op = parts.pop(0)
  782. value = {'$' + op: value}
  783. key = '.'.join(parts)
  784. cleaned_fields.append((key, value))
  785. # Sort fields by their values, explicitly excluded fields first, then
  786. # explicitly included, and then more complicated operators such as
  787. # $slice.
  788. def _sort_key(field_tuple):
  789. _, value = field_tuple
  790. if isinstance(value, int):
  791. return value # 0 for exclusion, 1 for inclusion
  792. return 2 # so that complex values appear last
  793. fields = sorted(cleaned_fields, key=_sort_key)
  794. # Clone the queryset, group all fields by their value, convert
  795. # each of them to db_fields, and set the queryset's _loaded_fields
  796. queryset = self.clone()
  797. for value, group in itertools.groupby(fields, lambda x: x[1]):
  798. fields = [field for field, value in group]
  799. fields = queryset._fields_to_dbfields(fields)
  800. queryset._loaded_fields += QueryFieldList(
  801. fields, value=value, _only_called=_only_called)
  802. return queryset
  803. def all_fields(self):
  804. """Include all fields. Reset all previously calls of .only() or
  805. .exclude(). ::
  806. post = BlogPost.objects.exclude('comments').all_fields()
  807. .. versionadded:: 0.5
  808. """
  809. queryset = self.clone()
  810. queryset._loaded_fields = QueryFieldList(
  811. always_include=queryset._loaded_fields.always_include)
  812. return queryset
  813. def order_by(self, *keys):
  814. """Order the :class:`~mongoengine.queryset.QuerySet` by the keys. The
  815. order may be specified by prepending each of the keys by a + or a -.
  816. Ascending order is assumed. If no keys are passed, existing ordering
  817. is cleared instead.
  818. :param keys: fields to order the query results by; keys may be
  819. prefixed with **+** or **-** to determine the ordering direction
  820. """
  821. queryset = self.clone()
  822. old_ordering = queryset._ordering
  823. new_ordering = queryset._get_order_by(keys)
  824. if queryset._cursor_obj:
  825. # If a cursor object has already been created, apply the sort to it
  826. if new_ordering:
  827. queryset._cursor_obj.sort(new_ordering)
  828. # If we're trying to clear a previous explicit ordering, we need
  829. # to clear the cursor entirely (because PyMongo doesn't allow
  830. # clearing an existing sort on a cursor).
  831. elif old_ordering:
  832. queryset._cursor_obj = None
  833. queryset._ordering = new_ordering
  834. return queryset
  835. def comment(self, text):
  836. """Add a comment to the query.
  837. See https://docs.mongodb.com/manual/reference/method/cursor.comment/#cursor.comment
  838. for details.
  839. """
  840. return self._chainable_method('comment', text)
  841. def explain(self, format=False):
  842. """Return an explain plan record for the
  843. :class:`~mongoengine.queryset.QuerySet`\ 's cursor.
  844. :param format: format the plan before returning it
  845. """
  846. plan = self._cursor.explain()
  847. # TODO remove this option completely - it's useless. If somebody
  848. # wants to pretty-print the output, they easily can.
  849. if format:
  850. msg = ('"format" param of BaseQuerySet.explain has been '
  851. 'deprecated and will be removed in future versions.')
  852. warnings.warn(msg, DeprecationWarning)
  853. plan = pprint.pformat(plan)
  854. return plan
  855. # DEPRECATED. Has no more impact on PyMongo 3+
  856. def snapshot(self, enabled):
  857. """Enable or disable snapshot mode when querying.
  858. :param enabled: whether or not snapshot mode is enabled
  859. ..versionchanged:: 0.5 - made chainable
  860. .. deprecated:: Ignored with PyMongo 3+
  861. """
  862. if IS_PYMONGO_3:
  863. msg = 'snapshot is deprecated as it has no impact when using PyMongo 3+.'
  864. warnings.warn(msg, DeprecationWarning)
  865. queryset = self.clone()
  866. queryset._snapshot = enabled
  867. return queryset
  868. def timeout(self, enabled):
  869. """Enable or disable the default mongod timeout when querying.
  870. :param enabled: whether or not the timeout is used
  871. ..versionchanged:: 0.5 - made chainable
  872. """
  873. queryset = self.clone()
  874. queryset._timeout = enabled
  875. return queryset
  876. # DEPRECATED. Has no more impact on PyMongo 3+
  877. def slave_okay(self, enabled):
  878. """Enable or disable the slave_okay when querying.
  879. :param enabled: whether or not the slave_okay is enabled
  880. .. deprecated:: Ignored with PyMongo 3+
  881. """
  882. if IS_PYMONGO_3:
  883. msg = 'slave_okay is deprecated as it has no impact when using PyMongo 3+.'
  884. warnings.warn(msg, DeprecationWarning)
  885. queryset = self.clone()
  886. queryset._slave_okay = enabled
  887. return queryset
  888. def read_preference(self, read_preference):
  889. """Change the read_preference when querying.
  890. :param read_preference: override ReplicaSetConnection-level
  891. preference.
  892. """
  893. validate_read_preference('read_preference', read_preference)
  894. queryset = self.clone()
  895. queryset._read_preference = read_preference
  896. queryset._cursor_obj = None # we need to re-create the cursor object whenever we apply read_preference
  897. return queryset
  898. def scalar(self, *fields):
  899. """Instead of returning Document instances, return either a specific
  900. value or a tuple of values in order.
  901. Can be used along with
  902. :func:`~mongoengine.queryset.QuerySet.no_dereference` to turn off
  903. dereferencing.
  904. .. note:: This effects all results and can be unset by calling
  905. ``scalar`` without arguments. Calls ``only`` automatically.
  906. :param fields: One or more fields to return instead of a Document.
  907. """
  908. queryset = self.clone()
  909. queryset._scalar = list(fields)
  910. if fields:
  911. queryset = queryset.only(*fields)
  912. else:
  913. queryset = queryset.all_fields()
  914. return queryset
  915. def values_list(self, *fields):
  916. """An alias for scalar"""
  917. return self.scalar(*fields)
  918. def as_pymongo(self):
  919. """Instead of returning Document instances, return raw values from
  920. pymongo.
  921. This method is particularly useful if you don't need dereferencing
  922. and care primarily about the speed of data retrieval.
  923. """
  924. queryset = self.clone()
  925. queryset._as_pymongo = True
  926. return queryset
  927. def max_time_ms(self, ms):
  928. """Wait `ms` milliseconds before killing the query on the server
  929. :param ms: the number of milliseconds before killing the query on the server
  930. """
  931. return self._chainable_method('max_time_ms', ms)
  932. # JSON Helpers
  933. def to_json(self, *args, **kwargs):
  934. """Converts a queryset to JSON"""
  935. return json_util.dumps(self.as_pymongo(), *args, **kwargs)
  936. def from_json(self, json_data):
  937. """Converts json data to unsaved objects"""
  938. son_data = json_util.loads(json_data)
  939. return [self._document._from_son(data, only_fields=self.only_fields) for data in son_data]
  940. def aggregate(self, *pipeline, **kwargs):
  941. """
  942. Perform a aggregate function based in your queryset params
  943. :param pipeline: list of aggregation commands,\
  944. see: http://docs.mongodb.org/manual/core/aggregation-pipeline/
  945. .. versionadded:: 0.9
  946. """
  947. initial_pipeline = []
  948. if self._query:
  949. initial_pipeline.append({'$match': self._query})
  950. if self._ordering:
  951. initial_pipeline.append({'$sort': dict(self._ordering)})
  952. if self._limit is not None:
  953. initial_pipeline.append({'$limit': self._limit})
  954. if self._skip is not None:
  955. initial_pipeline.append({'$skip': self._skip})
  956. pipeline = initial_pipeline + list(pipeline)
  957. if IS_PYMONGO_3 and self._read_preference is not None:
  958. return self._collection.with_options(read_preference=self._read_preference) \
  959. .aggregate(pipeline, cursor={}, **kwargs)
  960. return self._collection.aggregate(pipeline, cursor={}, **kwargs)
  961. # JS functionality
  962. def map_reduce(self, map_f, reduce_f, output, finalize_f=None, limit=None,
  963. scope=None):
  964. """Perform a map/reduce query using the current query spec
  965. and ordering. While ``map_reduce`` respects ``QuerySet`` chaining,
  966. it must be the last call made, as it does not return a maleable
  967. ``QuerySet``.
  968. See the :meth:`~mongoengine.tests.QuerySetTest.test_map_reduce`
  969. and :meth:`~mongoengine.tests.QuerySetTest.test_map_advanced`
  970. tests in ``tests.queryset.QuerySetTest`` for usage examples.
  971. :param map_f: map function, as :class:`~bson.code.Code` or string
  972. :param reduce_f: reduce function, as
  973. :class:`~bson.code.Code` or string
  974. :param output: output collection name, if set to 'inline' will try to
  975. use :class:`~pymongo.collection.Collection.inline_map_reduce`
  976. This can also be a dictionary containing output options
  977. see: http://docs.mongodb.org/manual/reference/command/mapReduce/#dbcmd.mapReduce
  978. :param finalize_f: finalize function, an optional function that
  979. performs any post-reduction processing.
  980. :param scope: values to insert into map/reduce global scope. Optional.
  981. :param limit: number of objects from current query to provide
  982. to map/reduce method
  983. Returns an iterator yielding
  984. :class:`~mongoengine.document.MapReduceDocument`.
  985. .. note::
  986. Map/Reduce changed in server version **>= 1.7.4**. The PyMongo
  987. :meth:`~pymongo.collection.Collection.map_reduce` helper requires
  988. PyMongo version **>= 1.11**.
  989. .. versionchanged:: 0.5
  990. - removed ``keep_temp`` keyword argument, which was only relevant
  991. for MongoDB server versions older than 1.7.4
  992. .. versionadded:: 0.3
  993. """
  994. queryset = self.clone()
  995. MapReduceDocument = _import_class('MapReduceDocument')
  996. if not hasattr(self._collection, 'map_reduce'):
  997. raise NotImplementedError('Requires MongoDB >= 1.7.1')
  998. map_f_scope = {}
  999. if isinstance(map_f, Code):
  1000. map_f_scope = map_f.scope
  1001. map_f = six.text_type(map_f)
  1002. map_f = Code(queryset._sub_js_fields(map_f), map_f_scope)
  1003. reduce_f_scope = {}
  1004. if isinstance(reduce_f, Code):
  1005. reduce_f_scope = reduce_f.scope
  1006. reduce_f = six.text_type(reduce_f)
  1007. reduce_f_code = queryset._sub_js_fields(reduce_f)
  1008. reduce_f = Code(reduce_f_code, reduce_f_scope)
  1009. mr_args = {'query': queryset._query}
  1010. if finalize_f:
  1011. finalize_f_scope = {}
  1012. if isinstance(finalize_f, Code):
  1013. finalize_f_scope = finalize_f.scope
  1014. finalize_f = six.text_type(finalize_f)
  1015. finalize_f_code = queryset._sub_js_fields(finalize_f)
  1016. finalize_f = Code(finalize_f_code, finalize_f_scope)
  1017. mr_args['finalize'] = finalize_f
  1018. if scope:
  1019. mr_args['scope'] = scope
  1020. if limit:
  1021. mr_args['limit'] = limit
  1022. if output == 'inline' and not queryset._ordering:
  1023. map_reduce_function = 'inline_map_reduce'
  1024. else:
  1025. map_reduce_function = 'map_reduce'
  1026. if isinstance(output, six.string_types):
  1027. mr_args['out'] = output
  1028. elif isinstance(output, dict):
  1029. ordered_output = []
  1030. for part in ('replace', 'merge', 'reduce'):
  1031. value = output.get(part)
  1032. if value:
  1033. ordered_output.append((part, value))
  1034. break
  1035. else:
  1036. raise OperationError('actionData not specified for output')
  1037. db_alias = output.get('db_alias')
  1038. remaing_args = ['db', 'sharded', 'nonAtomic']
  1039. if db_alias:
  1040. ordered_output.append(('db', get_db(db_alias).name))
  1041. del remaing_args[0]
  1042. for part in remaing_args:
  1043. value = output.get(part)
  1044. if value:
  1045. ordered_output.append((part, value))
  1046. mr_args['out'] = SON(ordered_output)
  1047. results = getattr(queryset._collection, map_reduce_function)(
  1048. map_f, reduce_f, **mr_args)
  1049. if map_reduce_function == 'map_reduce':
  1050. results = results.find()
  1051. if queryset._ordering:
  1052. results = results.sort(queryset._ordering)
  1053. for doc in results:
  1054. yield MapReduceDocument(queryset._document, queryset._collection,
  1055. doc['_id'], doc['value'])
  1056. def exec_js(self, code, *fields, **options):
  1057. """Execute a Javascript function on the server. A list of fields may be
  1058. provided, which will be translated to their correct names and supplied
  1059. as the arguments to the function. A few extra variables are added to
  1060. the function's scope: ``collection``, which is the name of the
  1061. collection in use; ``query``, which is an object representing the
  1062. current query; and ``options``, which is an object containing any
  1063. options specified as keyword arguments.
  1064. As fields in MongoEngine may use different names in the database (set
  1065. using the :attr:`db_field` keyword argument to a :class:`Field`
  1066. constructor), a mechanism exists for replacing MongoEngine field names
  1067. with the database field names in Javascript code. When accessing a
  1068. field, use square-bracket notation, and prefix the MongoEngine field
  1069. name with a tilde (~).
  1070. :param code: a string of Javascript code to execute
  1071. :param fields: fields that you will be using in your function, which
  1072. will be passed in to your function as arguments
  1073. :param options: options that you want available to the function
  1074. (accessed in Javascript through the ``options`` object)
  1075. """
  1076. queryset = self.clone()
  1077. code = queryset._sub_js_fields(code)
  1078. fields = [queryset._document._translate_field_name(f) for f in fields]
  1079. collection = queryset._document._get_collection_name()
  1080. scope = {
  1081. 'collection': collection,
  1082. 'options': options or {},
  1083. }
  1084. query = queryset._query
  1085. if queryset._where_clause:
  1086. query['$where'] = queryset._where_clause
  1087. scope['query'] = query
  1088. code = Code(code, scope=scope)
  1089. db = queryset._document._get_db()
  1090. return db.eval(code, *fields)
  1091. def where(self, where_clause):
  1092. """Filter ``QuerySet`` results with a ``$where`` clause (a Javascript
  1093. expression). Performs automatic field name substitution like
  1094. :meth:`mongoengine.queryset.Queryset.exec_js`.
  1095. .. note:: When using this mode of query, the database will call your
  1096. function, or evaluate your predicate clause, for each object
  1097. in the collection.
  1098. .. versionadded:: 0.5
  1099. """
  1100. queryset = self.clone()
  1101. where_clause = queryset._sub_js_fields(where_clause)
  1102. queryset._where_clause = where_clause
  1103. return queryset
  1104. def sum(self, field):
  1105. """Sum over the values of the specified field.
  1106. :param field: the field to sum over; use dot notation to refer to
  1107. embedded document fields
  1108. """
  1109. db_field = self._fields_to_dbfields([field]).pop()
  1110. pipeline = [
  1111. {'$match': self._query},
  1112. {'$group': {'_id': 'sum', 'total': {'$sum': '$' + db_field}}}
  1113. ]
  1114. # if we're performing a sum over a list field, we sum up all the
  1115. # elements in the list, hence we need to $unwind the arrays first
  1116. ListField = _import_class('ListField')
  1117. field_parts = field.split('.')
  1118. field_instances = self._document._lookup_field(field_parts)
  1119. if isinstance(field_instances[-1], ListField):
  1120. pipeline.insert(1, {'$unwind': '$' + field})
  1121. result = self._document._get_collection().aggregate(pipeline)
  1122. if IS_PYMONGO_3:
  1123. result = tuple(result)
  1124. else:
  1125. result = result.get('result')
  1126. if result:
  1127. return result[0]['total']
  1128. return 0
  1129. def average(self, field):
  1130. """Average over the values of the specified field.
  1131. :param field: the field to average over; use dot notation to refer to
  1132. embedded document fields
  1133. """
  1134. db_field = self._fields_to_dbfields([field]).pop()
  1135. pipeline = [
  1136. {'$match': self._query},
  1137. {'$group': {'_id': 'avg', 'total': {'$avg': '$' + db_field}}}
  1138. ]
  1139. # if we're performing an average over a list field, we average out
  1140. # all the elements in the list, hence we need to $unwind the arrays
  1141. # first
  1142. ListField = _import_class('ListField')
  1143. field_parts = field.split('.')
  1144. field_instances = self._document._lookup_field(field_parts)
  1145. if isinstance(field_instances[-1], ListField):
  1146. pipeline.insert(1, {'$unwind': '$' + field})
  1147. result = self._document._get_collection().aggregate(pipeline)
  1148. if IS_PYMONGO_3:
  1149. result = tuple(result)
  1150. else:
  1151. result = result.get('result')
  1152. if result:
  1153. return result[0]['total']
  1154. return 0
  1155. def item_frequencies(self, field, normalize=False, map_reduce=True):
  1156. """Returns a dictionary of all items present in a field across
  1157. the whole queried set of documents, and their corresponding frequency.
  1158. This is useful for generating tag clouds, or searching documents.
  1159. .. note::
  1160. Can only do direct simple mappings and cannot map across
  1161. :class:`~mongoengine.fields.ReferenceField` or
  1162. :class:`~mongoengine.fields.GenericReferenceField` for more complex
  1163. counting a manual map reduce call is required.
  1164. If the field is a :class:`~mongoengine.fields.ListField`, the items within
  1165. each list will be counted individually.
  1166. :param field: the field to use
  1167. :param normalize: normalize the results so they add to 1.0
  1168. :param map_reduce: Use map_reduce over exec_js
  1169. .. versionchanged:: 0.5 defaults to map_reduce and can handle embedded
  1170. document lookups
  1171. """
  1172. if map_reduce:
  1173. return self._item_frequencies_map_reduce(field,
  1174. normalize=normalize)
  1175. return self._item_frequencies_exec_js(field, normalize=normalize)
  1176. # Iterator helpers
  1177. def __next__(self):
  1178. """Wrap the result in a :class:`~mongoengine.Document` object.
  1179. """
  1180. if self._limit == 0 or self._none:
  1181. raise StopIteration
  1182. raw_doc = six.next(self._cursor)
  1183. if self._as_pymongo:
  1184. return self._get_as_pymongo(raw_doc)
  1185. doc = self._document._from_son(
  1186. raw_doc, _auto_dereference=self._auto_dereference,
  1187. only_fields=self.only_fields)
  1188. if self._scalar:
  1189. return self._get_scalar(doc)
  1190. return doc
  1191. next = __next__ # For Python2 support
  1192. def rewind(self):
  1193. """Rewind the cursor to its unevaluated state.
  1194. .. versionadded:: 0.3
  1195. """
  1196. self._iter = False
  1197. self._cursor.rewind()
  1198. # Properties
  1199. @property
  1200. def _collection(self):
  1201. """Property that returns the collection object. This allows us to
  1202. perform operations only if the collection is accessed.
  1203. """
  1204. return self._collection_obj
  1205. @property
  1206. def _cursor_args(self):
  1207. if not IS_PYMONGO_3:
  1208. fields_name = 'fields'
  1209. cursor_args = {
  1210. 'timeout': self._timeout,
  1211. 'snapshot': self._snapshot
  1212. }
  1213. if self._read_preference is not None:
  1214. cursor_args['read_preference'] = self._read_preference
  1215. else:
  1216. cursor_args['slave_okay'] = self._slave_okay
  1217. else:
  1218. fields_name = 'projection'
  1219. # snapshot is not handled at all by PyMongo 3+
  1220. # TODO: evaluate similar possibilities using modifiers
  1221. if self._snapshot:
  1222. msg = 'The snapshot option is not anymore available with PyMongo 3+'
  1223. warnings.warn(msg, DeprecationWarning)
  1224. cursor_args = {
  1225. 'no_cursor_timeout': not self._timeout
  1226. }
  1227. if self._loaded_fields:
  1228. cursor_args[fields_name] = self._loaded_fields.as_dict()
  1229. if self._search_text:
  1230. if fields_name not in cursor_args:
  1231. cursor_args[fields_name] = {}
  1232. cursor_args[fields_name]['_text_score'] = {'$meta': 'textScore'}
  1233. return cursor_args
  1234. @property
  1235. def _cursor(self):
  1236. """Return a PyMongo cursor object corresponding to this queryset."""
  1237. # If _cursor_obj already exists, return it immediately.
  1238. if self._cursor_obj is not None:
  1239. return self._cursor_obj
  1240. # Create a new PyMongo cursor.
  1241. # XXX In PyMongo 3+, we define the read preference on a collection
  1242. # level, not a cursor level. Thus, we need to get a cloned collection
  1243. # object using `with_options` first.
  1244. if IS_PYMONGO_3 and self._read_preference is not None:
  1245. self._cursor_obj = self._collection\
  1246. .with_options(read_preference=self._read_preference)\
  1247. .find(self._query, **self._cursor_args)
  1248. else:
  1249. self._cursor_obj = self._collection.find(self._query,
  1250. **self._cursor_args)
  1251. # Apply "where" clauses to cursor
  1252. if self._where_clause:
  1253. where_clause = self._sub_js_fields(self._where_clause)
  1254. self._cursor_obj.where(where_clause)
  1255. # Apply ordering to the cursor.
  1256. # XXX self._ordering can be equal to:
  1257. # * None if we didn't explicitly call order_by on this queryset.
  1258. # * A list of PyMongo-style sorting tuples.
  1259. # * An empty list if we explicitly called order_by() without any
  1260. # arguments. This indicates that we want to clear the default
  1261. # ordering.
  1262. if self._ordering:
  1263. # explicit ordering
  1264. self._cursor_obj.sort(self._ordering)
  1265. elif self._ordering is None and self._document._meta['ordering']:
  1266. # default ordering
  1267. order = self._get_order_by(self._document._meta['ordering'])
  1268. self._cursor_obj.sort(order)
  1269. if self._limit is not None:
  1270. self._cursor_obj.limit(self._limit)
  1271. if self._skip is not None:
  1272. self._cursor_obj.skip(self._skip)
  1273. if self._hint != -1:
  1274. self._cursor_obj.hint(self._hint)
  1275. if self._batch_size is not None:
  1276. self._cursor_obj.batch_size(self._batch_size)
  1277. if self._comment is not None:
  1278. self._cursor_obj.comment(self._comment)
  1279. return self._cursor_obj
  1280. def __deepcopy__(self, memo):
  1281. """Essential for chained queries with ReferenceFields involved"""
  1282. return self.clone()
  1283. @property
  1284. def _query(self):
  1285. if self._mongo_query is None:
  1286. self._mongo_query = self._query_obj.to_query(self._document)
  1287. if self._class_check and self._initial_query:
  1288. if '_cls' in self._mongo_query:
  1289. self._mongo_query = {'$and': [self._initial_query, self._mongo_query]}
  1290. else:
  1291. self._mongo_query.update(self._initial_query)
  1292. return self._mongo_query
  1293. @property
  1294. def _dereference(self):
  1295. if not self.__dereference:
  1296. self.__dereference = _import_class('DeReference')()
  1297. return self.__dereference
  1298. def no_dereference(self):
  1299. """Turn off any dereferencing for the results of this queryset."""
  1300. queryset = self.clone()
  1301. queryset._auto_dereference = False
  1302. return queryset
  1303. # Helper Functions
  1304. def _item_frequencies_map_reduce(self, field, normalize=False):
  1305. map_func = """
  1306. function() {
  1307. var path = '{{~%(field)s}}'.split('.');
  1308. var field = this;
  1309. for (p in path) {
  1310. if (typeof field != 'undefined')
  1311. field = field[path[p]];
  1312. else
  1313. break;
  1314. }
  1315. if (field && field.constructor == Array) {
  1316. field.forEach(function(item) {
  1317. emit(item, 1);
  1318. });
  1319. } else if (typeof field != 'undefined') {
  1320. emit(field, 1);
  1321. } else {
  1322. emit(null, 1);
  1323. }
  1324. }
  1325. """ % {'field': field}
  1326. reduce_func = """
  1327. function(key, values) {
  1328. var total = 0;
  1329. var valuesSize = values.length;
  1330. for (var i=0; i < valuesSize; i++) {
  1331. total += parseInt(values[i], 10);
  1332. }
  1333. return total;
  1334. }
  1335. """
  1336. values = self.map_reduce(map_func, reduce_func, 'inline')
  1337. frequencies = {}
  1338. for f in values:
  1339. key = f.key
  1340. if isinstance(key, float):
  1341. if int(key) == key:
  1342. key = int(key)
  1343. frequencies[key] = int(f.value)
  1344. if normalize:
  1345. count = sum(frequencies.values())
  1346. frequencies = {k: float(v) / count
  1347. for k, v in frequencies.items()}
  1348. return frequencies
  1349. def _item_frequencies_exec_js(self, field, normalize=False):
  1350. """Uses exec_js to execute"""
  1351. freq_func = """
  1352. function(path) {
  1353. var path = path.split('.');
  1354. var total = 0.0;
  1355. db[collection].find(query).forEach(function(doc) {
  1356. var field = doc;
  1357. for (p in path) {
  1358. if (field)
  1359. field = field[path[p]];
  1360. else
  1361. break;
  1362. }
  1363. if (field && field.constructor == Array) {
  1364. total += field.length;
  1365. } else {
  1366. total++;
  1367. }
  1368. });
  1369. var frequencies = {};
  1370. var types = {};
  1371. var inc = 1.0;
  1372. db[collection].find(query).forEach(function(doc) {
  1373. field = doc;
  1374. for (p in path) {
  1375. if (field)
  1376. field = field[path[p]];
  1377. else
  1378. break;
  1379. }
  1380. if (field && field.constructor == Array) {
  1381. field.forEach(function(item) {
  1382. frequencies[item] = inc + (isNaN(frequencies[item]) ? 0: frequencies[item]);
  1383. });
  1384. } else {
  1385. var item = field;
  1386. types[item] = item;
  1387. frequencies[item] = inc + (isNaN(frequencies[item]) ? 0: frequencies[item]);
  1388. }
  1389. });
  1390. return [total, frequencies, types];
  1391. }
  1392. """
  1393. total, data, types = self.exec_js(freq_func, field)
  1394. values = {types.get(k): int(v) for k, v in data.iteritems()}
  1395. if normalize:
  1396. values = {k: float(v) / total for k, v in values.items()}
  1397. frequencies = {}
  1398. for k, v in values.iteritems():
  1399. if isinstance(k, float):
  1400. if int(k) == k:
  1401. k = int(k)
  1402. frequencies[k] = v
  1403. return frequencies
  1404. def _fields_to_dbfields(self, fields):
  1405. """Translate fields' paths to their db equivalents."""
  1406. subclasses = []
  1407. if self._document._meta['allow_inheritance']:
  1408. subclasses = [get_document(x)
  1409. for x in self._document._subclasses][1:]
  1410. db_field_paths = []
  1411. for field in fields:
  1412. field_parts = field.split('.')
  1413. try:
  1414. field = '.'.join(
  1415. f if isinstance(f, six.string_types) else f.db_field
  1416. for f in self._document._lookup_field(field_parts)
  1417. )
  1418. db_field_paths.append(field)
  1419. except LookUpError as err:
  1420. found = False
  1421. # If a field path wasn't found on the main document, go
  1422. # through its subclasses and see if it exists on any of them.
  1423. for subdoc in subclasses:
  1424. try:
  1425. subfield = '.'.join(
  1426. f if isinstance(f, six.string_types) else f.db_field
  1427. for f in subdoc._lookup_field(field_parts)
  1428. )
  1429. db_field_paths.append(subfield)
  1430. found = True
  1431. break
  1432. except LookUpError:
  1433. pass
  1434. if not found:
  1435. raise err
  1436. return db_field_paths
  1437. def _get_order_by(self, keys):
  1438. """Given a list of MongoEngine-style sort keys, return a list
  1439. of sorting tuples that can be applied to a PyMongo cursor. For
  1440. example:
  1441. >>> qs._get_order_by(['-last_name', 'first_name'])
  1442. [('last_name', -1), ('first_name', 1)]
  1443. """
  1444. key_list = []
  1445. for key in keys:
  1446. if not key:
  1447. continue
  1448. if key == '$text_score':
  1449. key_list.append(('_text_score', {'$meta': 'textScore'}))
  1450. continue
  1451. direction = pymongo.ASCENDING
  1452. if key[0] == '-':
  1453. direction = pymongo.DESCENDING
  1454. if key[0] in ('-', '+'):
  1455. key = key[1:]
  1456. key = key.replace('__', '.')
  1457. try:
  1458. key = self._document._translate_field_name(key)
  1459. except Exception:
  1460. # TODO this exception should be more specific
  1461. pass
  1462. key_list.append((key, direction))
  1463. return key_list
  1464. def _get_scalar(self, doc):
  1465. def lookup(obj, name):
  1466. chunks = name.split('__')
  1467. for chunk in chunks:
  1468. obj = getattr(obj, chunk)
  1469. return obj
  1470. data = [lookup(doc, n) for n in self._scalar]
  1471. if len(data) == 1:
  1472. return data[0]
  1473. return tuple(data)
  1474. def _get_as_pymongo(self, doc):
  1475. """Clean up a PyMongo doc, removing fields that were only fetched
  1476. for the sake of MongoEngine's implementation, and return it.
  1477. """
  1478. # Always remove _cls as a MongoEngine's implementation detail.
  1479. if '_cls' in doc:
  1480. del doc['_cls']
  1481. # If the _id was not included in a .only or was excluded in a .exclude,
  1482. # remove it from the doc (we always fetch it so that we can properly
  1483. # construct documents).
  1484. fields = self._loaded_fields
  1485. if fields and '_id' in doc and (
  1486. (fields.value == QueryFieldList.ONLY and '_id' not in fields.fields) or
  1487. (fields.value == QueryFieldList.EXCLUDE and '_id' in fields.fields)
  1488. ):
  1489. del doc['_id']
  1490. return doc
  1491. def _sub_js_fields(self, code):
  1492. """When fields are specified with [~fieldname] syntax, where
  1493. *fieldname* is the Python name of a field, *fieldname* will be
  1494. substituted for the MongoDB name of the field (specified using the
  1495. :attr:`name` keyword argument in a field's constructor).
  1496. """
  1497. def field_sub(match):
  1498. # Extract just the field name, and look up the field objects
  1499. field_name = match.group(1).split('.')
  1500. fields = self._document._lookup_field(field_name)
  1501. # Substitute the correct name for the field into the javascript
  1502. return u'["%s"]' % fields[-1].db_field
  1503. def field_path_sub(match):
  1504. # Extract just the field name, and look up the field objects
  1505. field_name = match.group(1).split('.')
  1506. fields = self._document._lookup_field(field_name)
  1507. # Substitute the correct name for the field into the javascript
  1508. return '.'.join([f.db_field for f in fields])
  1509. code = re.sub(r'\[\s*~([A-z_][A-z_0-9.]+?)\s*\]', field_sub, code)
  1510. code = re.sub(r'\{\{\s*~([A-z_][A-z_0-9.]+?)\s*\}\}', field_path_sub,
  1511. code)
  1512. return code
  1513. def _chainable_method(self, method_name, val):
  1514. """Call a particular method on the PyMongo cursor call a particular chainable method
  1515. with the provided value.
  1516. """
  1517. queryset = self.clone()
  1518. # Get an existing cursor object or create a new one
  1519. cursor = queryset._cursor
  1520. # Find the requested method on the cursor and call it with the
  1521. # provided value
  1522. getattr(cursor, method_name)(val)
  1523. # Cache the value on the queryset._{method_name}
  1524. setattr(queryset, '_' + method_name, val)
  1525. return queryset
  1526. # Deprecated
  1527. def ensure_index(self, **kwargs):
  1528. """Deprecated use :func:`Document.ensure_index`"""
  1529. msg = ('Doc.objects()._ensure_index() is deprecated. '
  1530. 'Use Doc.ensure_index() instead.')
  1531. warnings.warn(msg, DeprecationWarning)
  1532. self._document.__class__.ensure_index(**kwargs)
  1533. return self
  1534. def _ensure_indexes(self):
  1535. """Deprecated use :func:`~Document.ensure_indexes`"""
  1536. msg = ('Doc.objects()._ensure_indexes() is deprecated. '
  1537. 'Use Doc.ensure_indexes() instead.')
  1538. warnings.warn(msg, DeprecationWarning)
  1539. self._document.__class__.ensure_indexes()