kopia lustrzana https://github.com/jaseg/gerbolyze
78 wiersze
2.9 KiB
Python
78 wiersze
2.9 KiB
Python
|
|
import json
|
|
import sqlite3
|
|
|
|
class JobQueue:
|
|
def __init__(self, dbfile):
|
|
self.dbfile = dbfile
|
|
self.db = sqlite3.connect(dbfile, check_same_thread=False)
|
|
self.db.row_factory = sqlite3.Row
|
|
with self.db as conn:
|
|
conn.execute('''CREATE TABLE IF NOT EXISTS jobs
|
|
(id INTEGER PRIMARY KEY,
|
|
type TEXT,
|
|
params TEXT,
|
|
client TEXT,
|
|
result INTEGER DEFAULT NULL,
|
|
created DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
consumed DATETIME DEFAULT NULL,
|
|
aborted DATETIME DEFAULT NULL,
|
|
finished DATETIME DEFAULT NULL);''')
|
|
|
|
def enqueue(self, task_type:str, client, **params):
|
|
""" Enqueue a job of the given type with the given params. Returns the new job ID. """
|
|
with self.db as conn:
|
|
return conn.execute('INSERT INTO jobs(type, client, params) VALUES (?, ?, ?)',
|
|
(task_type, client, json.dumps(params))).lastrowid
|
|
|
|
def pop(self, task_type):
|
|
""" Fetch the next job of the given type. Returns a sqlite3.Row object of the job or None if no jobs of the given
|
|
type are queued. """
|
|
with self.db as conn:
|
|
job = conn.execute('SELECT * FROM jobs WHERE type=? AND consumed IS NULL AND aborted IS NULL ORDER BY created ASC LIMIT 1',
|
|
(task_type,)).fetchone()
|
|
if job is None:
|
|
return None
|
|
|
|
# Atomically commit to this job
|
|
conn.execute('UPDATE jobs SET consumed=datetime("now") WHERE id=?', (job['id'],))
|
|
|
|
return Job(self.db, job)
|
|
|
|
def job_iter(self, task_type):
|
|
return iter(lambda: self.pop(task_type), None)
|
|
|
|
def __getitem__(self, key):
|
|
""" Return the job with the given ID, or raise a KeyError if the key cannot be found. """
|
|
with self.db as conn:
|
|
job = conn.execute('SELECT * FROM jobs WHERE id=?', (key,)).fetchone()
|
|
if job is None:
|
|
raise KeyError(f'Unknown job ID "{key}"')
|
|
|
|
return Job(self.db, job)
|
|
|
|
class Job(dict):
|
|
def __init__(self, db, row):
|
|
super().__init__(json.loads(row['params']))
|
|
self._db = db
|
|
self._row = row
|
|
self.id = row['id']
|
|
self.type = row['type']
|
|
self.client = row['client']
|
|
self.created = row['created']
|
|
self.consumed = row['consumed']
|
|
self.finished = row['finished']
|
|
self.result = row['result']
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, _exc_type, _exc_val, _exc_tb):
|
|
with self._db as conn:
|
|
conn.execute('UPDATE jobs SET finished=datetime("now"), result=? WHERE id=?', (self.result, self.id,))
|
|
|
|
def abort(self, job_id):
|
|
with self.db as conn:
|
|
conn.execute('UPDATE jobs SET aborted=datetime("now") WHERE id=?', (self.id,))
|
|
|