From bcaee98d6db7d24df8ccd25bc466ca6da3b4f75e Mon Sep 17 00:00:00 2001
From: schneefux <schneefux+commit@schneefux.xyz>
Date: Sun, 26 Feb 2017 11:32:25 +0100
Subject: acquire: return priority

---
 joblib.py      | 10 +++++-----
 test_joblib.py | 11 ++++++++---
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/joblib.py b/joblib.py
index e11ec52..6bde310 100644
--- a/joblib.py
+++ b/joblib.py
@@ -36,15 +36,15 @@ class JobQueue(object):
             """, 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."""
+        """Mark a job as running, return id, payload and priority.
+        Return (None, 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
+                            SELECT id, payload, priority
                             FROM jobs WHERE
                             type=$1 AND status='open'
                             ORDER BY priority DESC
@@ -52,13 +52,13 @@ class JobQueue(object):
                         if result is None:
                             # no jobs available
                             return None, None
-                        jobid, payload = result
+                        jobid, payload, priority = result
                         await con.execute("""
                             UPDATE jobs
                             SET status='running'
                             WHERE id=$1
                         """, jobid)
-                        return jobid, json.loads(payload)
+                        return jobid, json.loads(payload), priority
                 except asyncpg.exceptions.SerializationError:
                     # job is being picked up by another worker, try again
                     pass
diff --git a/test_joblib.py b/test_joblib.py
index 35ae90e..47eedec 100644
--- a/test_joblib.py
+++ b/test_joblib.py
@@ -51,17 +51,22 @@ class TestJoblib:
     async def test_cleanup(self, queue, payload):
         await queue.request(jobtype="testing", payload=payload)
         # mark job as processing
-        jobid_1, payload_1 = await queue.acquire(jobtype="testing")
+        jobid_1, payload_1, _ = await queue.acquire(jobtype="testing")
         assert payload_1 == payload
         await queue.cleanup()
         # same job should be available again
-        jobid_2, payload_2 = await queue.acquire(jobtype="testing")
+        jobid_2, payload_2, _ = await queue.acquire(jobtype="testing")
         assert jobid_1 == jobid_2 and payload_1 == payload_2
 
+    @pytest.mark.asyncio
+    async def test_priority(self, queue, payload):
+        await queue.request(jobtype="testing", payload=payload, priority=9)
+        assert 9 == (await queue.acquire(jobtype="testing"))[2]
+
     @pytest.mark.asyncio
     async def test_finish(self, queue, payload):
         await queue.request(jobtype="testing", payload=payload)
-        jobid, _ = await queue.acquire(jobtype="testing")
+        jobid, _, _ = await queue.acquire(jobtype="testing")
         await queue.finish(jobid)
         # job should not be available again
         assert None == (await queue.acquire(jobtype="testing"))[1]
-- 
cgit v1.3.1