source file: /opt/devel/celery/testproj/../celery/task.py
file stats: 105 lines, 95 executed: 90.5% covered
   1. from carrot.connection import DjangoAMQPConnection
   2. from celery.log import setup_logger
   3. from celery.conf import TASK_META_USE_DB
   4. from celery.registry import tasks
   5. from celery.messaging import TaskPublisher, TaskConsumer
   6. from celery.models import TaskMeta
   7. from django.core.cache import cache
   8. from datetime import timedelta
   9. import uuid
  10. import traceback
  11. 
  12. 
  13. def delay_task(task_name, **kwargs):
  14.     """Delay a task for execution by the ``celery`` daemon.
  15. 
  16.     Examples
  17.     --------
  18.         >>> delay_task("update_record", name="George Constanza", age=32)
  19. 
  20.     """
  21.     if task_name not in tasks:
  22.         raise tasks.NotRegistered(
  23.                 "Task with name %s not registered in the task registry." % (
  24.                     task_name))
  25.     publisher = TaskPublisher(connection=DjangoAMQPConnection)
  26.     task_id = publisher.delay_task(task_name, **kwargs)
  27.     publisher.close()
  28.     return task_id
  29. 
  30. 
  31. def discard_all():
  32.     """Discard all waiting tasks.
  33. 
  34.     This will ignore all tasks waiting for execution, and they will
  35.     be deleted from the messaging server.
  36. 
  37.     Returns the number of tasks discarded.
  38. 
  39.     """
  40.     consumer = TaskConsumer(connection=DjangoAMQPConnection)
  41.     discarded_count = consumer.discard_all()
  42.     consumer.close()
  43.     return discarded_count
  44. 
  45. 
  46. def gen_task_done_cache_key(task_id):
  47.     """Generate a cache key for marking a task as done."""
  48.     return "celery-task-done-marker-%s" % task_id
  49. 
  50. 
  51. def mark_as_done(task_id, result):
  52.     """Mark task as done (executed).
  53. 
  54.     if ``settings.TASK_META_USE_DB`` is ``True``, this will
  55.     use the :class:`celery.models.TaskMeta` model, if not memcached
  56.     is used.
  57. 
  58.     """
  59.     if result is None:
  60.         result = True
  61.     if TASK_META_USE_DB:
  62.         TaskMeta.objects.mark_as_done(task_id)
  63.     else:
  64.         cache_key = gen_task_done_cache_key(task_id)
  65.         cache.set(cache_key, result)
  66. 
  67. 
  68. def is_done(task_id):
  69.     """Returns ``True`` if task with ``task_id`` has been executed."""
  70.     if TASK_META_USE_DB:
  71.         return TaskMeta.objects.is_done(task_id)
  72.     else:
  73.         cache_key = gen_task_done_cache_key(task_id)
  74.         return bool(cache.get(cache_key))
  75. 
  76. 
  77. class Task(object):
  78.     """A task that can be delayed for execution by the ``celery`` daemon.
  79. 
  80.     All subclasses of ``Task`` has to define the ``name`` attribute, which is
  81.     the name of the task that can be passed to ``celery.task.delay_task``,
  82.     it also has to define the ``run`` method, which is the actual method the
  83.     ``celery`` daemon executes. This method does not support positional
  84.     arguments, only keyword arguments.
  85. 
  86.     Examples
  87.     --------
  88. 
  89.     This is a simple task just logging a message,
  90. 
  91.         >>> from celery.task import tasks, Task
  92.         >>> class MyTask(Task):
  93.         ...     name = "mytask"
  94.         ...
  95.         ...     def run(self, some_arg=None, **kwargs):
  96.         ...         logger = self.get_logger(**kwargs)
  97.         ...         logger.info("Running MyTask with arg some_arg=%s" %
  98.         ...                     some_arg))
  99.         ... tasks.register(MyTask)
 100. 
 101.     You can delay the task using the classmethod ``delay``...
 102. 
 103.         >>> MyTask.delay(some_arg="foo")
 104. 
 105.     ...or using the ``celery.task.delay_task`` function, by passing the
 106.     name of the task.
 107. 
 108.         >>> from celery.task import delay_task
 109.         >>> delay_task(MyTask.name, some_arg="foo")
 110. 
 111.     """
 112.     name = None
 113.     type = "regular"
 114. 
 115.     def __init__(self):
 116.         if not self.name:
 117.             raise NotImplementedError("Tasks must define a name attribute.")
 118. 
 119.     def __call__(self, **kwargs):
 120.         """The ``__call__`` is called when you do ``Task().run()`` and calls
 121.         the ``run`` method. It also catches any exceptions and logs them."""
 122.         try:
 123.             retval = self.run(**kwargs)
 124.         except Exception, e:
 125.             logger = self.get_logger(**kwargs)
 126.             logger.critical("Task got exception %s: %s\n%s" % (
 127.                                 e.__class__, e, traceback.format_exc()))
 128.             return
 129.         else:
 130.             return retval
 131. 
 132.     def run(self, **kwargs):
 133.         """The actual task. All subclasses of :class:`Task` must define
 134.         the run method, if not a ``NotImplementedError`` exception is raised.
 135.         """
 136.         raise NotImplementedError("Tasks must define a run method.")
 137. 
 138.     def get_logger(self, **kwargs):
 139.         """Get a process-aware logger object."""
 140.         return setup_logger(**kwargs)
 141. 
 142.     def get_publisher(self):
 143.         """Get a celery task message publisher."""
 144.         return TaskPublisher(connection=DjangoAMQPConnection)
 145. 
 146.     def get_consumer(self):
 147.         """Get a celery task message consumer."""
 148.         return TaskConsumer(connection=DjangoAMQPConnection)
 149. 
 150.     @classmethod
 151.     def delay(cls, **kwargs):
 152.         """Delay this task for execution by the ``celery`` daemon(s)."""
 153.         return delay_task(cls.name, **kwargs)
 154. 
 155. 
 156. class TaskSet(object):
 157.     """A task containing several subtasks, making it possible
 158.     to track how many, or when all of the tasks are completed.
 159. 
 160.     Example Usage
 161.     --------------
 162. 
 163.         >>> from djangofeeds.tasks import RefreshFeedTask
 164.         >>> taskset = TaskSet(RefreshFeedTask, args=[
 165.         ...                 {"feed_url": "http://cnn.com/rss"},
 166.         ...                 {"feed_url": "http://bbc.com/rss"},
 167.         ...                 {"feed_url": "http://xkcd.com/rss"}])
 168. 
 169.         >>> taskset_id, subtask_ids = taskset.run()
 170. 
 171. 
 172.     """
 173. 
 174.     def __init__(self, task, args):
 175.         """``task`` can be either a fully qualified task name, or a task
 176.         class, args is a list of arguments for the subtasks.
 177.         """
 178. 
 179.         try:
 180.             task_name = task.name
 181.         except AttributeError:
 182.             task_name = task
 183. 
 184.         self.task_name = task_name
 185.         self.arguments = args
 186.         self.total = len(args)
 187. 
 188.     def run(self):
 189.         """Run all tasks in the taskset.
 190. 
 191.         Returns a tuple with the taskset id, and a list of subtask id's.
 192. 
 193.         Examples
 194.         --------
 195.             >>> ts = RefreshFeeds(["http://foo.com/rss", http://bar.com/rss"])
 196.             >>> taskset_id, subtask_ids = ts.run()
 197.             >>> taskset_id
 198.             "d2c9b261-8eff-4bfb-8459-1e1b72063514"
 199.             >>> subtask_ids
 200.             ["b4996460-d959-49c8-aeb9-39c530dcde25",
 201.             "598d2d18-ab86-45ca-8b4f-0779f5d6a3cb"]
 202.             >>> time.sleep(10)
 203.             >>> is_done(taskset_id)
 204.             True
 205.         """
 206.         taskset_id = str(uuid.uuid4())
 207.         publisher = TaskPublisher(connection=DjangoAMQPConnection)
 208.         subtask_ids = []
 209.         for arg in self.arguments:
 210.             subtask_id = publisher.delay_task_in_set(task_name=self.task_name,
 211.                                                      taskset_id=taskset_id,
 212.                                                      task_kwargs=arg)
 213.             subtask_ids.append(subtask_id)
 214.         publisher.close()
 215.         return taskset_id, subtask_ids
 216. 
 217. 
 218. class PeriodicTask(Task):
 219.     """A periodic task is a task that behaves like a cron job.
 220. 
 221.     The ``run_every`` attribute defines how often the task is run (its
 222.     interval), it can be either a ``datetime.timedelta`` object or a integer
 223.     specifying the time in seconds.
 224. 
 225.     You have to register the periodic task in the task registry.
 226. 
 227.     Examples
 228.     --------
 229. 
 230.         >>> from celery.task import tasks, PeriodicTask
 231.         >>> from datetime import timedelta
 232.         >>> class MyPeriodicTask(PeriodicTask):
 233.         ...     name = "my_periodic_task"
 234.         ...     run_every = timedelta(seconds=30)
 235.         ...
 236.         ...     def run(self, **kwargs):
 237.         ...         logger = self.get_logger(**kwargs)
 238.         ...         logger.info("Running MyPeriodicTask")
 239.         >>> tasks.register(MyPeriodicTask)
 240. 
 241.     """
 242.     run_every = timedelta(days=1)
 243.     type = "periodic"
 244. 
 245.     def __init__(self):
 246.         if not self.run_every:
 247.             raise NotImplementedError(
 248.                     "Periodic tasks must have a run_every attribute")
 249. 
 250.         # If run_every is a integer, convert it to timedelta seconds.
 251.         if isinstance(self.run_every, int):
 252.             self.run_every = timedelta(seconds=self.run_every)
 253. 
 254.         super(PeriodicTask, self).__init__()
 255. 
 256. 
 257. class DeleteExpiredTaskMetaTask(PeriodicTask):
 258.     """A periodic task that deletes expired task metadata every day.
 259. 
 260.     It's only registered if ``settings.CELERY_TASK_META_USE_DB`` is set.
 261.     """
 262.     name = "celery.delete_expired_task_meta"
 263.     run_every = timedelta(days=1)
 264. 
 265.     def run(self, **kwargs):
 266.         logger = self.get_logger(**kwargs)
 267.         logger.info("Deleting expired task meta objects...")
 268.         TaskMeta.objects.delete_expired()
 269. if TASK_META_USE_DB:
 270.     tasks.register(DeleteExpiredTaskMetaTask)