diff options
Diffstat (limited to '.venv/lib/python3.12/site-packages/apscheduler/jobstores/mongodb.py')
-rw-r--r-- | .venv/lib/python3.12/site-packages/apscheduler/jobstores/mongodb.py | 158 |
1 files changed, 158 insertions, 0 deletions
diff --git a/.venv/lib/python3.12/site-packages/apscheduler/jobstores/mongodb.py b/.venv/lib/python3.12/site-packages/apscheduler/jobstores/mongodb.py new file mode 100644 index 00000000..102c0bd0 --- /dev/null +++ b/.venv/lib/python3.12/site-packages/apscheduler/jobstores/mongodb.py @@ -0,0 +1,158 @@ +import pickle +import warnings + +from apscheduler.job import Job +from apscheduler.jobstores.base import BaseJobStore, ConflictingIdError, JobLookupError +from apscheduler.util import ( + datetime_to_utc_timestamp, + maybe_ref, + utc_timestamp_to_datetime, +) + +try: + from bson.binary import Binary + from pymongo import ASCENDING, MongoClient + from pymongo.errors import DuplicateKeyError +except ImportError as exc: # pragma: nocover + raise ImportError("MongoDBJobStore requires PyMongo installed") from exc + + +class MongoDBJobStore(BaseJobStore): + """ + Stores jobs in a MongoDB database. Any leftover keyword arguments are directly passed to + pymongo's `MongoClient + <http://api.mongodb.org/python/current/api/pymongo/mongo_client.html#pymongo.mongo_client.MongoClient>`_. + + Plugin alias: ``mongodb`` + + :param str database: database to store jobs in + :param str collection: collection to store jobs in + :param client: a :class:`~pymongo.mongo_client.MongoClient` instance to use instead of + providing connection arguments + :param int pickle_protocol: pickle protocol level to use (for serialization), defaults to the + highest available + """ + + def __init__( + self, + database="apscheduler", + collection="jobs", + client=None, + pickle_protocol=pickle.HIGHEST_PROTOCOL, + **connect_args, + ): + super().__init__() + self.pickle_protocol = pickle_protocol + + if not database: + raise ValueError('The "database" parameter must not be empty') + if not collection: + raise ValueError('The "collection" parameter must not be empty') + + if client: + self.client = maybe_ref(client) + else: + connect_args.setdefault("w", 1) + self.client = MongoClient(**connect_args) + + self.collection = self.client[database][collection] + + def start(self, scheduler, alias): + super().start(scheduler, alias) + self.collection.create_index("next_run_time", sparse=True) + + @property + def connection(self): + warnings.warn( + 'The "connection" member is deprecated -- use "client" instead', + DeprecationWarning, + ) + return self.client + + def lookup_job(self, job_id): + document = self.collection.find_one(job_id, ["job_state"]) + return self._reconstitute_job(document["job_state"]) if document else None + + def get_due_jobs(self, now): + timestamp = datetime_to_utc_timestamp(now) + return self._get_jobs({"next_run_time": {"$lte": timestamp}}) + + def get_next_run_time(self): + document = self.collection.find_one( + {"next_run_time": {"$ne": None}}, + projection=["next_run_time"], + sort=[("next_run_time", ASCENDING)], + ) + return ( + utc_timestamp_to_datetime(document["next_run_time"]) if document else None + ) + + def get_all_jobs(self): + jobs = self._get_jobs({}) + self._fix_paused_jobs_sorting(jobs) + return jobs + + def add_job(self, job): + try: + self.collection.insert_one( + { + "_id": job.id, + "next_run_time": datetime_to_utc_timestamp(job.next_run_time), + "job_state": Binary( + pickle.dumps(job.__getstate__(), self.pickle_protocol) + ), + } + ) + except DuplicateKeyError: + raise ConflictingIdError(job.id) + + def update_job(self, job): + changes = { + "next_run_time": datetime_to_utc_timestamp(job.next_run_time), + "job_state": Binary(pickle.dumps(job.__getstate__(), self.pickle_protocol)), + } + result = self.collection.update_one({"_id": job.id}, {"$set": changes}) + if result and result.matched_count == 0: + raise JobLookupError(job.id) + + def remove_job(self, job_id): + result = self.collection.delete_one({"_id": job_id}) + if result and result.deleted_count == 0: + raise JobLookupError(job_id) + + def remove_all_jobs(self): + self.collection.delete_many({}) + + def shutdown(self): + self.client.close() + + def _reconstitute_job(self, job_state): + job_state = pickle.loads(job_state) + job = Job.__new__(Job) + job.__setstate__(job_state) + job._scheduler = self._scheduler + job._jobstore_alias = self._alias + return job + + def _get_jobs(self, conditions): + jobs = [] + failed_job_ids = [] + for document in self.collection.find( + conditions, ["_id", "job_state"], sort=[("next_run_time", ASCENDING)] + ): + try: + jobs.append(self._reconstitute_job(document["job_state"])) + except BaseException: + self._logger.exception( + 'Unable to restore job "%s" -- removing it', document["_id"] + ) + failed_job_ids.append(document["_id"]) + + # Remove all the jobs we failed to restore + if failed_job_ids: + self.collection.delete_many({"_id": {"$in": failed_job_ids}}) + + return jobs + + def __repr__(self): + return f"<{self.__class__.__name__} (client={self.client})>" |