add status()

2 files changed, 23 insertions, 2 deletions

diff --git a/joblib.py b/joblib.py
index 48666a6..de69cfd 100644
--- a/joblib.py
+++ b/joblib.py
@@ -28,11 +28,12 @@ class JobQueue(object):
                 """)
 
     async def request(self, jobtype, payload, priority=0):
-        """Create a new job."""
+        """Create a new job and return its id."""
         async with self._pool.acquire() as con:
-            await con.execute("""
+            return await con.fetchval("""
                 INSERT INTO jobs(type, payload, priority)
                 VALUES($1, $2, $3)
+                RETURNING id
             """, jobtype, json.dumps(payload), priority)
 
     async def acquire(self, jobtype):
@@ -62,6 +63,15 @@ class JobQueue(object):
                 except asyncpg.exceptions.SerializationError:
                     # job is being picked up by another worker, try again
                     pass
+    
+    async def status(self, jobid):
+        """Return the status of a job."""
+        async with self._pool.acquire() as con:
+            return await con.fetchval("""
+                SELECT status
+                FROM jobs WHERE
+                id=$1
+            """, jobid)
 
     async def finish(self, jobid):
         """Mark a job as completed."""
diff --git a/test_joblib.py b/test_joblib.py
index f9aa875..865cfab 100644
--- a/test_joblib.py
+++ b/test_joblib.py
@@ -78,6 +78,17 @@ class TestJoblib:
             assert json.loads(pl) == payload
 
     @pytest.mark.asyncio
+    async def test_status(self, queue, payload):
+        assert await queue.status(-1) == None
+        jobid = await queue.request(jobtype="testing",
+                                    payload=payload)
+        assert await queue.status(jobid) == "open"
+        await queue.acquire(jobtype="testing")
+        assert await queue.status(jobid) == "running"
+        await queue.finish(jobid)
+        assert await queue.status(jobid) == "finished"
+
+    @pytest.mark.asyncio
     async def test_finish(self, queue, payload):
         await queue.request(jobtype="testing", payload=payload)
         jobid, _, _ = await queue.acquire(jobtype="testing")