You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
190 lines
6.8 KiB
190 lines
6.8 KiB
4 years ago
|
import asyncio
|
||
|
|
|
||
|
|
from apscheduler.jobstores.base import ConflictingIdError, JobLookupError
|
||
|
|
from apscheduler.jobstores.memory import MemoryJobStore
|
||
|
|
from apscheduler.util import datetime_to_utc_timestamp
|
||
|
|
from redbot.core import Config
|
||
|
|
|
||
|
|
|
||
|
|
# TODO: use get_lock on config
|
||
|
|
from redbot.core.bot import Red
|
||
|
|
|
||
|
|
|
||
|
|
|
||
|
|
class RedConfigJobStore(MemoryJobStore):
|
||
|
|
def __init__(self, config: Config, bot: Red):
|
||
|
|
super().__init__()
|
||
|
|
self.config = config
|
||
|
|
# nest_asyncio.apply()
|
||
|
|
self.bot = bot
|
||
|
|
asyncio.ensure_future(self._load_from_config(), loop=self.bot.loop)
|
||
|
|
|
||
|
|
async def _load_from_config(self):
|
||
|
|
self._jobs = await self.config.jobs()
|
||
|
|
self._jobs_index = await self.config.jobs_index.all()
|
||
|
|
|
||
|
|
def add_job(self, job):
|
||
|
|
if job.id in self._jobs_index:
|
||
|
|
raise ConflictingIdError(job.id)
|
||
|
|
|
||
|
|
timestamp = datetime_to_utc_timestamp(job.next_run_time)
|
||
|
|
index = self._get_job_index(timestamp, job.id) # This is fine
|
||
|
|
self._jobs.insert(index, (job, timestamp))
|
||
|
|
self._jobs_index[job.id] = (job, timestamp)
|
||
|
|
asyncio.create_task(self._async_add_job(job, index, timestamp))
|
||
|
|
|
||
|
|
async def _async_add_job(self, job, index, timestamp):
|
||
|
|
async with self.config.jobs() as jobs:
|
||
|
|
jobs.insert(index, (job, timestamp))
|
||
|
|
await self.config.jobs_index.set_raw(job.id, value=(job, timestamp))
|
||
|
|
return True
|
||
|
|
|
||
|
|
def update_job(self, job):
|
||
|
|
old_job, old_timestamp = self._jobs_index.get(job.id, (None, None))
|
||
|
|
if old_job is None:
|
||
|
|
raise JobLookupError(job.id)
|
||
|
|
|
||
|
|
# If the next run time has not changed, simply replace the job in its present index.
|
||
|
|
# Otherwise, reinsert the job to the list to preserve the ordering.
|
||
|
|
old_index = self._get_job_index(old_timestamp, old_job.id)
|
||
|
|
new_timestamp = datetime_to_utc_timestamp(job.next_run_time)
|
||
|
|
asyncio.ensure_future(self._async_update_job(job, new_timestamp, old_index, old_job, old_timestamp), loop=self.bot.loop)
|
||
|
|
|
||
|
|
async def _async_update_job(self, job, new_timestamp, old_index, old_job, old_timestamp):
|
||
|
|
if old_timestamp == new_timestamp:
|
||
|
|
self._jobs[old_index] = (job, new_timestamp)
|
||
|
|
async with self.config.jobs() as jobs:
|
||
|
|
jobs[old_index] = (job, new_timestamp)
|
||
|
|
else:
|
||
|
|
del self._jobs[old_index]
|
||
|
|
new_index = self._get_job_index(new_timestamp, job.id) # This is fine
|
||
|
|
self._jobs.insert(new_index, (job, new_timestamp))
|
||
|
|
async with self.config.jobs() as jobs:
|
||
|
|
del jobs[old_index]
|
||
|
|
jobs.insert(new_index, (job, new_timestamp))
|
||
|
|
self._jobs_index[old_job.id] = (job, new_timestamp)
|
||
|
|
await self.config.jobs_index.set_raw(old_job.id, value=(job, new_timestamp))
|
||
|
|
|
||
|
|
def remove_job(self, job_id):
|
||
|
|
job, timestamp = self._jobs_index.get(job_id, (None, None))
|
||
|
|
if job is None:
|
||
|
|
raise JobLookupError(job_id)
|
||
|
|
|
||
|
|
index = self._get_job_index(timestamp, job_id)
|
||
|
|
del self._jobs[index]
|
||
|
|
del self._jobs_index[job.id]
|
||
|
|
asyncio.create_task(self._async_remove_job(index, job))
|
||
|
|
|
||
|
|
async def _async_remove_job(self, index, job):
|
||
|
|
async with self.config.jobs() as jobs:
|
||
|
|
del jobs[index]
|
||
|
|
await self.config.jobs_index.clear_raw(job.id)
|
||
|
|
|
||
|
|
def remove_all_jobs(self):
|
||
|
|
super().remove_all_jobs()
|
||
|
|
asyncio.create_task(self._async_remove_all_jobs())
|
||
|
|
|
||
|
|
async def _async_remove_all_jobs(self):
|
||
|
|
await self.config.jobs.clear()
|
||
|
|
await self.config.jobs_index.clear()
|
||
|
|
|
||
|
|
|
||
|
|
# import asyncio
|
||
|
|
#
|
||
|
|
# from apscheduler.jobstores.base import BaseJobStore, ConflictingIdError
|
||
|
|
# from apscheduler.util import datetime_to_utc_timestamp
|
||
|
|
# from redbot.core import Config
|
||
|
|
# from redbot.core.utils import AsyncIter
|
||
|
|
#
|
||
|
|
#
|
||
|
|
# class RedConfigJobStore(BaseJobStore):
|
||
|
|
# def __init__(self, config: Config, loop):
|
||
|
|
# super().__init__()
|
||
|
|
# self.config = config
|
||
|
|
# self.loop: asyncio.BaseEventLoop = loop
|
||
|
|
#
|
||
|
|
# self._jobs = []
|
||
|
|
# self._jobs_index = {} # id -> (job, timestamp) lookup table
|
||
|
|
#
|
||
|
|
# def lookup_job(self, job_id):
|
||
|
|
# return asyncio.run(self._async_lookup_job(job_id))
|
||
|
|
#
|
||
|
|
# async def _async_lookup_job(self, job_id):
|
||
|
|
# return (await self.config.jobs_index.get_raw(job_id, default=(None, None)))[0]
|
||
|
|
#
|
||
|
|
# def get_due_jobs(self, now):
|
||
|
|
# return asyncio.run(self._async_get_due_jobs(now))
|
||
|
|
#
|
||
|
|
# async def _async_get_due_jobs(self, now):
|
||
|
|
# now_timestamp = datetime_to_utc_timestamp(now)
|
||
|
|
# pending = []
|
||
|
|
# all_jobs = await self.config.jobs()
|
||
|
|
# async for job, timestamp in AsyncIter(all_jobs, steps=100):
|
||
|
|
# if timestamp is None or timestamp > now_timestamp:
|
||
|
|
# break
|
||
|
|
# pending.append(job)
|
||
|
|
#
|
||
|
|
# return pending
|
||
|
|
#
|
||
|
|
# def get_next_run_time(self):
|
||
|
|
# return asyncio.run(self._async_get_next_run_time())
|
||
|
|
#
|
||
|
|
# async def _async_get_next_run_time(self):
|
||
|
|
# _jobs = await self.config.jobs()
|
||
|
|
# return _jobs[0][0].next_run_time if _jobs else None
|
||
|
|
#
|
||
|
|
# def get_all_jobs(self):
|
||
|
|
# return asyncio.run(self._async_get_all_jobs())
|
||
|
|
#
|
||
|
|
# async def _async_get_all_jobs(self):
|
||
|
|
# return [j[0] for j in (await self.config.jobs())]
|
||
|
|
#
|
||
|
|
# def add_job(self, job):
|
||
|
|
# return asyncio.run(self._async_add_job(job))
|
||
|
|
#
|
||
|
|
# async def _async_add_job(self, job):
|
||
|
|
# if await self.config.jobs_index.get_raw(job.id, default=None) is not None:
|
||
|
|
# raise ConflictingIdError(job.id)
|
||
|
|
#
|
||
|
|
# timestamp = datetime_to_utc_timestamp(job.next_run_time)
|
||
|
|
# index = self._get_job_index(timestamp, job.id)
|
||
|
|
# self._jobs.insert(index, (job, timestamp))
|
||
|
|
# self._jobs_index[job.id] = (job, timestamp)
|
||
|
|
#
|
||
|
|
# def update_job(self, job):
|
||
|
|
# pass
|
||
|
|
#
|
||
|
|
# def remove_job(self, job_id):
|
||
|
|
# pass
|
||
|
|
#
|
||
|
|
# def remove_all_jobs(self):
|
||
|
|
# pass
|
||
|
|
#
|
||
|
|
# def _get_job_index(self, timestamp, job_id):
|
||
|
|
# """
|
||
|
|
# Returns the index of the given job, or if it's not found, the index where the job should be
|
||
|
|
# inserted based on the given timestamp.
|
||
|
|
#
|
||
|
|
# :type timestamp: int
|
||
|
|
# :type job_id: str
|
||
|
|
#
|
||
|
|
# """
|
||
|
|
# lo, hi = 0, len(self._jobs)
|
||
|
|
# timestamp = float('inf') if timestamp is None else timestamp
|
||
|
|
# while lo < hi:
|
||
|
|
# mid = (lo + hi) // 2
|
||
|
|
# mid_job, mid_timestamp = self._jobs[mid]
|
||
|
|
# mid_timestamp = float('inf') if mid_timestamp is None else mid_timestamp
|
||
|
|
# if mid_timestamp > timestamp:
|
||
|
|
# hi = mid
|
||
|
|
# elif mid_timestamp < timestamp:
|
||
|
|
# lo = mid + 1
|
||
|
|
# elif mid_job.id > job_id:
|
||
|
|
# hi = mid
|
||
|
|
# elif mid_job.id < job_id:
|
||
|
|
# lo = mid + 1
|
||
|
|
# else:
|
||
|
|
# return mid
|
||
|
|
#
|
||
|
|
# return lo
|