Django @retry_transaction decorator

retry_transaction_decorator.py

from django.db import transaction, IntegrityError, DatabaseError

from functools import wraps
import logging

_transaction_logger = logging.getLogger("db.transactions")

def retry_transaction(attempts=2, on_failure=None, logger=_transaction_logger):
	"""
	Decorator for catching db errors and retrying the operation.
	
	There are a few reasons a transaction may fail:
	- Serialization errors or the like, which should of course be retried
	- Many db errors have equivalent checks in django code - as db errors raised may well be caused by
	  something changing _after_ the django checks have been made but before the transaction has been
	  committed. Trying it again should at least give us the proper error.

	Being careful to also enforce a maximum number of attempts lest something else is causing the DatabaseErrors
	and we go into an infinite loop. django < 1.7 is stupid and doesn't give us access to the pgcode of
	DatabaseErrors which would allow us to actually distinguish the serialization errors we're on the lookout
	for from other DatabaseErrors so this is a possibility.
	"""
	def decorator(func):
		def wrapping_function(*args, **kwargs):
			attempts_left = attempts
			while True:
				try:
					return func(*args, **kwargs)
				except (transaction.TransactionManagementError, DatabaseError), e:
					attempts_left -= 1
					if callable(on_failure):
						ret = on_failure(e, attempts_left)
						try:
							# allow an integer-convertable return value to override our attempts_left
							attempts_left = int(ret)
						except:
							pass
					if attempts_left > 0:
						logger.warning("""Caught %r in retry_transaction decorator, retrying...""", e)
						continue
					else:
						raise e
		
		return wraps(func)(wrapping_function)
	
	return decorator