Development of an internal social media platform with personalised dashboards for students
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

mongodb.py 8.5KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264
  1. # -*- coding: utf-8 -*-
  2. """
  3. celery.backends.mongodb
  4. ~~~~~~~~~~~~~~~~~~~~~~~
  5. MongoDB result store backend.
  6. """
  7. from __future__ import absolute_import
  8. from datetime import datetime
  9. from kombu.syn import detect_environment
  10. from kombu.utils import cached_property
  11. from kombu.utils.url import maybe_sanitize_url
  12. from celery import states
  13. from celery.exceptions import ImproperlyConfigured
  14. from celery.five import items, string_t
  15. from celery.utils.timeutils import maybe_timedelta
  16. from .base import BaseBackend
  17. try:
  18. import pymongo
  19. except ImportError: # pragma: no cover
  20. pymongo = None # noqa
  21. if pymongo:
  22. try:
  23. from bson.binary import Binary
  24. except ImportError: # pragma: no cover
  25. from pymongo.binary import Binary # noqa
  26. else: # pragma: no cover
  27. Binary = None # noqa
  28. __all__ = ['MongoBackend']
  29. class MongoBackend(BaseBackend):
  30. """MongoDB result backend.
  31. :raises celery.exceptions.ImproperlyConfigured: if
  32. module :mod:`pymongo` is not available.
  33. """
  34. host = 'localhost'
  35. port = 27017
  36. user = None
  37. password = None
  38. database_name = 'celery'
  39. taskmeta_collection = 'celery_taskmeta'
  40. max_pool_size = 10
  41. options = None
  42. supports_autoexpire = False
  43. _connection = None
  44. def __init__(self, app=None, url=None, **kwargs):
  45. self.options = {}
  46. super(MongoBackend, self).__init__(app, **kwargs)
  47. self.expires = kwargs.get('expires') or maybe_timedelta(
  48. self.app.conf.CELERY_TASK_RESULT_EXPIRES)
  49. if not pymongo:
  50. raise ImproperlyConfigured(
  51. 'You need to install the pymongo library to use the '
  52. 'MongoDB backend.')
  53. config = self.app.conf.get('CELERY_MONGODB_BACKEND_SETTINGS')
  54. if config is not None:
  55. if not isinstance(config, dict):
  56. raise ImproperlyConfigured(
  57. 'MongoDB backend settings should be grouped in a dict')
  58. config = dict(config) # do not modify original
  59. self.host = config.pop('host', self.host)
  60. self.port = int(config.pop('port', self.port))
  61. self.user = config.pop('user', self.user)
  62. self.password = config.pop('password', self.password)
  63. self.database_name = config.pop('database', self.database_name)
  64. self.taskmeta_collection = config.pop(
  65. 'taskmeta_collection', self.taskmeta_collection,
  66. )
  67. self.options = dict(config, **config.pop('options', None) or {})
  68. # Set option defaults
  69. for key, value in items(self._prepare_client_options()):
  70. self.options.setdefault(key, value)
  71. self.url = url
  72. if self.url:
  73. # Specifying backend as an URL
  74. self.host = self.url
  75. def _prepare_client_options(self):
  76. if pymongo.version_tuple >= (3, ):
  77. return {'maxPoolSize': self.max_pool_size}
  78. else: # pragma: no cover
  79. options = {
  80. 'max_pool_size': self.max_pool_size,
  81. 'auto_start_request': False
  82. }
  83. if detect_environment() != 'default':
  84. options['use_greenlets'] = True
  85. return options
  86. def _get_connection(self):
  87. """Connect to the MongoDB server."""
  88. if self._connection is None:
  89. from pymongo import MongoClient
  90. # The first pymongo.Connection() argument (host) can be
  91. # a list of ['host:port'] elements or a mongodb connection
  92. # URI. If this is the case, don't use self.port
  93. # but let pymongo get the port(s) from the URI instead.
  94. # This enables the use of replica sets and sharding.
  95. # See pymongo.Connection() for more info.
  96. url = self.host
  97. if isinstance(url, string_t) \
  98. and not url.startswith('mongodb://'):
  99. url = 'mongodb://{0}:{1}'.format(url, self.port)
  100. if url == 'mongodb://':
  101. url = url + 'localhost'
  102. self._connection = MongoClient(host=url, **self.options)
  103. return self._connection
  104. def process_cleanup(self):
  105. if self._connection is not None:
  106. # MongoDB connection will be closed automatically when object
  107. # goes out of scope
  108. del(self.collection)
  109. del(self.database)
  110. self._connection = None
  111. def _store_result(self, task_id, result, status,
  112. traceback=None, request=None, **kwargs):
  113. """Store return value and status of an executed task."""
  114. meta = {'_id': task_id,
  115. 'status': status,
  116. 'result': Binary(self.encode(result)),
  117. 'date_done': datetime.utcnow(),
  118. 'traceback': Binary(self.encode(traceback)),
  119. 'children': Binary(self.encode(
  120. self.current_task_children(request),
  121. ))}
  122. self.collection.save(meta)
  123. return result
  124. def _get_task_meta_for(self, task_id):
  125. """Get task metadata for a task by id."""
  126. obj = self.collection.find_one({'_id': task_id})
  127. if not obj:
  128. return {'status': states.PENDING, 'result': None}
  129. meta = {
  130. 'task_id': obj['_id'],
  131. 'status': obj['status'],
  132. 'result': self.decode(obj['result']),
  133. 'date_done': obj['date_done'],
  134. 'traceback': self.decode(obj['traceback']),
  135. 'children': self.decode(obj['children']),
  136. }
  137. return meta
  138. def _save_group(self, group_id, result):
  139. """Save the group result."""
  140. meta = {'_id': group_id,
  141. 'result': Binary(self.encode(result)),
  142. 'date_done': datetime.utcnow()}
  143. self.collection.save(meta)
  144. return result
  145. def _restore_group(self, group_id):
  146. """Get the result for a group by id."""
  147. obj = self.collection.find_one({'_id': group_id})
  148. if not obj:
  149. return
  150. meta = {
  151. 'task_id': obj['_id'],
  152. 'result': self.decode(obj['result']),
  153. 'date_done': obj['date_done'],
  154. }
  155. return meta
  156. def _delete_group(self, group_id):
  157. """Delete a group by id."""
  158. self.collection.remove({'_id': group_id})
  159. def _forget(self, task_id):
  160. """Remove result from MongoDB.
  161. :raises celery.exceptions.OperationsError:
  162. if the task_id could not be removed.
  163. """
  164. # By using safe=True, this will wait until it receives a response from
  165. # the server. Likewise, it will raise an OperationsError if the
  166. # response was unable to be completed.
  167. self.collection.remove({'_id': task_id})
  168. def cleanup(self):
  169. """Delete expired metadata."""
  170. self.collection.remove(
  171. {'date_done': {'$lt': self.app.now() - self.expires}},
  172. )
  173. def __reduce__(self, args=(), kwargs={}):
  174. return super(MongoBackend, self).__reduce__(
  175. args, dict(kwargs, expires=self.expires, url=self.url),
  176. )
  177. def _get_database(self):
  178. conn = self._get_connection()
  179. db = conn[self.database_name]
  180. if self.user and self.password:
  181. if not db.authenticate(self.user,
  182. self.password):
  183. raise ImproperlyConfigured(
  184. 'Invalid MongoDB username or password.')
  185. return db
  186. @cached_property
  187. def database(self):
  188. """Get database from MongoDB connection and perform authentication
  189. if necessary."""
  190. return self._get_database()
  191. @cached_property
  192. def collection(self):
  193. """Get the metadata task collection."""
  194. collection = self.database[self.taskmeta_collection]
  195. # Ensure an index on date_done is there, if not process the index
  196. # in the background. Once completed cleanup will be much faster
  197. collection.ensure_index('date_done', background='true')
  198. return collection
  199. def as_uri(self, include_password=False):
  200. """Return the backend as an URI.
  201. :keyword include_password: Censor passwords.
  202. """
  203. if not self.url:
  204. return 'mongodb://'
  205. if include_password:
  206. return self.url
  207. if ',' not in self.url:
  208. return maybe_sanitize_url(self.url)
  209. uri1, remainder = self.url.split(',', 1)
  210. return ','.join([maybe_sanitize_url(uri1), remainder])