<file_sep>import sys
import os
import django
from django.conf import settings
from django.core.management import call_command


opts = {
    'INSTALLED_APPS': ['pgqueue'],
    'DATABASES': {
        'default': {
            'ENGINE': 'django.db.backends.postgresql',
            'HOST': '127.0.0.1',
            'PORT': '5432',
            'NAME': 'pgqueue',
            'USER': 'postgres',
            'PASSWORD': '',
        },
    },
}


if __name__ == '__main__':
    sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
    settings.configure(**opts)
    django.setup()

    call_command('test', 'pgqueue')
<file_sep># Generated by Django 2.0.2 on 2018-03-19 13:01

import django.contrib.postgres.fields.jsonb
import django.contrib.postgres.functions
from django.db import migrations, models


class Migration(migrations.Migration):

    initial = True

    dependencies = [
    ]

    operations = [
        migrations.CreateModel(
            name='Job',
            fields=[
                ('id', models.BigAutoField(primary_key=True, serialize=False)),
                ('task', models.CharField(max_length=255)),
                ('created_at', models.DateTimeField(default=django.contrib.postgres.functions.TransactionNow)),
                ('execute_at', models.DateTimeField(default=django.contrib.postgres.functions.TransactionNow)),
                ('priority', models.PositiveIntegerField(default=0)),
                ('kwargs', django.contrib.postgres.fields.jsonb.JSONField(blank=True, default={})),
            ],
            options={
                'abstract': False,
            },
        ),
        migrations.AddIndex(
            model_name='job',
            index=models.Index(fields=['-priority', 'created_at'], name='pgqueue_job_priorit_c6c76e_idx'),
        ),
    ]
<file_sep>from django.contrib.postgres.functions import TransactionNow
from django.contrib.postgres.fields import JSONField
from django.db import models


def lazy_empty_dict():
    return {}


class BaseJob(models.Model):
    id = models.BigAutoField(primary_key=True)
    task = models.CharField(max_length=255)
    created_at = models.DateTimeField(default=TransactionNow)
    execute_at = models.DateTimeField(default=TransactionNow)
    priority = models.PositiveIntegerField(default=0)
    context = JSONField(default=lazy_empty_dict, blank=True)
    kwargs = JSONField(default=lazy_empty_dict, blank=True)

    class Meta:
        indexes = [
            models.Index(fields=['-priority', 'created_at']),
        ]
        abstract = True

    def __str__(self):
        return '{} {}'.format(self.task, self.kwargs)

    def as_dict(self):
        return {
            'id': self.id,
            'created_at': self.created_at,
            'execute_at': self.execute_at,
            'priority': self.priority,
            'kwargs': self.kwargs,
            'task': self.task,
        }


class Job(BaseJob):
    pass
<file_sep># Django PGQueue

The project was initially forked from [django-postgres-queue][dpq] for internal use
in some of my projects. After some changes and refactoring, adding tests, features 
like warm shutdown I decided to put in on GitHub and PyPI as a standalone package
just in case.

## Installation

```
pip install django-pgqueue
```

Then add `pqueue` into `INSTALLED_APPS` and run `manage.py migrate` to create the jobs table.

## Usage

Initiate a queue object with defined tasks. This can go wherever you like and can 
be called whatever you like. For example:

```
# someapp/task_queue.py

from pqueue.queue import Queue


def say_hello(queue, job):
    name = job.kwargs['name']
    print('Hello, {}!'.format(name))


task_queue = Queue(
    tasks={
        'say_hello': say_hello,
    },
    notify_channel='someapp_task_queue',
)
```

Now define the worker command

```
# someapp/management/commands/pgqueue_worker.py

from pgqueue.worker import WorkerCommand
from someapp.queue import task_queue


class Command(WorkerCommand):
    queue = task_queue
```

And call the task this way:

```
from someapp.queue import task_queue

task_queue.enqueue('say_hello', {'name': 'Django'})
```

Please note that only primitives can be used in job’s arguments as they are stored
as JSON in the database. When you try to pass any complex non-json-serializable object,
you will get an error saying `object is not json serializable`.

## Periodic tasks

There is no built in way to run jobs periodically like celerybeat in Celery.
But you still can use cron. For example you can create a universal command
to execute any task. Something like this:

```
import json

from django.core.management import BaseCommand
from someapp.queue import task_queue


class Command(BaseCommand):
    def add_arguments(self, parser):
        parser.add_argument('task_name')
        parser.add_argument('task_kwargs')

    def handle(self, task_name, task_kwargs, **options):
        task_queue.enqueue(task_name, json.loads(task_kwargs))
```

And then put it into your cron records:

```
0 0 * * *  /path/to/python manage.py run_task say_hello '{"name": "Django!"}'
```


[dpq]: https://github.com/gavinwahl/django-postgres-queue
<file_sep>import select
import logging
import time

from django.utils import translation
from django.db import connection

from .models import Job


class Queue:
    job_model = Job
    logger = logging.getLogger(__name__)

    def __init__(self, tasks, notify_channel):
        self.tasks = tasks
        self.notify_channel = notify_channel

    def notify(self, job):
        with connection.cursor() as cur:
            cur.execute('NOTIFY "{}", %s;'.format(self.notify_channel), [str(job.pk)])

    def get_job_context(self):
        language = translation.get_language()
        return {'language': language}

    def in_atomic_block(self):
        return connection.in_atomic_block

    def enqueue(self, task, kwargs=None, execute_at=None, priority=0, context=None):
        assert not self.in_atomic_block(), 'Task cannot be executed inside a trasnaction'
        assert task in self.tasks, 'Task "{}" not found in the task list'
        kwargs = kwargs or {}

        job_context = self.get_job_context()
        if context is not None:
            job_context.update(context)

        job = Job(
            task=task,
            kwargs=kwargs,
            priority=priority,
            context=job_context,
        )
        if execute_at:
            job.execute_at = execute_at
        job.save()

        if self.notify_channel:
            self.notify(job)

        self.logger.info('New job scheduled %r', job)
        return job

    def enqueue_once(self, task, kwargs=None, *args_, **kwargs_):
        job = self.job_model.objects.filter(task=task, kwargs=kwargs or {}).first()
        if job is not None:
            return job
        return self.enqueue(task, kwargs, *args_, **kwargs_)

    def listen(self):
        with connection.cursor() as cur:
            cur.execute('LISTEN "{}";'.format(self.notify_channel))

    def wait(self, timeout):
        connection.connection.poll()
        notifies = self._get_notifies()
        if notifies:
            return notifies

        select.select([connection.connection], [], [], timeout)
        connection.connection.poll()
        return self._get_notifies()

    def _get_notifies(self):
        notifies = [
            i for i in connection.connection.notifies
            if i.channel == self.notify_channel]

        connection.connection.notifies = [
            i for i in connection.connection.notifies
            if i.channel != self.notify_channel]

        return notifies

    def dequeue(self):
        query = """
        DELETE FROM {table}
        WHERE id = (
            SELECT id
            FROM {table}
            WHERE execute_at <= NOW()
            ORDER BY priority DESC, created_at, id
            FOR UPDATE SKIP LOCKED
            LIMIT 1
        )
        RETURNING *;
        """
        query = query.format(table=self.job_model._meta.db_table)
        results = list(self.job_model.objects.raw(query))
        assert len(results) <= 1

        if not results:
            return None

        return results[0]

    def run_job(self, job):
        language = job.context.get('language')
        with translation.override(language):
            return self._run_job(job)

    def _run_job(self, job):
        task = self.tasks[job.task]
        start_time = time.time()
        retval = task(self, job)
        self.logger.info(
            'Processing %r took %0.4f seconds. Task returned %r.',
            job, time.time() - start_time, retval,
            extra={
                'data': {
                    'job': job.as_dict(),
                    'retval': retval,
                },
            },
        )
        return retval
<file_sep># Generated by Django 2.0.5 on 2018-11-04 22:13

import django.contrib.postgres.fields.jsonb
from django.db import migrations
import pgqueue.models


class Migration(migrations.Migration):

    dependencies = [
        ('pgqueue', '0002_job_context'),
    ]

    operations = [
        migrations.AlterField(
            model_name='job',
            name='context',
            field=django.contrib.postgres.fields.jsonb.JSONField(blank=True, default=pgqueue.models.lazy_empty_dict),
        ),
        migrations.AlterField(
            model_name='job',
            name='kwargs',
            field=django.contrib.postgres.fields.jsonb.JSONField(blank=True, default=pgqueue.models.lazy_empty_dict),
        ),
    ]
<file_sep>import logging
import signal
import sys
import os

from django.core.management import BaseCommand


class Worker:
    logger = logging.getLogger(__name__)
    wait_timeout = 10
    queue = None

    def __init__(self):
        self._shutdown = False
        self._is_waiting = False

    def handle_shutdown(self, sig, frame):
        if self._is_waiting:
            # If there're no active tasks let's exit immediately.
            sys.exit(0)

        self.logger.info('Waiting for active tasks to finish...')
        self._shutdown = True

    def bind_signals(self):
        """Handle the signals for warm shutdown."""
        signal.signal(signal.SIGINT, self.handle_shutdown)
        signal.signal(signal.SIGTERM, self.handle_shutdown)

    def start(self):
        self.bind_signals()
        self.queue.listen()
        self.logger.info('Worker is started with pid {}'.format(os.getpid()))

        while True:
            if self._shutdown:
                return

            job = self.queue.dequeue()
            if job is None:
                self.wait()
                continue

            try:
                self.queue.run_job(job)
            except Exception as e:
                self.logger.exception('Error in %r: %r.', job, e, extra={
                    'data': {
                        'job': job.as_dict(),
                    },
                })

    def wait(self):
        self._is_waiting = True
        self.queue.wait(self.wait_timeout)
        self._is_waiting = False


class WorkerCommand(Worker, BaseCommand):
    def before_start(self):
        pass

    def handle(self, *args, **options):
        self.before_start()
        self.start()
<file_sep># -*- coding: utf-8 -*-
# Generated by Django 1.11 on 2018-05-03 17:09
from __future__ import unicode_literals

import django.contrib.postgres.fields.jsonb
from django.db import migrations


class Migration(migrations.Migration):

    dependencies = [
        ('pgqueue', '0001_initial'),
    ]

    operations = [
        migrations.AddField(
            model_name='job',
            name='context',
            field=django.contrib.postgres.fields.jsonb.JSONField(blank=True, default={}),
        ),
    ]
<file_sep>import time
from contextlib import contextmanager
from datetime import timedelta
from threading import Thread
from unittest import mock

from django.test import SimpleTestCase
from django.utils import timezone

from .decorators import retry, repeat
from .models import Job
from .queue import Queue
from .worker import Worker


def demo_task(queue, job):
    return job.kwargs['value']


class QueueTest(SimpleTestCase):
    allow_database_queries = True

    def setUp(self):
        self.db_objects = []
        self.queue = Queue(
            tasks={'demo_task': demo_task},
            notify_channel='demo_queue',
        )

    def tearDown(self):
        for obj in self.db_objects:
            obj.delete()

    def test_enqueue(self):
        self.queue.enqueue('demo_task', kwargs={'value': 'val1'})

        job = Job.objects.first()
        self.db_objects.append(job)
        self.assertEqual(job.task, 'demo_task')
        self.assertEqual(job.kwargs, {'value': 'val1'})

    def test_enqueue_once(self):
        old_job = Job.objects.create(task='demo_task', kwargs={'value': 'val1'})
        job_1 = self.queue.enqueue_once('demo_task', {'value': 'val1'})
        self.assertEqual(job_1.id, old_job.id)
        job_2 = self.queue.enqueue_once('demo_task', {'value': 'val2'})
        self.assertNotEqual(job_2, old_job.id)

    def test_dequeue(self):
        job1 = Job.objects.create(task='demo_task')
        job2 = Job.objects.create(task='demo_task')
        self.db_objects += [job1, job2]
        self.assertEqual(job1.id, self.queue.dequeue().id)
        self.assertEqual(job2.id, self.queue.dequeue().id)

    def test_run_job(self):
        job = Job.objects.create(task='demo_task', kwargs={'value': 'hello'})
        self.db_objects.append(job)
        retval = self.queue.run_job(job)
        self.assertEqual(retval, 'hello')

    def test_listen_notify(self):
        self.queue.listen()

        job = Job.objects.create(task='demo_task', kwargs={'value': 'hello'})
        self.db_objects.append(job)
        self.queue.notify(job)

        notifies = self.queue.wait(1)
        notify = notifies[0]

        self.assertEqual(notify.channel, 'demo_queue')
        self.assertEqual(notify.payload, str(job.id))


class WorkerTest(SimpleTestCase):
    allow_database_queries = True

    class DemoWorker(Worker):
        wait_timeout = 1

        def __init__(self, queue, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self.queue = queue

        def start(self):
            super().start()

            # Close database connection when worker is stopped
            # top prevent "database is being accessed by other users".
            from django.db import connection
            connection.close()

        def stop(self):
            self._shutdown = True

        def bind_signals(self):
            # Signals doesn't work if the worker is running not in
            # the main thread, also there are no signals during
            # tests running so let’s ignore them.
            pass

    def setUp(self):
        self.db_objects = []

    def tearDown(self):
        for obj in self.db_objects:
            obj.delete()

    @contextmanager
    def start_worker(self, queue):
        worker = self.DemoWorker(queue)
        thread = Thread(target=worker.start, daemon=True)
        thread.start()
        yield worker
        time.sleep(1)  # wait worker to finish the tasks
        worker.stop()
        thread.join()

    def test_enqueue(self):
        queue = Queue(
            tasks={
                'task1': mock.Mock(),
                'task2': mock.Mock(),
            },
            notify_channel='test_channel',
        )

        with self.start_worker(queue):
            job1 = queue.enqueue('task1')
            job2 = queue.enqueue('task2')
            self.db_objects += [job1, job2]

        queue.tasks['task1'].assert_called()
        queue.tasks['task2'].assert_called()


class RepeatDecoratorTest(SimpleTestCase):
    allow_database_queries = True

    def tearDown(self):
        Job.objects.all().delete()

    def test_repeat(self):
        task = mock.Mock()
        wrapped_task = repeat(timedelta(minutes=1))(task)

        queue = Queue(tasks={'task': wrapped_task}, notify_channel='test_channel')
        job = Job(task='task')  # job should not be saved in database to simulate dequeue call
        wrapped_task(queue, job)

        new_job = Job.objects.last()
        delta = new_job.execute_at - timezone.now()
        self.assertAlmostEqual(delta.total_seconds(), timedelta(minutes=1).total_seconds(), places=1)


class RetryDecoratorTest(SimpleTestCase):
    allow_database_queries = True

    def tearDown(self):
        Job.objects.all().delete()

    def test_retry(self):
        task = mock.Mock()
        task.side_effect = Exception()
        wrapped_task = retry([Exception], delay=timedelta(seconds=1), max_attempts=3)(task)

        queue = Queue(tasks={'task': task}, notify_channel='test_channel')
        job = Job.objects.create(task='task')

        wrapped_task(queue, job)
        retry_job = Job.objects.last()

        self.assertEqual(retry_job.task, 'task')
        self.assertEqual(retry_job.context['retry_attempt'], 1)

    def test_retry_delay(self):
        task = mock.Mock()
        task.side_effect = Exception()
        wrapped_task = retry([Exception], delay=timedelta(seconds=10), max_attempts=10)(task)

        queue = Queue(tasks={'task': task}, notify_channel='test_channel')
        job = Job.objects.create(task='task')

        for n in range(10):
            wrapped_task(queue, job)
            job = Job.objects.order_by('created_at').last()
            delta = job.execute_at - timezone.now()
            self.assertAlmostEqual(delta.total_seconds(), timedelta(seconds=10).total_seconds(), places=1)

        jobs = Job.objects.all()
        self.assertEqual(jobs.count(), 11)

    def test_retry_max_attempts(self):
        task = mock.Mock()
        task.side_effect = Exception()
        wrapped_task = retry([Exception], delay=timedelta(seconds=1), max_attempts=10)(task)

        queue = Queue(tasks={'task': task}, notify_channel='test_channel')
        job = Job.objects.create(task='task')

        for n in range(10):
            wrapped_task(queue, job)
            job = Job.objects.order_by('created_at').last()
            self.assertEqual(job.context['retry_attempt'], n + 1)

        with self.assertRaises(Exception):
            wrapped_task(queue, job)

        jobs = Job.objects.all()
        self.assertEqual(jobs.count(), 11)
<file_sep>from functools import wraps
import copy

from django.utils import timezone

RETRY_ATTEMPT = 'retry_attempt'


def repeat(interval):
    """
    Endlessly repeats a task, every `delay` (a timedelta).

        @repeat(datetime.timedelta(minutes=5)
        def task(queue, job):
            pass

    This will run `task` every 5 minutes. It's up to you to kick off the first
    task, though.
    """
    def wrapper(func):
        @wraps(func)
        def decorator(queue, job):
            queue.enqueue_once(job.task, job.kwargs, execute_at=(timezone.now() + interval))
            return func(queue, job)
        return decorator
    return wrapper


def retry(exceptions, delay, max_attempts):
    """
    Retries failed task every `interval` (a timedelta).

        @retry([Exception], delay=timedelta(seconds=30), max_attempts=10)
        def task(queue, job):
            pass

    This will repeat `task` every 30 seconds if it fails with `Exception`
    but not more than ten times. When `max_attempts` is reached,
    the original exception will be raised.
    """
    def wrapper(func):
        @wraps(func)
        def decorator(queue, job):
            try:
                return func(queue, job)
            except Exception as e:
                if not isinstance(e, tuple(exceptions)) \
                        or queue is None:  # for testing reasons
                    raise

                attempt = job.context.get(RETRY_ATTEMPT, 0) + 1
                if max_attempts is not None and attempt > max_attempts:
                    raise

                context = copy.copy(job.context)
                context[RETRY_ATTEMPT] = attempt
                execute_at = timezone.now() + delay
                queue.enqueue(job.task, job.kwargs, context=context, execute_at=execute_at)

        return decorator
    return wrapper


def job_kwargs(func):
    @wraps(func)
    def wrapper(queue, job):
        return func(**job.kwargs)
    return wrapper
<file_sep>from unittest import mock


def call_task(task, kwargs):
    job = mock.Mock()
    job.kwargs = kwargs
    queue = mock.Mock()
    return task(queue, job)
<file_sep>from django.contrib import admin
from .models import Job


@admin.register(Job)
class JobAdmin(admin.ModelAdmin):
    list_display = ('id', 'task', 'execute_at', 'priority')
    ordering = ('-priority', 'execute_at')
