first implementation, fixes #1

3 files changed, 252 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..9a05e2d
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,91 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# dotenv
+.env
+
+# virtualenv
+.venv/
+venv/
+ENV/
+
+# Spyder project settings
+.spyderproject
+
+# Rope project settings
+.ropeproject
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
diff --git a/test_joblib.py b/test_joblib.py
new file mode 100644
index 0000000..35ae90e
--- /dev/null
+++ b/test_joblib.py
@@ -0,0 +1,67 @@
+#!/usr/bin/python3
+
+import os
+import asyncio
+import asyncpg
+import pytest
+import joblib
+
+class TestJoblib:
+    # async fixtures are not available yet, so we use a workaround
+    async def queue_helper(self, q):
+        await q.connect(
+            host=os.environ["POSTGRESQL_HOST"],
+            port=os.environ["POSTGRESQL_PORT"],
+            user=os.environ["POSTGRESQL_USER"],
+            password=os.environ["POSTGRESQL_PASSWORD"],
+            database=os.environ["POSTGRESQL_DB"]
+        )
+        async with q._pool.acquire() as con:
+            await con.execute("DROP TABLE IF EXISTS jobs")
+
+        # clean db table
+        await q.setup()
+
+    @pytest.fixture
+    def queue(self, event_loop):
+        queue = joblib.JobQueue()
+        event_loop.run_until_complete(self.queue_helper(queue))
+        return queue
+
+    @pytest.fixture
+    def payload(self):
+        return {
+            "key": "value",
+            "dict": {
+                "foo": 1,
+                "bar": "baz"
+            }
+        }
+
+    @pytest.mark.asyncio
+    async def test_request_and_acquire(self, queue, payload):
+        # request should succeed
+        await queue.request(jobtype="testing", payload=payload)
+        # acquire should return same payload
+        assert payload == (await queue.acquire(jobtype="testing"))[1]
+        # there should not be another job
+        assert None == (await queue.acquire(jobtype="testing"))[1]
+
+    @pytest.mark.asyncio
+    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")
+        assert payload_1 == payload
+        await queue.cleanup()
+        # same job should be available again
+        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_finish(self, queue, payload):
+        await queue.request(jobtype="testing", payload=payload)
+        jobid, _ = await queue.acquire(jobtype="testing")
+        await queue.finish(jobid)
+        # job should not be available again
+        assert None == (await queue.acquire(jobtype="testing"))[1]