redis.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150
  1. from __future__ import absolute_import
  2. from datetime import datetime
  3. from pytz import utc
  4. import six
  5. from apscheduler.jobstores.base import BaseJobStore, JobLookupError, ConflictingIdError
  6. from apscheduler.util import datetime_to_utc_timestamp, utc_timestamp_to_datetime
  7. from apscheduler.job import Job
  8. try:
  9. import cPickle as pickle
  10. except ImportError: # pragma: nocover
  11. import pickle
  12. try:
  13. from redis import Redis
  14. except ImportError: # pragma: nocover
  15. raise ImportError('RedisJobStore requires redis installed')
  16. class RedisJobStore(BaseJobStore):
  17. """
  18. Stores jobs in a Redis database. Any leftover keyword arguments are directly passed to redis's
  19. :class:`~redis.StrictRedis`.
  20. Plugin alias: ``redis``
  21. :param int db: the database number to store jobs in
  22. :param str jobs_key: key to store jobs in
  23. :param str run_times_key: key to store the jobs' run times in
  24. :param int pickle_protocol: pickle protocol level to use (for serialization), defaults to the
  25. highest available
  26. """
  27. def __init__(self, db=0, jobs_key='apscheduler.jobs', run_times_key='apscheduler.run_times',
  28. pickle_protocol=pickle.HIGHEST_PROTOCOL, **connect_args):
  29. super(RedisJobStore, self).__init__()
  30. if db is None:
  31. raise ValueError('The "db" parameter must not be empty')
  32. if not jobs_key:
  33. raise ValueError('The "jobs_key" parameter must not be empty')
  34. if not run_times_key:
  35. raise ValueError('The "run_times_key" parameter must not be empty')
  36. self.pickle_protocol = pickle_protocol
  37. self.jobs_key = jobs_key
  38. self.run_times_key = run_times_key
  39. self.redis = Redis(db=int(db), **connect_args)
  40. def lookup_job(self, job_id):
  41. job_state = self.redis.hget(self.jobs_key, job_id)
  42. return self._reconstitute_job(job_state) if job_state else None
  43. def get_due_jobs(self, now):
  44. timestamp = datetime_to_utc_timestamp(now)
  45. job_ids = self.redis.zrangebyscore(self.run_times_key, 0, timestamp)
  46. if job_ids:
  47. job_states = self.redis.hmget(self.jobs_key, *job_ids)
  48. return self._reconstitute_jobs(six.moves.zip(job_ids, job_states))
  49. return []
  50. def get_next_run_time(self):
  51. next_run_time = self.redis.zrange(self.run_times_key, 0, 0, withscores=True)
  52. if next_run_time:
  53. return utc_timestamp_to_datetime(next_run_time[0][1])
  54. def get_all_jobs(self):
  55. job_states = self.redis.hgetall(self.jobs_key)
  56. jobs = self._reconstitute_jobs(six.iteritems(job_states))
  57. paused_sort_key = datetime(9999, 12, 31, tzinfo=utc)
  58. return sorted(jobs, key=lambda job: job.next_run_time or paused_sort_key)
  59. def add_job(self, job):
  60. if self.redis.hexists(self.jobs_key, job.id):
  61. raise ConflictingIdError(job.id)
  62. with self.redis.pipeline() as pipe:
  63. pipe.multi()
  64. pipe.hset(self.jobs_key, job.id, pickle.dumps(job.__getstate__(),
  65. self.pickle_protocol))
  66. if job.next_run_time:
  67. pipe.zadd(self.run_times_key,
  68. {job.id: datetime_to_utc_timestamp(job.next_run_time)})
  69. pipe.execute()
  70. def update_job(self, job):
  71. if not self.redis.hexists(self.jobs_key, job.id):
  72. raise JobLookupError(job.id)
  73. with self.redis.pipeline() as pipe:
  74. pipe.hset(self.jobs_key, job.id, pickle.dumps(job.__getstate__(),
  75. self.pickle_protocol))
  76. if job.next_run_time:
  77. pipe.zadd(self.run_times_key,
  78. {job.id: datetime_to_utc_timestamp(job.next_run_time)})
  79. else:
  80. pipe.zrem(self.run_times_key, job.id)
  81. pipe.execute()
  82. def remove_job(self, job_id):
  83. if not self.redis.hexists(self.jobs_key, job_id):
  84. raise JobLookupError(job_id)
  85. with self.redis.pipeline() as pipe:
  86. pipe.hdel(self.jobs_key, job_id)
  87. pipe.zrem(self.run_times_key, job_id)
  88. pipe.execute()
  89. def remove_all_jobs(self):
  90. with self.redis.pipeline() as pipe:
  91. pipe.delete(self.jobs_key)
  92. pipe.delete(self.run_times_key)
  93. pipe.execute()
  94. def shutdown(self):
  95. self.redis.connection_pool.disconnect()
  96. def _reconstitute_job(self, job_state):
  97. job_state = pickle.loads(job_state)
  98. job = Job.__new__(Job)
  99. job.__setstate__(job_state)
  100. job._scheduler = self._scheduler
  101. job._jobstore_alias = self._alias
  102. return job
  103. def _reconstitute_jobs(self, job_states):
  104. jobs = []
  105. failed_job_ids = []
  106. for job_id, job_state in job_states:
  107. try:
  108. jobs.append(self._reconstitute_job(job_state))
  109. except BaseException:
  110. self._logger.exception('Unable to restore job "%s" -- removing it', job_id)
  111. failed_job_ids.append(job_id)
  112. # Remove all the jobs we failed to restore
  113. if failed_job_ids:
  114. with self.redis.pipeline() as pipe:
  115. pipe.hdel(self.jobs_key, *failed_job_ids)
  116. pipe.zrem(self.run_times_key, *failed_job_ids)
  117. pipe.execute()
  118. return jobs
  119. def __repr__(self):
  120. return '<%s>' % self.__class__.__name__