Reset Connection._top_xact; disallow managed xacts inside manual ones

Author: 1st1

asyncpg/_testbase.py

@@ -12,6 +12,7 @@
 import inspect
 import logging
 import os
+import re
 import time
 import unittest
 
@@ -95,6 +96,30 @@ def assertRunUnder(self, delta):
  raise AssertionError(
  'running block took longer than {}'.format(delta))
 
+ @contextlib.contextmanager
+ def assertLoopErrorHandlerCalled(self, msg_re: str):
+ contexts = []
+
+ def handler(loop, ctx):
+ contexts.append(ctx)
+
+ old_handler = self.loop.get_exception_handler()
+ self.loop.set_exception_handler(handler)
+ try:
+ yield
+
+ for ctx in contexts:
+ msg = ctx.get('message')
+ if msg and re.search(msg_re, msg):
+ return
+
+ raise AssertionError(
+ 'no message matching {!r} was logged with '
+ 'loop.call_exception_handler()'.format(msg_re))
+
+ finally:
+ self.loop.set_exception_handler(old_handler)
+
 
 _default_cluster = None
 

asyncpg/connection.py

@@ -513,11 +513,12 @@ def _get_reset_query(self):
  caps = self._server_caps
 
  _reset_query = []
- if self._protocol.is_in_transaction():
+ if self._protocol.is_in_transaction() or self._top_xact is not None:
  self._loop.call_exception_handler({
  'message': 'Resetting connection with an '
  'active transaction {!r}'.format(self)
  })
+ self._top_xact = None
  _reset_query.append('ROLLBACK;')
  if caps.advisory_locks:
  _reset_query.append('SELECT pg_advisory_unlock_all();')

asyncpg/transaction.py

@@ -84,6 +84,10 @@ async def start(self):
  con = self._connection
 
  if con._top_xact is None:
+ if con._protocol.is_in_transaction():
+ raise apg_errors.InterfaceError(
+ 'cannot use Connection.transaction() in '
+ 'a manually started transaction')
  con._top_xact = self
  else:
  # Nested transaction block

tests/test_pool.py

@@ -353,32 +353,28 @@ async def test_pool_release_in_xact(self):
  """Test that Connection.reset() closes any open transaction."""
  async with self.create_pool(database='postgres',
  min_size=1, max_size=1) as pool:
-
- last_error_ctx = None
-
  async def get_xact_id(con):
  return await con.fetchval('select txid_current()')
 
- def logger(loop, ctx):
- nonlocal last_error_ctx
- last_error_ctx = ctx
-
- async with pool.acquire() as con:
- id1 = await get_xact_id(con)
+ with self.assertLoopErrorHandlerCalled('an active transaction'):
+ async with pool.acquire() as con:
+ real_con = con._con # unwrap PoolConnectionProxy
 
- tr = con.transaction()
- await tr.start()
+ id1 = await get_xact_id(con)
 
- id2 = await get_xact_id(con)
- self.assertNotEqual(id1, id2)
+ tr = con.transaction()
+ self.assertIsNone(con._con._top_xact)
+ await tr.start()
+ self.assertIs(real_con._top_xact, tr)
 
- self.loop.set_exception_handler(logger)
+ id2 = await get_xact_id(con)
+ self.assertNotEqual(id1, id2)
 
- self.assertIsNotNone(last_error_ctx)
- self.assertIn('an active transaction', last_error_ctx['message'])
- self.loop.set_exception_handler(None)
+ self.assertIsNone(real_con._top_xact)
 
  async with pool.acquire() as con:
+ self.assertIs(con._con, real_con)
+ self.assertIsNone(con._con._top_xact)
  id3 = await get_xact_id(con)
  self.assertNotEqual(id2, id3)
 

tests/test_test.py

@@ -33,3 +33,14 @@ def test_tests_fail_1(self):
  suite.run(result)
 
  self.assertIn('ZeroDivisionError', result.errors[0][1])
+
+
+class TestHelpers(tb.TestCase):
+
+ async def test_tests_assertLoopErrorHandlerCalled_01(self):
+ with self.assertRaisesRegex(AssertionError, r'no message.*was logged'):
+ with self.assertLoopErrorHandlerCalled('aa'):
+ self.loop.call_exception_handler({'message': 'bb a bb'})
+
+ with self.assertLoopErrorHandlerCalled('aa'):
+ self.loop.call_exception_handler({'message': 'bbaabb'})

tests/test_transaction.py

@@ -139,3 +139,21 @@ async def test_transaction_interface_errors(self):
  async with tr:
  async with tr:
  pass
+
+ async def test_transaction_within_manual_transaction(self):
+ self.assertIsNone(self.con._top_xact)
+
+ await self.con.execute('BEGIN')
+
+ tr = self.con.transaction()
+ self.assertIsNone(self.con._top_xact)
+
+ with self.assertRaisesRegex(asyncpg.InterfaceError,
+ 'cannot use Connection.transaction'):
+ await tr.start()
+
+ with self.assertLoopErrorHandlerCalled(
+ 'Resetting connection with an active transaction'):
+ await self.con.reset()
+
+ self.assertIsNone(self.con._top_xact)