Refining exceptions with Context Decorators
Tags: python, celery, boto3, django
Published: 2020-01-07 17:59:31

Very often, 3rd party API clients are designed to return a generic "API Exception", with details contained inside some property. One notorious example of this is `boto3`: most exceptions come as a `ClientError` with a `response`.

Usually this is fine, but occasionally I need finer-grained exceptions than what the client is giving me.

One common scenario is configuring a celery task to retry a while after the rate limiting: I want to get errors for all exceptions, but for rate-limiting, I just want to retry the task later. Celery offers a convenient `autoretry_for` argument for that, but we can't use it _just_ for throttling because `boto3` does not return an exception specific enough.


I could wrap the login in a `try ... except` clause and inspect the exception, but that get repetitive pretty quickly, especially as I add more and more tasks.


For these kind of situations, I create a [context decorator](https://docs.python.org/3/library/contextlib.html#contextlib.ContextDecorator). The decorator inspects the exception for me and, if it's the one I'm looking, raises a custom exception that then I can catch:


```python
# myproject/exceptions.py

import contexlib

from botocore.exceptions import ClientError


class ThrottleException(ClientError):
    pass


class raise_throttle(contextlib.ContextDecorator):
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, tb):
        if exc_type == ClientError:
            if exc_value.response.get("Error", {}).get("Code") == "Throttling":
                raise ThrottleException(
                    exc_value.response, exc_value.operation_name
                ) from exc_value
        # returning `False` makes the decorator
        # raise the original exception, if any.
        return False
```

This decorator can be used on any function or method that calls the API:

```python
# myproject/myapp/aws.py
import boto3


@raise_throttle()
def upload(content):
    s3 = boto3.resource("s3")
    bucket = s3.Bucket("my-bucket")
    bucket.put_object(
        Body=b"lorem ipsum",
        Key="Hamlet.txt",
    )
```

And the exception can be caught by celery's `autoretry_for`:

```python
# myproject/myapp/tasks.py

from myproject.celery import app
from myproject.exceptions import ThrottleException
from myproject.myapp.aws import upload


@app.task(autoretry_for=(ThrottleException,))
def my_task(content):
    upload(content)
```

But I don't want to retry the task immediately, since the throttling takes a while to be lifted. Therefore I create a [Task base class](https://docs.celeryproject.org/en/latest/userguide/tasks.html#task-inheritance) that retries with exponential backoff:

```python
# myproject/tasks.py

import random

from celery import Task


def jitter(jitter_max=1.4):
    return random.uniform(1, jitter_max)


def exponential(retries, factor=3):
    return (factor ** (retries + 1))


def exp_jitter(retries, exp_factor=3, jitter_max=1.4):
    return exponential(retries, exp_factor) * jitter(jitter_max)


class ExpBackoffTask(Task):
    abstract = True
    max_retries = 10

    def retry(self, *args, **kwargs):
        countdown = kwargs.get('countdown', None)
        if countdown is None:
            # if no explicit countdown is given, use an exponential backoff
            # with some random jitter, giving us a top-end under 24 hours at
            # which the last retry will be attempted.
            kwargs['countdown'] = int(
                exp_jitter(self.request.retries)
            )
        return super().retry(*args, **kwargs)
```

I can then use the custom class as base:

```python
# myproject/myapp/tasks.py

from myproject.celery import app
from myproject.exceptions import ThrottleException
from myproject.tasks import ExpBackoffTask
from myproject.myapp.aws import upload


@app.task(base=ExpBackoffTask, autoretry_for=(ThrottleException,))
def my_task(content):
    upload(content)
```