monco.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309
  1. # -*- coding: utf-8 -*-
  2. """Monco: a MongoDB database backend
  3. Classes and functions used to issue queries to a MongoDB database.
  4. Copyright 2016-2017 Davide Alberani <da@erlug.linux.it>
  5. RaspiBO <info@raspibo.org>
  6. Licensed under the Apache License, Version 2.0 (the "License");
  7. you may not use this file except in compliance with the License.
  8. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
  9. Unless required by applicable law or agreed to in writing, software
  10. distributed under the License is distributed on an "AS IS" BASIS,
  11. WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. See the License for the specific language governing permissions and
  13. limitations under the License.
  14. """
  15. import re
  16. import pymongo
  17. from bson.objectid import ObjectId
  18. re_objectid = re.compile(r'[0-9a-f]{24}')
  19. _force_conversion = {
  20. '_id': ObjectId
  21. }
  22. def convert_obj(obj):
  23. """Convert an object in a format suitable to be stored in MongoDB.
  24. :param obj: object to convert
  25. :returns: object that can be stored in MongoDB.
  26. """
  27. if obj is None:
  28. return None
  29. if isinstance(obj, bool):
  30. return obj
  31. try:
  32. if re_objectid.match(obj):
  33. return ObjectId(obj)
  34. except:
  35. pass
  36. return obj
  37. def convert(seq):
  38. """Convert an object to a format suitable to be stored in MongoDB,
  39. descending lists, tuples and dictionaries (a copy is returned).
  40. :param seq: sequence or object to convert
  41. :returns: object that can be stored in MongoDB.
  42. """
  43. if isinstance(seq, dict):
  44. d = {}
  45. for key, item in seq.items():
  46. if key in _force_conversion:
  47. try:
  48. d[key] = _force_conversion[key](item)
  49. except:
  50. d[key] = item
  51. else:
  52. d[key] = convert(item)
  53. return d
  54. if isinstance(seq, (list, tuple)):
  55. return [convert(x) for x in seq]
  56. return convert_obj(seq)
  57. class MoncoError(Exception):
  58. """Base class for Monco exceptions."""
  59. pass
  60. class MoncoConnectionError(MoncoError):
  61. """Monco exceptions raise when a connection problem occurs."""
  62. pass
  63. class Monco(object):
  64. """MongoDB connector."""
  65. db = None
  66. connection = None
  67. # map operations on lists of items.
  68. _operations = {
  69. 'update': '$set',
  70. 'append': '$push',
  71. 'appendUnique': '$addToSet',
  72. 'delete': '$pull',
  73. 'increment': '$inc'
  74. }
  75. def __init__(self, dbName, url=None):
  76. """Initialize the instance, connecting to the database.
  77. :param dbName: name of the database
  78. :type dbName: str (or None to use the dbName passed at initialization)
  79. :param url: URL of the database
  80. :type url: str (or None to connect to localhost)
  81. """
  82. self._url = url
  83. self._dbName = dbName
  84. self.connect()
  85. def connect(self, dbName=None, url=None):
  86. """Connect to the database.
  87. :param dbName: name of the database
  88. :type dbName: str (or None to use the dbName passed at initialization)
  89. :param url: URL of the database
  90. :type url: str (or None to connect to localhost)
  91. :returns: the database we're connected to
  92. :rtype: :class:`~pymongo.database.Database`
  93. """
  94. if self.db is not None:
  95. return self.db
  96. if url:
  97. self._url = url
  98. if dbName:
  99. self._dbName = dbName
  100. if not self._dbName:
  101. raise MoncoConnectionError('no database name specified')
  102. self.connection = pymongo.MongoClient(self._url)
  103. self.db = self.connection[self._dbName]
  104. return self.db
  105. def getOne(self, collection, query=None):
  106. """Get a single document with the specified `query`.
  107. :param collection: search the document in this collection
  108. :type collection: str
  109. :param query: query to filter the documents
  110. :type query: dict or None
  111. :returns: the first document matching the query
  112. :rtype: dict
  113. """
  114. results = self.query(collection, convert(query))
  115. return results and results[0] or {}
  116. def get(self, collection, _id):
  117. """Get a single document with the specified `_id`.
  118. :param collection: search the document in this collection
  119. :type collection: str
  120. :param _id: unique ID of the document
  121. :type _id: str or :class:`~bson.objectid.ObjectId`
  122. :returns: the document with the given `_id`
  123. :rtype: dict
  124. """
  125. return self.getOne(collection, {'_id': _id})
  126. def query(self, collection, query=None, condition='or'):
  127. """Get multiple documents matching a query.
  128. :param collection: search for documents in this collection
  129. :type collection: str
  130. :param query: search for documents with those attributes
  131. :type query: dict, list or None
  132. :returns: list of matching documents
  133. :rtype: list
  134. """
  135. db = self.connect()
  136. query = convert(query or {})
  137. if isinstance(query, (list, tuple)):
  138. query = {'$%s' % condition: query}
  139. return list(db[collection].find(query))
  140. def add(self, collection, data, _id=None):
  141. """Insert a new document.
  142. :param collection: insert the document in this collection
  143. :type collection: str
  144. :param data: the document to store
  145. :type data: dict
  146. :param _id: the _id of the document to store; if None, it's generated
  147. :type _id: object
  148. :returns: the document, as created in the database
  149. :rtype: dict
  150. """
  151. db = self.connect()
  152. data = convert(data)
  153. if _id is not None:
  154. data['_id'] = _id
  155. _id = db[collection].insert(data)
  156. return self.get(collection, _id)
  157. def insertOne(self, collection, data):
  158. """Insert a document, avoiding duplicates.
  159. :param collection: update a document in this collection
  160. :type collection: str
  161. :param data: the document information to store
  162. :type data: dict
  163. :returns: True if the document was already present
  164. :rtype: bool
  165. """
  166. db = self.connect()
  167. data = convert(data)
  168. ret = db[collection].update(data, {'$set': data}, upsert=True)
  169. return ret['updatedExisting']
  170. def _buildSearchPattern(self, data, searchBy):
  171. """Return an OR condition."""
  172. _or = []
  173. for searchPattern in searchBy:
  174. try:
  175. _or.append(dict([(k, data[k]) for k in searchPattern if k in data]))
  176. except KeyError:
  177. continue
  178. return _or
  179. def update(self, collection, _id_or_query, data, operation='update',
  180. updateList=None, create=True):
  181. """Update an existing document or create it, if requested.
  182. _id_or_query can be an ID, a dict representing a query or a list of tuples.
  183. In the latter case, the tuples are put in OR; a tuple match if all of its
  184. items from 'data' are contained in the document.
  185. :param collection: update a document in this collection
  186. :type collection: str
  187. :param _id_or_query: ID of the document to be updated, or a query or a list of attributes in the data that must match
  188. :type _id_or_query: str or :class:`~bson.objectid.ObjectId` or iterable
  189. :param data: the updated information to store
  190. :type data: dict
  191. :param operation: operation used to update the document or a portion of it, like a list (update, append, appendUnique, delete, increment)
  192. :type operation: str
  193. :param updateList: if set, it's considered the name of a list (the first matching element will be updated)
  194. :type updateList: str
  195. :param create: if True, the document is created if no document matches
  196. :type create: bool
  197. :returns: a boolean (True if an existing document was updated) and the document after the update
  198. :rtype: tuple of (bool, dict)
  199. """
  200. db = self.connect()
  201. data = convert(data or {})
  202. _id_or_query = convert(_id_or_query)
  203. if isinstance(_id_or_query, (list, tuple)):
  204. _id_or_query = {'$or': self._buildSearchPattern(data, _id_or_query)}
  205. elif not isinstance(_id_or_query, dict):
  206. _id_or_query = {'_id': _id_or_query}
  207. if '_id' in data:
  208. del data['_id']
  209. operator = self._operations.get(operation)
  210. if updateList:
  211. newData = {}
  212. for key, value in data.items():
  213. newData['%s.$.%s' % (updateList, key)] = value
  214. data = newData
  215. res = db[collection].find_and_modify(query=_id_or_query,
  216. update={operator: data}, full_response=True, new=True, upsert=create)
  217. lastErrorObject = res.get('lastErrorObject') or {}
  218. return lastErrorObject.get('updatedExisting', False), res.get('value') or {}
  219. def updateMany(self, collection, query, data):
  220. """Update multiple existing documents.
  221. query can be an ID or a dict representing a query.
  222. :param collection: update documents in this collection
  223. :type collection: str
  224. :param query: a query or a list of attributes in the data that must match
  225. :type query: str or :class:`~bson.objectid.ObjectId` or iterable
  226. :param data: the updated information to store
  227. :type data: dict
  228. :returns: a dict with the success state and number of updated items
  229. :rtype: dict
  230. """
  231. db = self.connect()
  232. data = convert(data or {})
  233. query = convert(query)
  234. if not isinstance(query, dict):
  235. query = {'_id': query}
  236. if '_id' in data:
  237. del data['_id']
  238. return db[collection].update(query, {'$set': data}, multi=True)
  239. def delete(self, collection, _id_or_query=None, force=False):
  240. """Remove one or more documents from a collection.
  241. :param collection: search the documents in this collection
  242. :type collection: str
  243. :param _id_or_query: unique ID of the document or query to match multiple documents
  244. :type _id_or_query: str or :class:`~bson.objectid.ObjectId` or dict
  245. :param force: force the deletion of all documents, when `_id_or_query` is empty
  246. :type force: bool
  247. :returns: dictionary with the number or removed documents
  248. :rtype: dict
  249. """
  250. if not _id_or_query and not force:
  251. return
  252. db = self.connect()
  253. if not isinstance(_id_or_query, dict):
  254. _id_or_query = {'_id': _id_or_query}
  255. _id_or_query = convert(_id_or_query)
  256. return db[collection].remove(_id_or_query)