Make the undocumented retry_on_deadlock function more useful.

CHANGES.rst

@@ -1,6 +1,13 @@
 Changelog
 =========
 
+Version 0.3.5
+-------------
+
+- Make the undocumented `retry_on_deadlock` function more useful.
+- Add new tests
+
+
 Version 0.3.4
 -------------
 

flask_signalbus/__init__.py

@@ -18,6 +18,10 @@
 FLUSHMANY_LIMIT = 1000
 
 
+class DBSerializationError(Exception):
+    """The transaction is rolled back due to a race condition."""
+
+
 def get_db_error_code(exception):
     """Return 5-character SQLSTATE code, or '' if not available.
 
@@ -43,9 +47,13 @@ def f(*args, **kwargs):
             while True:
                 try:
                     return action(*args, **kwargs)
-                except DBAPIError as e:
+                except (DBAPIError, DBSerializationError) as e:
                     num_failures += 1
-                    if num_failures > retries or get_db_error_code(e.orig) not in DEADLOCK_ERROR_CODES:
+                    is_serialization_error = (
+                        isinstance(e, DBSerializationError) or
+                        get_db_error_code(e.orig) in DEADLOCK_ERROR_CODES,
+                    )
+                    if num_failures > retries or not is_serialization_error:
                         raise
                 session.rollback()
                 wait_seconds = min(max_wait, min_wait * 2 ** (num_failures - 1))

tests/test_basics.py

@@ -1,6 +1,7 @@
+import sys
 import pytest
 import flask_sqlalchemy as fsa
-from flask_signalbus import SignalBus
+from flask_signalbus import SignalBus, retry_on_deadlock, DBSerializationError
 from .conftest import SignalBusAlchemy
 
 
@@ -175,3 +176,17 @@ def test_flush_signal_with_props(db, signalbus, send_mock, Signal, SignalPropert
     assert props['last_name'] == 'Doe'
     assert Signal.query.count() == 0
     assert SignalProperty.query.count() == 0
+
+
+def test_retry_on_deadlock(db):
+    retry = retry_on_deadlock(db.session, retries=5, max_wait=0.0)
+    executions = []
+
+    @retry
+    def f():
+        executions.append(1)
+        raise DBSerializationError
+
+    with pytest.raises(DBSerializationError):
+        f()
+    assert len(executions) == 6