Execute handle_error()/handle_error_async() on exception

Author: seankerr

postgres_lock/asyncpg.py

@@ -45,16 +45,22 @@ async def acquire_async(lock: Lock, block: bool = True) -> bool:
  return False if result is False else True
 
 
-def handle_error(lock: Lock) -> None:
+def handle_error(lock: Lock, exc: BaseException) -> None:
  """
  Handle an error.
+
+ Parameters:
+ exc (Exception): Exception.
  """
  raise NotImplementedError("ascynpg interface does not support handle_error()")
 
 
-async def handle_error_async(lock: Lock) -> None:
+async def handle_error_async(lock: Lock, exc: BaseException) -> None:
  """
  Handle an error asynchronously.
+
+ Parameters:
+ exc (Exception): Exception.
  """
  if not lock.rollback_on_error:
  return

postgres_lock/lock.py

@@ -151,17 +151,23 @@ async def acquire_async(self, block: bool = True) -> bool:
 
  return self._locked
 
- def handle_error(self) -> None:
+ def handle_error(self, exc: BaseException) -> None:
  """
  Handle an error.
+
+ Parameters:
+ exc (Exception): Exception.
  """
- return self.impl.handle_error(self)
+ return self.impl.handle_error(self, exc)
 
- async def handle_error_async(self) -> None:
+ async def handle_error_async(self, exc: BaseException) -> None:
  """
  Handle an error asynchronously.
+
+ Parameters:
+ exc (Exception): Exception.
  """
- return await self.impl.handle_error_async(self)
+ return await self.impl.handle_error_async(self, exc)
 
  @property
  def locked(self) -> bool:
@@ -304,6 +310,9 @@ async def __aexit__(
  Exception (BaseException): Exception that was raised.
  Traceback (TracebackType): Traceback.
  """
+ if exc:
+ await self.impl.handle_error_async(self, exc)
+
  await self.impl.release_async(self)
 
  if exc:
@@ -329,6 +338,9 @@ def __exit__(
  Exception (BaseException): Exception that was raised.
  Traceback (TracebackType): Traceback.
  """
+ if exc:
+ self.impl.handle_error(self, exc)
+
  self.impl.release(self)
 
  if exc:

postgres_lock/psycopg2.py

@@ -48,19 +48,25 @@ async def acquire_async(lock: Lock, block: bool = True) -> bool:
  raise NotImplementedError("psycopg2 interface does not support acquire_async()")
 
 
-def handle_error(lock: Lock) -> None:
+def handle_error(lock: Lock, exc: BaseException) -> None:
  """
  Handle an error.
+
+ Parameters:
+ exc (Exception): Exception.
  """
  if not lock.rollback_on_error:
  return
 
  lock.conn.rollback()
 
 
-async def handle_error_async(lock: Lock) -> None:
+async def handle_error_async(lock: Lock, exc: BaseException) -> None:
  """
  Handle an error asynchronously.
+
+ Parameters:
+ exc (Exception): Exception.
  """
  raise NotImplementedError(
  "psycopg2 interface does not support handle_error_async()"

postgres_lock/psycopg3.py

@@ -28,6 +28,7 @@ def acquire(lock: Lock, block: bool = True) -> bool:
  cursor = lock.conn.cursor()
  cursor.execute(f"SELECT pg_catalog.{lock_func}({lock.lock_id})")
  result, *_ = cursor.fetchone()
+ cursor.close()
 
  # lock function returns True/False in unblocking mode, and always None in blocking mode
  return False if result is False else True
@@ -54,24 +55,31 @@ async def acquire_async(lock: Lock, block: bool = True) -> bool:
  await cursor.execute(f"SELECT pg_catalog.{lock_func}({lock.lock_id})")
 
  result, *_ = await cursor.fetchone()
+ cursor.close()
 
  # lock function returns True/False in unblocking mode, and always None in blocking mode
  return False if result is False else True
 
 
-def handle_error(lock: Lock) -> None:
+def handle_error(lock: Lock, exc: BaseException) -> None:
  """
  Handle an error.
+
+ Parameters:
+ exc (Exception): Exception.
  """
  if not lock.rollback_on_error:
  return
 
  lock.conn.rollback()
 
 
-async def handle_error_async(lock: Lock) -> None:
+async def handle_error_async(lock: Lock, exc: BaseException) -> None:
  """
  Handle an error asynchronously.
+
+ Parameters:
+ exc (Exception): Exception.
  """
  if not lock.rollback_on_error:
  return
@@ -92,6 +100,7 @@ def release(lock: Lock) -> bool:
  cursor = lock.conn.cursor()
  cursor.execute(f"SELECT pg_catalog.{lock.unlock_func}({lock.lock_id})")
  result, *_ = cursor.fetchone()
+ cursor.close()
 
  return result
 
@@ -111,5 +120,6 @@ async def release_async(lock: Lock) -> bool:
  await cursor.execute(f"SELECT pg_catalog.{lock.unlock_func}({lock.lock_id})")
 
  result, *_ = await cursor.fetchone()
+ cursor.close()
 
  return result

postgres_lock/sqlalchemy.py

@@ -59,19 +59,25 @@ async def acquire_async(lock: Lock, block: bool = True) -> bool:
  return False if result is False else True
 
 
-def handle_error(lock: Lock) -> None:
+def handle_error(lock: Lock, exc: BaseException) -> None:
  """
  Handle an error.
+
+ Parameters:
+ exc (Exception): Exception.
  """
  if not lock.rollback_on_error:
  return
 
  lock.conn.rollback()
 
 
-async def handle_error_async(lock: Lock) -> None:
+async def handle_error_async(lock: Lock, exc: BaseException) -> None:
  """
  Handle an error asynchronously.
+
+ Parameters:
+ exc (Exception): Exception.
  """
  if not lock.rollback_on_error:
  return

tests/test_asyncpg.py

@@ -87,7 +87,7 @@ async def test_acquire_async__block_true(result):
 
 def test_handle_error():
  with raises(NotImplementedError) as exc:
- handle_error(None)
+ handle_error(None, None)
 
  assert str(exc.value) == "ascynpg interface does not support handle_error()"
 
@@ -97,14 +97,14 @@ async def test_handle_error_async():
  lock = Mock()
  lock.conn.execute = AsyncMock()
 
- await handle_error_async(lock)
+ await handle_error_async(lock, None)
 
 
 @mark.asyncio
 async def test_handle_error_async__rollback_disabled():
  lock = Mock(rollback_on_error=False)
 
- await handle_error_async(lock)
+ await handle_error_async(lock, None)
 
 
 def test_release():

tests/test_lock.py

@@ -204,6 +204,7 @@ def test_context_manager__raises_exception(_load_impl):
  assert str(exc.value) == "Inner exception"
 
  lock.impl.acquire.assert_called_with(lock, block=True)
+ lock.impl.handle_error.assert_called_with(lock, exc.value)
  lock.impl.release.assert_called_with(lock)
 
 
@@ -226,6 +227,7 @@ async def test_context_manager_async(_load_impl):
 async def test_context_manager_async__raises_exception(_load_impl):
  lock = Lock(None, "key")
  lock.impl.acquire_async = AsyncMock()
+ lock.impl.handle_error_async = AsyncMock()
  lock.impl.release_async = AsyncMock()
 
  with raises(Exception) as exc:
@@ -235,27 +237,32 @@ async def test_context_manager_async__raises_exception(_load_impl):
  assert str(exc.value) == "Inner exception"
 
  lock.impl.acquire_async.assert_called_with(lock, block=True)
+ lock.impl.handle_error_async.assert_called_with(lock, exc.value)
  lock.impl.release_async.assert_called_with(lock)
 
 
 @patch(f"{PATH}.Lock._load_impl")
 def test_handle_error(_load_impl):
  lock = Lock(None, "key")
+ exc = Exception()
 
- assert lock.impl.handle_error.return_value == lock.handle_error()
+ assert lock.impl.handle_error.return_value == lock.handle_error(exc)
 
- lock.impl.handle_error.assert_called_with(lock)
+ lock.impl.handle_error.assert_called_with(lock, exc)
 
 
 @mark.asyncio
 @patch(f"{PATH}.Lock._load_impl")
 async def test_handle_error_async(_load_impl):
  lock = Lock(None, "key")
  lock.impl.handle_error_async = AsyncMock()
+ exc = Exception()
 
- assert lock.impl.handle_error_async.return_value == await lock.handle_error_async()
+ assert lock.impl.handle_error_async.return_value == await lock.handle_error_async(
+ exc
+ )
 
- lock.impl.handle_error_async.assert_called_with(lock)
+ lock.impl.handle_error_async.assert_called_with(lock, exc)
 
 
 @patch(f"{PATH}.Lock._load_impl")

tests/test_psycopg2.py

@@ -30,6 +30,7 @@ def test_acquire__defaults(result):
  assert not acquire(lock)
 
  cursor.execute.assert_called_with(f"SELECT pg_catalog.{lock_func}({lock.lock_id})")
+ cursor.close.assert_called_once()
 
 
 @mark.parametrize("result", [None, True, False])
@@ -72,6 +73,7 @@ def test_acquire__block_true(result):
  assert not acquire(lock, block=True)
 
  cursor.execute.assert_called_with(f"SELECT pg_catalog.{lock_func}({lock.lock_id})")
+ cursor.close.assert_called_once()
 
 
 @mark.asyncio
@@ -85,21 +87,21 @@ async def test_acquire_async():
 def test_handle_error():
  lock = Mock()
 
- handle_error(lock)
+ handle_error(lock, None)
 
  lock.conn.rollback.assert_called_once()
 
 
 def test_handle_error__rollback_disabled():
  lock = Mock(rollback_on_error=False)
 
- handle_error(lock)
+ handle_error(lock, None)
 
 
 @mark.asyncio
 async def test_handle_error_async():
  with raises(NotImplementedError) as exc:
- await handle_error_async(None)
+ await handle_error_async(None, None)
 
  assert str(exc.value) == "psycopg2 interface does not support handle_error_async()"
 
@@ -117,6 +119,7 @@ def test_release(result):
  cursor.execute.assert_called_with(
  f"SELECT pg_catalog.{lock.unlock_func}({lock.lock_id})"
  )
+ cursor.close.assert_called_once()
 
 
 @mark.asyncio