djanquiltdb-plugin-celery 1.0.0__py3-none-any.whl

djanquiltdb_plugin_celery/__init__.py

@@ -0,0 +1,52 @@
+"""
+This package defines the plugin object, which DjanQuiltDB finds through the ``djanquiltdb.plugins`` entry-point group,
+and exports the public API of the plugin. DjanQuiltDB serves the decorators as ``djanquiltdb.decorators.celery``, and
+the helpers as ``djanquiltdb.utils.celery``.
+"""
+
+from types import SimpleNamespace
+
+from djanquiltdb_plugin_celery import decorators, utils
+from djanquiltdb_plugin_celery.decorators import (
+    task_for_active_shard,
+    task_for_every_node,
+    task_for_every_shard,
+    task_for_node,
+    task_for_shard,
+)
+from djanquiltdb_plugin_celery.utils import (
+    aggregate_per_shard,
+    fan_out,
+    fan_out_to_nodes,
+    fanout_shards,
+    per_node_signatures,
+    per_shard_signatures,
+)
+
+__version__ = '1.0.0'
+
+
+def install():
+    """
+    Registers the plugin's system checks. DjanQuiltDB calls install() once the app registry is ready.
+    """
+    from djanquiltdb_plugin_celery import checks  # noqa: F401
+
+
+plugin = SimpleNamespace(decorators=decorators, utils=utils, install=install)
+
+__all__ = [
+    'aggregate_per_shard',
+    'fan_out',
+    'fan_out_to_nodes',
+    'fanout_shards',
+    'install',
+    'per_node_signatures',
+    'per_shard_signatures',
+    'plugin',
+    'task_for_active_shard',
+    'task_for_every_node',
+    'task_for_every_shard',
+    'task_for_node',
+    'task_for_shard',
+]

djanquiltdb_plugin_celery/checks.py

@@ -0,0 +1,79 @@
+from django.conf import settings
+from django.core import checks
+from djanquiltdb.utils import get_mapping_class
+
+from djanquiltdb_plugin_celery.settings import (
+    QUEUE_DESCRIPTION,
+    QUILT_CELERY_DEFAULTS,
+    SPREAD_DESCRIPTION,
+    is_spread,
+)
+from djanquiltdb_plugin_celery.targets import NODE_NAME_KWARG, SHARD_ID_KWARG
+
+#: For each key of QUILT_CELERY: a function that returns whether a value of the key is valid, and how a message
+#: describes the valid values. A boolean is an instance of int, so the keys that accept an integer and not a boolean
+#: exclude it.
+QUILT_CELERY_VALIDATORS = {
+    'MAPPING_KWARG': (lambda value: isinstance(value, (str, type(None))), 'a string or None'),
+    'MAX_RETRIES': (
+        lambda value: value is None or (isinstance(value, int) and not isinstance(value, bool)),
+        'an integer or None',
+    ),
+    'RETRY_BACKOFF': (lambda value: isinstance(value, (bool, int)), 'a boolean or an integer'),
+    'RETRY_BACKOFF_MAX': (lambda value: isinstance(value, int) and not isinstance(value, bool), 'an integer'),
+    'RETRY_JITTER': (lambda value: isinstance(value, bool), 'a boolean'),
+    'FANOUT_QUEUE': (lambda value: isinstance(value, (str, type(None))), QUEUE_DESCRIPTION),
+    'FANOUT_SPREAD': (is_spread, SPREAD_DESCRIPTION),
+}
+
+
+@checks.register(checks.Tags.compatibility)
+def check_quilt_celery_setting(app_configs, **kwargs):
+    """
+    Reports a QUILT_CELERY setting, key or value that the celery plugin cannot use.
+    """
+    quilt_celery = getattr(settings, 'QUILT_CELERY', {})
+
+    if not isinstance(quilt_celery, dict):
+        return [checks.Error('QUILT_CELERY is not a dict.', id='djanquiltdb_plugin_celery.E002')]
+
+    messages = [
+        checks.Error(
+            'QUILT_CELERY has the key {!r}, which the celery plugin does not read.'.format(name),
+            hint='The plugin reads the keys {}.'.format(', '.join(QUILT_CELERY_DEFAULTS)),
+            id='djanquiltdb_plugin_celery.E001',
+        )
+        for name in quilt_celery
+        if name not in QUILT_CELERY_DEFAULTS
+    ]
+
+    for name, (is_valid, description) in QUILT_CELERY_VALIDATORS.items():
+        if name in quilt_celery and not is_valid(quilt_celery[name]):
+            messages.append(
+                checks.Error(
+                    "QUILT_CELERY['{}'] is {!r}, and it accepts {}.".format(name, quilt_celery[name], description),
+                    id='djanquiltdb_plugin_celery.E002',
+                )
+            )
+
+    if quilt_celery.get('MAPPING_KWARG') in (SHARD_ID_KWARG, NODE_NAME_KWARG):
+        messages.append(
+            checks.Error(
+                "QUILT_CELERY['MAPPING_KWARG'] is {!r}, which already identifies a target in another way.".format(
+                    quilt_celery['MAPPING_KWARG']
+                ),
+                hint='Set it to the keyword argument that has the mapping value, such as tenant_id.',
+                id='djanquiltdb_plugin_celery.E002',
+            )
+        )
+    elif quilt_celery.get('MAPPING_KWARG') and get_mapping_class() is None:
+        messages.append(
+            checks.Error(
+                "QUILT_CELERY['MAPPING_KWARG'] is {!r}, but QUILT_DB has no MAPPING_MODEL in which a task can look up "
+                'the shard of a mapping value.'.format(quilt_celery['MAPPING_KWARG']),
+                hint="Set QUILT_DB['MAPPING_MODEL'], or remove QUILT_CELERY['MAPPING_KWARG'].",
+                id='djanquiltdb_plugin_celery.E002',
+            )
+        )
+
+    return messages

djanquiltdb_plugin_celery/decorators.py

@@ -0,0 +1,289 @@
+"""
+The task decorators of the plugin. DjanQuiltDB serves this module as ``djanquiltdb.decorators.celery``. Because
+``SERVE_UNQUALIFIED`` is False, DjanQuiltDB does not also serve these decorators as attributes of
+``djanquiltdb.decorators``, so that a call site names the plugin, as in ``@celery.task_for_shard``.
+"""
+
+import annotationlib
+import functools
+import inspect
+
+from celery import shared_task
+from djanquiltdb.utils import StateException
+
+from djanquiltdb_plugin_celery.settings import ensure_valid_queue, ensure_valid_spread, quilt_celery_setting
+from djanquiltdb_plugin_celery.targets import use_target
+from djanquiltdb_plugin_celery.tasks import QuiltTask
+
+SERVE_UNQUALIFIED = False
+
+#: The signature of the wrapper of each function: the task, and any arguments.
+WRAPPER_SIGNATURE = inspect.Signature(
+    [
+        inspect.Parameter('self', inspect.Parameter.POSITIONAL_ONLY),
+        inspect.Parameter('args', inspect.Parameter.VAR_POSITIONAL),
+        inspect.Parameter('kwargs', inspect.Parameter.VAR_KEYWORD),
+    ]
+)
+
+
+def _retry_options(options):
+    """
+    Returns the Celery options that retry a task while its shard is in maintenance, with the values of ``options``
+    taking precedence over the QUILT_CELERY settings.
+    """
+    return {
+        'max_retries': quilt_celery_setting('MAX_RETRIES'),
+        'retry_backoff': quilt_celery_setting('RETRY_BACKOFF'),
+        'retry_backoff_max': quilt_celery_setting('RETRY_BACKOFF_MAX'),
+        'retry_jitter': quilt_celery_setting('RETRY_JITTER'),
+        **options,
+        'autoretry_for': (StateException, *options.get('autoretry_for', ())),
+    }
+
+
+def _signature_of_function(func, bind):
+    """
+    Returns the signature of ``func`` without the parameter for the task, which ``func`` takes when ``bind`` is True.
+    The arguments of each message of the task are checked against this signature.
+    """
+    signature = inspect.signature(func, annotation_format=annotationlib.Format.STRING)
+    parameters = list(signature.parameters.values())
+
+    return signature.replace(parameters=parameters[1:] if bind else parameters)
+
+
+def _create_quilt_task(func, task_attributes, options):
+    """
+    Returns the task of ``func``, with ``task_attributes`` as attributes of its QuiltTask class and ``options`` as
+    the options of ``celery.shared_task()``.
+    """
+    ensure_valid_spread(task_attributes.get('fanout_spread'), func.__qualname__)
+    ensure_valid_queue(task_attributes.get('fanout_queue'), func.__qualname__)
+
+    options = dict(options)
+    bind = options.pop('bind', False)
+    base = options.pop('base', QuiltTask)
+
+    if not issubclass(base, QuiltTask):
+        raise TypeError(
+            'The base of {} is {}, which does not inherit from QuiltTask. The tasks of the celery plugin need '
+            'the methods of QuiltTask to send and to enter their targets.'.format(func.__qualname__, base.__qualname__)
+        )
+
+    # The wrapper is always bound, because it reads the request of the task. The function receives the task only when it
+    # is decorated with bind=True.
+    @functools.wraps(func)
+    def run_in_target(self, *args, **kwargs):
+        target = self.target_of_message(self.request.headers, kwargs)
+
+        if target is None and self.fans_out:
+            return self.fan_out_to_every_target(args, kwargs)
+
+        if target is None:
+            raise self.missing_target_error(kwargs)
+
+        with use_target(target, lock=self.lock_target):
+            function_kwargs = self.kwargs_for_function(kwargs, self.entered_target_value(target))
+
+            return func(self, *args, **function_kwargs) if bind else func(*args, **function_kwargs)
+
+    # functools.wraps makes the wrapper report the signature of the function, and Celery checks the arguments of a
+    # message against the signature of the wrapper. The function does not receive a keyword argument of the target that
+    # it does not declare, so the wrapper takes any arguments, and QuiltTask checks the arguments in the form in which
+    # the function receives them.
+    run_in_target.__signature__ = WRAPPER_SIGNATURE
+    task_attributes = {**task_attributes, 'function_signature': _signature_of_function(func, bind)}
+
+    return shared_task(bind=True, base=base, **task_attributes, **_retry_options(options))(run_in_target)
+
+
+def _task_decorator_or_task(task_attributes, func, options):
+    """
+    Returns the task of ``func``, or, when ``func`` is None, the decorator that returns the task, so that each decorator
+    works with and without arguments.
+    """
+
+    def decorate(func):
+        return _create_quilt_task(func, task_attributes, options)
+
+    return decorate if func is None else decorate(func)
+
+
+def task_for_shard(func=None, *, inherit_active_shard=False, lock=None, **options):
+    """
+    Turns a function into a Celery task that runs in a shard: the shard that ``for_shard()`` is given, else the shard
+    that its keyword arguments identify. That shard is the target of the task. ``shard_id`` identifies a shard by its
+    primary key. The mapping field of the ``MAPPING_MODEL``, such as ``tenant_id``, identifies a shard by its
+    mapping value, and ``QUILT_CELERY['MAPPING_KWARG']`` can set another keyword argument for the mapping value.
+
+    The function receives ``shard_id`` and the keyword argument for a mapping value only when it declares them by name,
+    and ``**kwargs`` does not receive them. A function that declares ``shard_id`` receives the primary key of the shard
+    that the task runs in, also when the task is sent without ``shard_id``. ``node_name`` is an ordinary keyword
+    argument for the function: a task that runs on a node is decorated with :func:`task_for_node`. The arguments of the
+    task are checked against the function when the task is sent, unless the task is decorated with ``typing=False``.
+
+    Sending the task without a shard raises TaskShardError, unless the task inherits the active shard and a shard is
+    active where the task is sent. Sending the task with two shards raises TaskShardError too. While the shard is in
+    maintenance, the task is retried with an exponential backoff, and after the last retry the task fails with the
+    StateException.
+
+    :param inherit_active_shard: When True, a task that is sent without a shard in its keyword arguments runs in the
+        shard that is active where it is sent, as :func:`task_for_active_shard` does.
+    :param lock: Passed on to ``use_shard()`` as ``lock``. By default, ``use_shard()`` locks the shard while the task
+        runs.
+    :param options: The options of ``celery.shared_task()``. ``max_retries``, ``retry_backoff``, ``retry_backoff_max``
+        and ``retry_jitter`` default to the QUILT_CELERY settings, and StateException is added to the exceptions of
+        ``autoretry_for``. ``base`` defaults to :class:`~djanquiltdb_plugin_celery.tasks.QuiltTask`, and a
+        ``base`` that does not inherit from that class raises TypeError.
+
+    :Example:
+        .. code-block:: python
+
+            from djanquiltdb.decorators.celery import task_for_shard
+
+
+            @task_for_shard
+            def send_invoice(invoice_id):
+                Invoice.objects.get(pk=invoice_id).send()
+
+
+            send_invoice.for_shard(shard).delay(invoice_id=12)
+            send_invoice.delay_for_shard(shard, invoice_id=12)  # the same
+            send_invoice.delay(invoice_id=12, tenant_id=3)
+    """
+    task_attributes = {'runs_in': 'shard', 'lock_target': lock, 'inherit_active_shard': inherit_active_shard}
+
+    return _task_decorator_or_task(task_attributes, func, options)
+
+
+def task_for_node(func=None, **options):
+    """
+    Turns a function into a Celery task that runs in the public schema of a node: the node that ``for_node()`` is given,
+    else the node that its ``node_name`` keyword argument identifies. That schema is the target of the task. The
+    function receives ``node_name`` only when it declares it by name, and then also when the task is sent without
+    ``node_name``. ``shard_id`` and the keyword argument for a mapping value are ordinary keyword arguments for the
+    function. Sending the task without a node raises TaskNodeError. The public schema of a node has no lock, so the task
+    takes no ``lock``. ``options`` work as they do for :func:`task_for_shard`.
+
+    :Example:
+        .. code-block:: python
+
+            from djanquiltdb.decorators.celery import task_for_node
+
+
+            @task_for_node
+            def refresh_exchange_rates():
+                ExchangeRate.objects.refresh()
+
+
+            refresh_exchange_rates.for_node('other').delay()
+            refresh_exchange_rates.delay(node_name='other')  # the same
+    """
+    return _task_decorator_or_task({'runs_in': 'node'}, func, options)
+
+
+def task_for_active_shard(func=None, **options):
+    """
+    :func:`task_for_shard` for a task that runs in the shard that is active where the task is sent: inside
+    ``use_shard()``, ``use_shard_for()`` or a request that the DjanQuiltDB middleware wraps. The shard is resolved when
+    the task is sent, and is sent in the message header. The shard of ``for_shard()`` and a shard in the keyword
+    arguments take precedence over the active shard. Sending the task without a shard while no shard is active, or while
+    only the public schema of a node is active, raises TaskShardError.
+
+    :Example:
+        .. code-block:: python
+
+            from djanquiltdb.decorators.celery import task_for_active_shard
+
+
+            @task_for_active_shard
+            def recount_invoices():
+                Invoice.objects.recount()
+
+
+            with use_shard(shard):
+                recount_invoices.delay()  # runs in shard
+                recount_invoices.delay_for_shard(other_shard)  # runs in other_shard
+    """
+    return task_for_shard(func, inherit_active_shard=True, **options)
+
+
+def task_for_every_shard(func=None, *, queue=None, spread=None, include_maintenance=False, lock=None, **options):
+    """
+    :func:`task_for_shard` for a task that runs once in every active shard when it is sent without a shard. The message
+    that is sent without a shard starts a fan-out where it runs: the task sends itself once for each active shard, with
+    the ``shard_id`` of the shard added to the keyword arguments. So a periodic task in the beat schedule needs only the
+    name of the task. A task that is sent with a shard, in ``for_shard()`` or in its keyword arguments, runs in that
+    shard only. The task does not inherit the active shard, so that it also runs in every active shard when it is sent
+    inside ``use_shard()``.
+
+    The function receives ``shard_id`` only when it declares it by name, as for :func:`task_for_shard`. A function that
+    declares ``shard_id`` can declare it without a default, because the check of the arguments of a message counts
+    ``shard_id`` as given. The parameters after ``shard_id`` are passed by keyword, as ``shard_id`` is.
+
+    :param queue: The name of the queue of each task that the fan-out sends. Defaults to
+        ``QUILT_CELERY['FANOUT_QUEUE']``, else the routing of the Celery app, whose ``task_routes`` can route each task
+        by its ``shard_id``. A queue that is not a string raises ValueError.
+    :param spread: Delays each task that the fan-out sends by a random number of seconds: from 0 to an integer, or
+        between a pair of integers, which are not negative and in ascending order. Defaults to
+        ``QUILT_CELERY['FANOUT_SPREAD']``. Any other value raises ValueError.
+    :param include_maintenance: When True, the fan-out also sends the task to the shards in maintenance, where it is
+        retried as in any shard in maintenance, and fails with the StateException after the last retry.
+
+    :Example:
+        .. code-block:: python
+
+            from djanquiltdb.decorators.celery import task_for_every_shard
+
+
+            @task_for_every_shard(queue='fanout', spread=(0, 300))
+            def clean_up_drafts():
+                Draft.objects.filter(updated_at__lt=a_month_ago()).delete()
+
+
+            clean_up_drafts.delay()  # runs once in every active shard
+    """
+    task_attributes = {
+        'runs_in': 'shard',
+        'lock_target': lock,
+        'fans_out': True,
+        'fanout_queue': queue,
+        'fanout_spread': spread,
+        'fanout_include_maintenance': include_maintenance,
+    }
+
+    return _task_decorator_or_task(task_attributes, func, options)
+
+
+def task_for_every_node(func=None, *, queue=None, spread=None, **options):
+    """
+    :func:`task_for_node` for a task that runs once in the public schema of every node when it is sent without a
+    ``node_name``. The message that is sent without a ``node_name`` starts a fan-out where it runs: the task sends
+    itself once for each node in ``DATABASES``, with the ``node_name`` of the node added to the keyword arguments. A
+    task that is sent with a node, in ``for_node()`` or in its keyword arguments, runs on that node only. The function
+    receives ``node_name`` only when it declares it by name. A function that declares ``node_name`` can declare it
+    without a default, and the parameters after ``node_name`` are passed by keyword. ``queue`` and ``spread`` work as
+    they do for :func:`task_for_every_shard`.
+
+    :Example:
+        .. code-block:: python
+
+            from djanquiltdb.decorators.celery import task_for_every_node
+
+
+            @task_for_every_node
+            def refresh_exchange_rates():
+                ExchangeRate.objects.refresh()
+    """
+    task_attributes = {
+        'runs_in': 'node',
+        'fans_out': True,
+        'fanout_queue': queue,
+        'fanout_spread': spread,
+    }
+
+    return _task_decorator_or_task(task_attributes, func, options)
+
+
+__all__ = ['task_for_active_shard', 'task_for_every_node', 'task_for_every_shard', 'task_for_node', 'task_for_shard']

djanquiltdb_plugin_celery/settings.py

@@ -0,0 +1,75 @@
+"""
+The ``QUILT_CELERY`` setting, an optional dict whose keys override the defaults below.
+"""
+
+from django.conf import settings
+
+QUILT_CELERY_DEFAULTS = {
+    # The keyword argument that identifies a shard by its mapping value, which is not shard_id or node_name. None means
+    # the mapping field of the MAPPING_MODEL.
+    'MAPPING_KWARG': None,
+    # Celery's max_retries, retry_backoff, retry_backoff_max and retry_jitter options for a task whose shard is in
+    # maintenance. A decorator argument overrides each of them. The defaults retry after 5, 10, 20, 40 and 60 seconds.
+    'MAX_RETRIES': 5,
+    'RETRY_BACKOFF': 5,
+    'RETRY_BACKOFF_MAX': 60,
+    'RETRY_JITTER': False,
+    # The queue of each task that a fan-out sends. None means the routing of the Celery app.
+    'FANOUT_QUEUE': None,
+    # The countdown of each task that a fan-out sends: a random number of seconds from 0 to an integer, or between a
+    # pair of integers, which are not negative and in ascending order. None means no countdown.
+    'FANOUT_SPREAD': None,
+}
+
+#: How a message describes the values that a spread accepts.
+SPREAD_DESCRIPTION = 'an integer that is not negative, a pair of such integers in ascending order, or None'
+
+#: How a message describes the values that a queue accepts.
+QUEUE_DESCRIPTION = 'a string or None'
+
+
+def quilt_celery_setting(name):
+    """
+    Returns the value of ``name`` in QUILT_CELERY, else the default of ``name``. When QUILT_CELERY is not a dict, the
+    function returns the default, so that a tasks module can be imported and the system check can report the setting.
+    """
+    quilt_celery = getattr(settings, 'QUILT_CELERY', {})
+
+    if not isinstance(quilt_celery, dict):
+        quilt_celery = {}
+
+    return quilt_celery.get(name, QUILT_CELERY_DEFAULTS[name])
+
+
+def is_spread(value):
+    """
+    Returns whether ``value`` is a spread of the countdown of a fan-out: None, an integer that is not negative, or a
+    pair of such integers in ascending order. A boolean is an instance of int, and is not a spread.
+    """
+
+    def is_bound(bound):
+        return isinstance(bound, int) and not isinstance(bound, bool) and bound >= 0
+
+    if value is None:
+        return True
+
+    if isinstance(value, int):
+        return is_bound(value)
+
+    return isinstance(value, (list, tuple)) and len(value) == 2 and all(map(is_bound, value)) and value[0] <= value[1]
+
+
+def ensure_valid_spread(spread, owner):
+    """
+    Raises ValueError, whose message includes ``owner``, when ``spread`` is not a spread of the countdown of a fan-out.
+    """
+    if not is_spread(spread):
+        raise ValueError('The spread of {} is {!r}, and it accepts {}.'.format(owner, spread, SPREAD_DESCRIPTION))
+
+
+def ensure_valid_queue(queue, owner):
+    """
+    Raises ValueError, whose message includes ``owner``, when ``queue`` is neither a string nor None.
+    """
+    if not isinstance(queue, (str, type(None))):
+        raise ValueError('The queue of {} is {!r}, and it accepts {}.'.format(owner, queue, QUEUE_DESCRIPTION))