1 files changed, 94 insertions, 0 deletions
diff --git a/joblib.py b/joblib.py
new file mode 100644
index 0000000..6891e11
--- /dev/null
+++ b/joblib.py
@@ -0,0 +1,94 @@
+#!/usr/bin/python3
+
+import asyncio
+import asyncpg
+import json
+
+
+class JobQueue(object):
+    def __init__(self):
+        self._pool = None
+
+    async def connect(self, **args):
+        """Connect the database."""
+        self._pool = await asyncpg.create_pool(**args)
+
+    async def setup(self):
+        """Initialize the database."""
+        async with self._pool.acquire() as con:
+            await con.execute("""
+                CREATE TABLE IF NOT EXISTS
+                jobs (
+                    id SERIAL,
+                    priority INT DEFAULT 0,
+                    status TEXT DEFAULT 'open',
+                    type TEXT,
+                    payload JSONB
+                )
+                """)
+
+    async def request(self, jobtype, payload, priority=0):
+        """Create a new job."""
+        async with self._pool.acquire() as con:
+            await con.execute("""
+                INSERT INTO jobs(type, payload, priority)
+                VALUES($1, $2, $3)
+            """, jobtype, json.dumps(payload), priority)
+
+    async def acquire(self, jobtype):
+        """Mark a job as running, return payload and return id.
+        Return (None, None) if no job is available."""
+        async with self._pool.acquire() as con:
+            while True:
+                try:
+                    # do not allow async access
+                    async with con.transaction(isolation="serializable"):
+                        result = await con.fetchrow("""
+                            SELECT id, payload
+                            FROM jobs WHERE
+                            type=$1 AND status='open'
+                            ORDER BY priority ASC
+                        """, jobtype)
+                        if result is None:
+                            # no jobs available
+                            return None, None
+                        jobid, payload = result
+                        await con.execute("""
+                            UPDATE jobs
+                            SET status='running'
+                            WHERE id=$1
+                        """, jobid)
+                        return jobid, json.loads(payload)
+                except asyncpg.exceptions.SerializationError:
+                    # job is being picked up by another worker, try again
+                    pass
+
+    async def finish(self, jobid):
+        """Mark a job as completed."""
+        async with self._pool.acquire() as con:
+            while True:
+                try:
+                    async with con.transaction(isolation="serializable"):
+                        await con.execute("""
+                            UPDATE jobs
+                            SET status='finished'
+                            WHERE id=$1
+                        """, jobid)
+                        return
+                except asyncpg.exceptions.SerializationError:
+                    pass
+
+    async def cleanup(self):
+        """Reopen all unfinished jobs."""
+        async with self._pool.acquire() as con:
+            while True:
+                try:
+                    async with con.transaction(isolation="serializable"):
+                        await con.execute("""
+                            UPDATE jobs
+                            SET status='open'
+                            WHERE status='running'
+                        """)
+                        return
+                except asyncpg.exceptions.SerializationError:
+                    pass