add job fails

2 files changed, 29 insertions, 0 deletions

diff --git a/joblib.py b/joblib.py
index a0ce985..48666a6 100644
--- a/joblib.py
+++ b/joblib.py
@@ -78,6 +78,21 @@ class JobQueue(object):
                 except asyncpg.exceptions.SerializationError:
                     pass
 
+    async def fail(self, jobid, reason):
+        """Mark a job as failed."""
+        async with self._pool.acquire() as con:
+            while True:
+                try:
+                    async with con.transaction(isolation="serializable"):
+                        await con.execute("""
+                            UPDATE jobs
+                            SET status='failed', payload=payload||$2::jsonb
+                            WHERE id=$1
+                        """, jobid, json.dumps({"error": reason}))
+                        return
+                except asyncpg.exceptions.SerializationError:
+                    pass
+
     async def cleanup(self):
         """Reopen all unfinished jobs."""
         async with self._pool.acquire() as con:
diff --git a/test_joblib.py b/test_joblib.py
index 47eedec..f9aa875 100644
--- a/test_joblib.py
+++ b/test_joblib.py
@@ -2,6 +2,7 @@
 
 import os
 import asyncio
+import json
 import asyncpg
 import pytest
 import joblib
@@ -64,6 +65,19 @@ class TestJoblib:
         assert 9 == (await queue.acquire(jobtype="testing"))[2]
 
     @pytest.mark.asyncio
+    async def test_fail(self, queue, payload):
+        err = "testing errors"
+        await queue.request(jobtype="testing", payload=payload)
+        jobid, _, _ = await queue.acquire(jobtype="testing")
+        await queue.fail(jobid, err)
+        async with queue._pool.acquire() as con:
+            jid, pl = await con.fetchrow(
+                "SELECT id, payload FROM jobs WHERE status='failed'")
+            assert jid == jobid
+            payload["error"] = err
+            assert json.loads(pl) == payload
+
+    @pytest.mark.asyncio
     async def test_finish(self, queue, payload):
         await queue.request(jobtype="testing", payload=payload)
         jobid, _, _ = await queue.acquire(jobtype="testing")