Improve wrap_query_exceptions code

Based on @ewjoachim's feedback.
procrastinate-org · Jun 18, 2020 · 082b819 · 082b819
1 parent 894b4b6
commit 082b819
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/procrastinate/aiopg_connector.py b/procrastinate/aiopg_connector.py
@@ -58,9 +58,10 @@ def wrap_query_exceptions(coro: CoroutineFunction) -> CoroutineFunction:
     @functools.wraps(coro)
     async def wrapped(*args, **kwargs):
         final_exc = None
-        max_tries = (
-            args[0]._pool.maxsize + 1 if args and getattr(args[0], "_pool", None) else 1
-        )
+        try:
+            max_tries = args[0]._pool.maxconn + 1
+        except Exception:
+            max_tries = 1
         for _ in range(max_tries):
             try:
                 return await coro(*args, **kwargs)

diff --git a/procrastinate/psycopg2_connector.py b/procrastinate/psycopg2_connector.py
@@ -50,9 +50,10 @@ def wrap_query_exceptions(func: Callable) -> Callable:
     @functools.wraps(func)
     def wrapped(*args, **kwargs):
         final_exc = None
-        max_tries = (
-            args[0]._pool.maxconn + 1 if args and getattr(args[0], "_pool", None) else 1
-        )
+        try:
+            max_tries = args[0]._pool.maxconn + 1
+        except Exception:
+            max_tries = 1
         for _ in range(max_tries):
             try:
                 return func(*args, **kwargs)