CODEWITHSUNDEEP
pythonexceptionwith-statementcontextmanager
Michael Waterfall
Michael Waterfall

Reputation: 20569

Encapsulating retries into `with` block

I'm looking to encapsulate logic for database transactions into a with block; wrapping the code in a transaction and handling various exceptions (locking issues). This is simple enough, however I'd like to also have the block encapsulate the retrying of the code block following certain exceptions. I can't see a way to package this up neatly into the context manager.

Is it possible to repeat the code within a with statement?

I'd like to use it as simply as this, which is really neat.

def do_work():
    ...
    # This is ideal!
    with transaction(retries=3):
        # Atomic DB statements
        ...
    ...

I'm currently handling this with a decorator, but I'd prefer to offer the context manager (or in fact both), so I can choose to wrap a few lines of code in the with block instead of an inline function wrapped in a decorator, which is what I do at the moment:

def do_work():
    ...
    # This is not ideal!
    @transaction(retries=3)
    def _perform_in_transaction():
        # Atomic DB statements
        ...
    _perform_in_transaction()
    ...

Upvotes: 41

Views: 12470

Answers (5)

gtalarico
gtalarico

Reputation: 4689

This question is a few years old but after reading the answers I decided to give this a shot.

This solution requires the use of a "helper" class, but I I think it does provide an interface with retries configured through a context manager.

class Client:
    def _request(self):
        # do request stuff
        print("tried")
        raise Exception()

    def request(self):
        retry = getattr(self, "_retry", None)
        if not retry:
            return self._request()
        else:
            for n in range(retry.tries):
                try:
                    return self._request()
                except Exception:
                    retry.attempts += 1


class Retry:
    def __init__(self, client, tries=1):
        self.client = client
        self.tries = tries
        self.attempts = 0

    def __enter__(self):
        self.client._retry = self

    def __exit__(self, *exc):
        print(f"Tried {self.attempts} times")
        del self.client._retry



>>> client = Client()
>>> with Retry(client, tries=3):
    ... # will try 3 times
    ... response = client.request()

tried once
tried once
tried once
Tried 3 times

Upvotes: -1

Oddthinking
Oddthinking

Reputation: 25282

While I agree it can't be done with a context manager... it can be done with two context managers!

The result is a little awkward, and I am not sure whether I approve of my own code yet, but this is what it looks like as the client:

with RetryManager(retries=3) as rm:
    while rm:
        with rm.protect:
            print("Attempt #%d of %d" % (rm.attempt_count, rm.max_retries))
             # Atomic DB statements

There is an explicit while loop still, and not one, but two, with statements, which leaves a little too much opportunity for mistakes for my liking.

Here's the code:

class RetryManager(object):
    """ Context manager that counts attempts to run statements without
        exceptions being raised.
        - returns True when there should be more attempts
    """

    class _RetryProtector(object):
        """ Context manager that only raises exceptions if its parent
            RetryManager has given up."""
        def __init__(self, retry_manager):
            self._retry_manager = retry_manager

        def __enter__(self):
            self._retry_manager._note_try()
            return self

        def __exit__(self, exc_type, exc_val, traceback):
            if exc_type is None:
                self._retry_manager._note_success()
            else:
                # This would be a good place to implement sleep between
                # retries.
                pass

            # Suppress exception if the retry manager is still alive.
            return self._retry_manager.is_still_trying()

    def __init__(self, retries=1):

        self.max_retries = retries
        self.attempt_count = 0 # Note: 1-based.
        self._success = False

        self.protect = RetryManager._RetryProtector(self)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, traceback):
        pass

    def _note_try(self):
        self.attempt_count += 1

    def _note_success(self):
        self._success = True

    def is_still_trying(self):
        return not self._success and self.attempt_count < self.max_retries

    def __bool__(self):
        return self.is_still_trying()

Bonus: I know you don't want to separate your work off into separate functions wrapped with decorators... but if you were happy with that, the redo package from Mozilla offers the decorators to do that, so you don't have to roll your own. There is even a Context Manager that effective acts as temporary decorator for your function, but it still relies on your retrievable code to be factored out into a single function.

Upvotes: 3

Henry Keiter
Henry Keiter

Reputation: 17168

The way that occurs to me to do this is just to implement a standard database transaction context manager, but allow it to take a retries argument in the constructor. Then I'd just wrap that up in your method implementations. Something like this:

class transaction(object):
    def __init__(self, retries=0):
        self.retries = retries
    def __enter__(self):
        return self
    def __exit__(self, exc_type, exc_val, traceback):
        pass

    # Implementation...
    def execute(self, query):
        err = None
        for _ in range(self.retries):
            try:
                return self._cursor.execute(query)
            except Exception as e:
                err = e # probably ought to save all errors, but hey
        raise err

with transaction(retries=3) as cursor:
    cursor.execute('BLAH')

Upvotes: 7

JAB
JAB

Reputation: 21079

As decorators are just functions themselves, you could do the following:

with transaction(_perform_in_transaction, retries=3) as _perf:
    _perf()

For the details, you'd need to implement transaction() as a factory method that returns an object with __callable__() set to call the original method and repeat it up to retries number of times on failure; __enter__() and __exit__() would be defined as normal for database transaction context managers.

You could alternatively set up transaction() such that it itself executes the passed method up to retries number of times, which would probably require about the same amount of work as implementing the context manager but would mean actual usage would be reduced to just transaction(_perform_in_transaction, retries=3) (which is, in fact, equivalent to the decorator example delnan provided).

Upvotes: 4

user395760
user395760

Reputation:

Is it possible to repeat the code within a with statement?

No.

As pointed out earlier in that mailing list thread, you can reduce a bit of duplication by making the decorator call the passed function:

def do_work():
    ...
    # This is not ideal!
    @transaction(retries=3)
    def _perform_in_transaction():
        # Atomic DB statements
        ...
    # called implicitly
    ...

Upvotes: 15

Related Questions