diff --git a/asyncpg/protocol/protocol.pyx b/asyncpg/protocol/protocol.pyx index 61d48449..f7445638 100644 --- a/asyncpg/protocol/protocol.pyx +++ b/asyncpg/protocol/protocol.pyx @@ -510,8 +510,13 @@ cdef class BaseProtocol(CoreProtocol): else: raise apg_exc.InternalClientError('TimoutError was not raised') - except (Exception, asyncio.CancelledError) as e: - self._write_copy_fail_msg(str(e)) + except (Exception, asyncio.CancelledError): + # CopyFail is limited to 10000 bytes, including its length and + # trailing NUL. Use a fixed reason so formatting or encoding the + # application exception cannot break COPY cleanup. The original + # exception is re-raised below. + self._write_copy_fail_msg( + 'COPY aborted due to an exception in input generator') self._request_cancel() # Make asyncio shut up about unretrieved QueryCanceledError waiter.add_done_callback(lambda f: f.exception()) diff --git a/tests/test_copy.py b/tests/test_copy.py index ed6b2223..724b23a1 100644 --- a/tests/test_copy.py +++ b/tests/test_copy.py @@ -700,6 +700,62 @@ async def record_generator(): finally: await self.con.execute('DROP TABLE copytab_async') + async def test_copy_records_to_table_generator_error(self): + # A 9996-byte reason exceeds PostgreSQL's CopyFail limit. + await self._test_copy_records_generator_error( + ValueError('a' * 9996), asynchronous=False, transaction=True) + + async def test_copy_records_to_table_unprintable_error(self): + class UnprintableError(ValueError): + def __str__(self): + raise AssertionError('cannot format exception') + + await self._test_copy_records_generator_error( + UnprintableError(), asynchronous=True, transaction=False) + + async def _test_copy_records_generator_error( + self, error, asynchronous, transaction): + con = await self.connect() + + def records(): + # Flush a CopyData message before the generator fails. + yield (b'x' * 524288,) + raise error + + async def async_records(): + for row in records(): + yield row + + async def copy(): + await con.copy_records_to_table( + 'copytab', + records=async_records() if asynchronous else records()) + + async def run(): + await con.execute('CREATE TEMP TABLE copytab (a bytea)') + with self.assertRaises(ValueError) as caught: + if transaction: + async with con.transaction(): + await copy() + else: + await copy() + self.assertIs(caught.exception, error) + self.assertFalse(con.is_in_transaction()) + self.assertEqual( + await con.fetchval('SELECT count(*) FROM copytab'), 0) + self.assertEqual(await con.copy_records_to_table( + 'copytab', records=[(b'recovered',)]), 'COPY 1') + self.assertEqual( + await con.fetchval('SELECT a FROM copytab'), b'recovered') + await con.close() + + try: + # Bound recovery and close as well as COPY itself: cancellation + # completion used to hang after an oversized CopyFail message. + await asyncio.wait_for(run(), timeout=5) + finally: + con.terminate() + async def test_copy_records_to_table_stmt_cache(self): # The introspection statement must be taken from the statement # cache, otherwise every call pays for an extra round-trip.