Fix COPY cleanup when a records generator raises (#1358)

Use a fixed short CopyFail reason instead of formatting the application
exception. PostgreSQL limits this message to 10000 bytes, including its
length field and trailing NUL. An oversized reason loses protocol
synchronisation and can leave transaction rollback waiting indefinitely.

Keep the original application exception intact, including when formatting
or encoding it would fail. Cover encoded-byte boundaries, multibyte text,
sync and async generators, COPY atomicity, rollback and connection reuse.

Fixes #1354

Co-authored-by: Elvis Pranskevichus <el@prans.net>

Author: CodingCossack

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())

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.