Encapsulating retries into `with` block
Asked Answered
M

5

41

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()
    ...
Multiplier answered 4/6, 2013 at 13:46 Comment(1)
docs.python.org/release/2.5/whatsnew/pep-343.html looks like it has examples on how to implement a context manager.Christmann
V
15

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
    ...
Vannesavanness answered 4/6, 2013 at 13:57 Comment(1)
Ah, shame it's not supported. Thanks for the link to the thread. I like the idea of having the call implicit to make it cleaner. If I want to set/modify vars within _perform_in_transaction, I guess I'll have to call it manually anyway and return what I need to continue the rest of the do_work function.Multiplier
C
7

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')
Collin answered 4/6, 2013 at 14:0 Comment(7)
Can you elaborate where _cursor in self._cursor comes from?Phatic
@MikeMüller I'm trying to draw on some common database API customs without getting bogged down in implementation details. _cursor is meant to be a Cursor object, as appropriate for the particular database connection involved. A full implementation would need to create and contain a Connection object of some kind, in order to actually carry out the database transactions.Collin
@HenryKeller I would have done something like this def __init__(self, cursor, retries=0): and inside the __init__ this self._cursor = cursor'. Usage: with transaction(cursor, retries=3) as cursor:`. Does this make sense?Phatic
@MikeMüller Certainly. That's what I mean by my comment regarding "a full implementation": to do this fully, it would probably be best either to make a space in the constructor for a Connection or Cursor, or else make the constructor something like def __init__(self, dbhost, dbname, user, password): and create a Connection object from there. I'm not including that stuff in the answer because it's not really relevant to the OP's question, which is specifically about repeating code automatically with a context manager, not creating a DB context manager in the first place.Collin
This is a rather elaborate way to pass 3 as a parameter to execute(). Why not just def execute(self, query, retries=1), and call it cursor.execute(query, 3)?Dorella
@Dorella You're right that that's equivalent. The reason (in my mind) for making it a transaction parameter is just so that you don't have to specify retries for every query you might want to execute within that transaction.Collin
@HenryKeiter: Ah, I see. Thanks.Dorella
P
4

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).

Philomel answered 4/6, 2013 at 13:52 Comment(0)
D
3

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.

Dorella answered 17/9, 2018 at 5:27 Comment(2)
Hmm... I think we could simplify this slightly to just be a for retry in Retry(retries=3): with retry:, so only two outer blocks.Holcman
@user295691: I can't fault that. Seems like an improvement.Dorella
B
-1

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
Brachiopod answered 16/8, 2020 at 6:41 Comment(1)
It seems as an overcomplication for me: moving tries and attempts into Client class makes Retry helper redundant. Moreover, this will make the whole implementation readability much better ;)Anagrammatize

© 2022 - 2024 — McMap. All rights reserved.