1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
|
import pickle
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 rethinkdb import RethinkDB
except ImportError as exc: # pragma: nocover
raise ImportError("RethinkDBJobStore requires rethinkdb installed") from exc
class RethinkDBJobStore(BaseJobStore):
"""
Stores jobs in a RethinkDB database. Any leftover keyword arguments are directly passed to
rethinkdb's `RethinkdbClient <http://www.rethinkdb.com/api/#connect>`_.
Plugin alias: ``rethinkdb``
:param str database: database to store jobs in
:param str collection: collection to store jobs in
:param client: a :class:`rethinkdb.net.Connection` 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",
table="jobs",
client=None,
pickle_protocol=pickle.HIGHEST_PROTOCOL,
**connect_args,
):
super().__init__()
if not database:
raise ValueError('The "database" parameter must not be empty')
if not table:
raise ValueError('The "table" parameter must not be empty')
self.database = database
self.table_name = table
self.table = None
self.client = client
self.pickle_protocol = pickle_protocol
self.connect_args = connect_args
self.r = RethinkDB()
self.conn = None
def start(self, scheduler, alias):
super().start(scheduler, alias)
if self.client:
self.conn = maybe_ref(self.client)
else:
self.conn = self.r.connect(db=self.database, **self.connect_args)
if self.database not in self.r.db_list().run(self.conn):
self.r.db_create(self.database).run(self.conn)
if self.table_name not in self.r.table_list().run(self.conn):
self.r.table_create(self.table_name).run(self.conn)
if "next_run_time" not in self.r.table(self.table_name).index_list().run(
self.conn
):
self.r.table(self.table_name).index_create("next_run_time").run(self.conn)
self.table = self.r.db(self.database).table(self.table_name)
def lookup_job(self, job_id):
results = list(self.table.get_all(job_id).pluck("job_state").run(self.conn))
return self._reconstitute_job(results[0]["job_state"]) if results else None
def get_due_jobs(self, now):
return self._get_jobs(
self.r.row["next_run_time"] <= datetime_to_utc_timestamp(now)
)
def get_next_run_time(self):
results = list(
self.table.filter(self.r.row["next_run_time"] != None)
.order_by(self.r.asc("next_run_time"))
.map(lambda x: x["next_run_time"])
.limit(1)
.run(self.conn)
)
return utc_timestamp_to_datetime(results[0]) if results else None
def get_all_jobs(self):
jobs = self._get_jobs()
self._fix_paused_jobs_sorting(jobs)
return jobs
def add_job(self, job):
job_dict = {
"id": job.id,
"next_run_time": datetime_to_utc_timestamp(job.next_run_time),
"job_state": self.r.binary(
pickle.dumps(job.__getstate__(), self.pickle_protocol)
),
}
results = self.table.insert(job_dict).run(self.conn)
if results["errors"] > 0:
raise ConflictingIdError(job.id)
def update_job(self, job):
changes = {
"next_run_time": datetime_to_utc_timestamp(job.next_run_time),
"job_state": self.r.binary(
pickle.dumps(job.__getstate__(), self.pickle_protocol)
),
}
results = self.table.get_all(job.id).update(changes).run(self.conn)
skipped = False in map(lambda x: results[x] == 0, results.keys())
if results["skipped"] > 0 or results["errors"] > 0 or not skipped:
raise JobLookupError(job.id)
def remove_job(self, job_id):
results = self.table.get_all(job_id).delete().run(self.conn)
if results["deleted"] + results["skipped"] != 1:
raise JobLookupError(job_id)
def remove_all_jobs(self):
self.table.delete().run(self.conn)
def shutdown(self):
self.conn.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, predicate=None):
jobs = []
failed_job_ids = []
query = (
self.table.filter(self.r.row["next_run_time"] != None).filter(predicate)
if predicate
else self.table
)
query = query.order_by("next_run_time", "id").pluck("id", "job_state")
for document in query.run(self.conn):
try:
jobs.append(self._reconstitute_job(document["job_state"]))
except Exception:
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.r.expr(failed_job_ids).for_each(
lambda job_id: self.table.get_all(job_id).delete()
).run(self.conn)
return jobs
def __repr__(self):
connection = self.conn
return f"<{self.__class__.__name__} (connection={connection})>"
|
