def inmem_queue():
with connection() as conn:
e = Worker(queues="pytest", connection=conn)
c = Queue(queue="pytest", connection=conn)
c.e = e
c.storage.clear(force=True)
yield c
e.shutdown()
def inmem_queue():
with tempfile.NamedTemporaryFile() as f:
connection = create_engine(
"sqlite:///{path}".format(path=f.name),
connect_args={"check_same_thread": False},
poolclass=NullPool,
)
e = Worker(queues="pytest", connection=connection)
c = Queue(queue="pytest", connection=connection)
c.e = e
yield c
e.shutdown()
def __init__(self, queue=None, connection=None):
if connection is None and not isinstance(queue, Queue):
raise ValueError(
"One of either connection or queue must be specified")
elif isinstance(queue, Queue):
self.queue = queue
if connection is None:
connection = self.queue.storage.engine
elif connection:
self.queue = Queue(connection=connection)
self._schedule_checker = None
super(Scheduler, self).__init__(connection, Base=Base)
def inmem_queue():
fd, filepath = tempfile.mkstemp()
connection = create_engine(
"sqlite:///{path}".format(path=filepath),
connect_args={"check_same_thread": False},
poolclass=NullPool,
)
e = Worker(queues="pytest", connection=connection)
c = Queue(queue="pytest", connection=connection)
c.e = e
yield c
e.shutdown()
os.close(fd)
try:
os.remove(filepath)
except OSError:
pass
def queue():
with tempfile.NamedTemporaryFile() as f:
connection = create_engine(
"sqlite:///{path}".format(path=f.name),
connect_args={"check_same_thread": False},
poolclass=NullPool,
)
q = Queue("pytest", connection)
yield q
def queue():
fd, filepath = tempfile.mkstemp()
connection = create_engine(
"sqlite:///{path}".format(path=filepath),
connect_args={"check_same_thread": False},
poolclass=NullPool,
)
q = Queue("pytest", connection)
yield q
os.close(fd)
os.remove(filepath)
def __queue():
return Queue(task_queue_name, connection=connection)
def __priority_queue():
return Queue(priority_queue_name, connection=connection)
def queue_no_worker():
with connection() as conn:
c = Queue(queue="pytest", connection=conn)
c.storage.clear(force=True)
yield c
# Add multiprocessing safeguards as recommended by
# https://docs.sqlalchemy.org/en/13/core/pooling.html#using-connection-pools-with-multiprocessing
@event.listens_for(connection, "connect")
def connect(dbapi_connection, connection_record):
connection_record.info["pid"] = os.getpid()
@event.listens_for(connection, "checkout")
def checkout(dbapi_connection, connection_record, connection_proxy):
pid = os.getpid()
if connection_record.info["pid"] != pid:
connection_record.connection = connection_proxy.connection = None
raise exc.DisconnectionError(
"Connection record belongs to pid %s, attempting to check out in pid %s"
% (connection_record.info["pid"], pid)
)
queue = Queue(app, connection=connection)
scheduler = Scheduler(queue=queue, connection=connection)
def initialize_worker():
worker = Worker(app, connection=connection, num_workers=1)
atexit.register(worker.shutdown)
def __facility_queue():
return Queue(facility_queue_name, connection=connection)
@event.listens_for(connection, "checkout")
def checkout(dbapi_connection, connection_record, connection_proxy):
pid = os.getpid()
if connection_record.info["pid"] != pid:
connection_record.connection = connection_proxy.connection = None
raise exc.DisconnectionError(
"Connection record belongs to pid %s, attempting to check out in pid %s"
% (connection_record.info["pid"], pid))
task_queue_name = "kolibri"
priority_queue_name = "no_waiting"
priority_queue = Queue(priority_queue_name, connection=connection)
queue = Queue(task_queue_name, connection=connection)
scheduler = Scheduler(queue=queue, connection=connection)
def initialize_workers():
regular_worker = Worker(task_queue_name,
connection=connection,
num_workers=1)
priority_worker = Worker(priority_queue_name,
connection=connection,
num_workers=3)
return regular_worker, priority_worker
class Scheduler(StorageMixin):
def __init__(self, queue=None, connection=None):
if connection is None and not isinstance(queue, Queue):
raise ValueError(
"One of either connection or queue must be specified")
elif isinstance(queue, Queue):
self.queue = queue
if connection is None:
connection = self.queue.storage.engine
elif connection:
self.queue = Queue(connection=connection)
self._schedule_checker = None
super(Scheduler, self).__init__(connection, Base=Base)
def __contains__(self, item):
"""
Returns a boolean indicating whether the given job instance or job id
is scheduled for execution.
"""
job_id = item
if isinstance(item, Job):
job_id = item.job_id
with self.session_scope() as session:
return session.query(
self._ns_query(session).filter_by(
id=job_id).exists()).scalar()
def change_execution_time(self, job, date_time):
"""
Change a job's execution time.
"""
if date_time.tzinfo is None:
raise ValueError(
"Must use a timezone aware datetime object for scheduling tasks"
)
with self.session_scope() as session:
scheduled_job = (session.query(ScheduledJob).filter_by(
id=job.job_id).one_or_none())
if scheduled_job:
scheduled_job.scheduled_time = naive_utc_datetime(date_time)
session.merge(scheduled_job)
else:
raise ValueError("Job not in scheduled jobs queue")
def start_schedule_checker(self):
"""
Starts up the job checker thread, that starts scheduled jobs when there are workers free,
and checks for cancellation requests for jobs currently assigned to a worker.
Returns: the Thread object.
"""
t = InfiniteLoopThread(self.check_schedule,
thread_name="SCHEDULECHECKER",
wait_between_runs=0.5)
t.start()
return t
def run(self):
"""
Start the schedule checker in a blocking way to be parallel to the rq-scheduler method.
"""
thread = self.start_schedule_checker()
thread.join()
def start_scheduler(self):
if not (self._schedule_checker and self._schedule_checker.is_alive()):
self._schedule_checker = self.start_schedule_checker()
def shutdown_scheduler(self):
if self._schedule_checker:
self._schedule_checker.stop()
def enqueue_at(self, dt, func, interval=0, repeat=0, *args, **kwargs):
"""
Add the job to the scheduler for the specified time
"""
return self.schedule(dt,
func,
interval=interval,
repeat=repeat,
*args,
**kwargs)
def enqueue_in(self, delta_t, func, interval=0, repeat=0, *args, **kwargs):
"""
Add the job to the scheduler in the specified time delta
"""
if not isinstance(delta_t, timedelta):
raise ValueError("Time argument must be a timedelta object.")
dt = self._now() + delta_t
return self.schedule(dt,
func,
interval=interval,
repeat=repeat,
*args,
**kwargs)
def schedule(self, dt, func, interval=0, repeat=0, *args, **kwargs):
"""
Add the job to the scheduler for the specified time, interval, and number of repeats.
Repeat of None with a specified interval means the job will repeat forever at that
interval.
"""
if not isinstance(dt, datetime):
raise ValueError("Time argument must be a datetime object.")
if not interval and repeat != 0:
raise ValueError(
"Must specify an interval if the task is repeating")
if dt.tzinfo is None:
raise ValueError(
"Must use a timezone aware datetime object for scheduling tasks"
)
if isinstance(func, Job):
job = func
# else, turn it into a job first.
else:
job = Job(func, *args, **kwargs)
job.state = State.SCHEDULED
with self.session_scope() as session:
scheduled_job = ScheduledJob(
id=job.job_id,
queue=self.queue.name,
interval=interval,
repeat=repeat,
scheduled_time=naive_utc_datetime(dt),
obj=job,
)
session.merge(scheduled_job)
return job.job_id
def get_jobs(self):
with self.session_scope() as s:
scheduled_jobs = self._ns_query(s).all()
return [o.obj for o in scheduled_jobs]
def count(self):
with self.session_scope() as s:
return self._ns_query(s).count()
def get_job(self, job_id):
with self.session_scope() as session:
scheduled_job = session.query(ScheduledJob).get(job_id)
if scheduled_job is None:
raise JobNotFound()
return scheduled_job.obj
def cancel(self, job_id):
"""
Clear a scheduled job.
:type job_id: NoneType or str
:param job_id: the job_id to clear. If None, clear all jobs.
"""
with self.session_scope() as s:
q = self._ns_query(s)
if job_id:
q = q.filter_by(id=job_id)
q.delete(synchronize_session=False)
def clear_scheduler(self):
"""
Clear all scheduled jobs
"""
self.cancel(None)
def check_schedule(self):
start = time.time()
naive_utc_now = datetime.utcnow()
with self.session_scope() as s:
scheduled_jobs = self._ns_query(s).filter(
ScheduledJob.scheduled_time <= naive_utc_now)
for scheduled_job in scheduled_jobs:
new_repeat = 0
repeat = False
if scheduled_job.repeat is None:
new_repeat = None
repeat = True
elif scheduled_job.repeat > 0:
new_repeat = scheduled_job.repeat - 1
repeat = True
job_for_queue = scheduled_job.obj
self.queue.enqueue(job_for_queue)
if repeat:
# Update this scheduled job to repeat this
self.schedule(
self._now() +
timedelta(seconds=scheduled_job.interval),
job_for_queue,
interval=scheduled_job.interval,
repeat=new_repeat,
)
else:
s.delete(scheduled_job)
return time.time() - start
def _ns_query(self, session):
"""
Return a SQLAlchemy query that is already namespaced by the queue.
Returns: a SQLAlchemy query object
"""
return session.query(ScheduledJob).filter(
ScheduledJob.queue == self.queue.name)
def _now(self):
return local_now()
def queue():
with connection() as c:
q = Queue("pytest", c)
yield q
