From 4a9f2a967949a56013e278d6d38138dd8b0035e5 Mon Sep 17 00:00:00 2001
From: Dima <dimastasy@protonmail.com>
Date: Sat, 12 Sep 2026 11:43:58 +0100
Subject: [PATCH] Fix COPY cleanup when a records generator raises

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
---
 asyncpg/protocol/protocol.pyx |  8 +++-
 tests/test_copy.py            | 84 +++++++++++++++++++++++++++++++++++
 2 files changed, 90 insertions(+), 2 deletions(-)

diff --git a/asyncpg/protocol/protocol.pyx b/asyncpg/protocol/protocol.pyx
index acce4e9f..f5db8c18 100644
--- a/asyncpg/protocol/protocol.pyx
+++ b/asyncpg/protocol/protocol.pyx
@@ -546,8 +546,12 @@ 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 source failed')
             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 e119e6d8..36d3e683 100644
--- a/tests/test_copy.py
+++ b/tests/test_copy.py
@@ -700,6 +700,90 @@ async def record_generator():
         finally:
             await self.con.execute('DROP TABLE copytab_async')
 
+    async def test_copy_records_to_table_generator_error(self):
+        # CopyFail allows 9995 encoded bytes, excluding its length and NUL.
+        messages = (
+            'failure in source',
+            'a' * 9994,
+            'a' * 9995,
+            'a' * 9996,
+            '\u20ac' * 3331 + 'aa',  # 9995 UTF-8 bytes
+            '\u20ac' * 3332,         # 9996 UTF-8 bytes, only 3332 characters
+            'a' * 9994 + '\u20ac',   # Byte truncation would split a character.
+        )
+
+        for message in messages:
+            for asynchronous in (False, True):
+                for transaction in (False, True):
+                    with self.subTest(chars=len(message),
+                                      bytes=len(message.encode('utf-8')),
+                                      asynchronous=asynchronous,
+                                      transaction=transaction):
+                        await self._test_copy_records_generator_error(
+                            ValueError(message), asynchronous, transaction)
+
+    async def test_copy_records_to_table_unprintable_error(self):
+        class UnprintableError(ValueError):
+            def __str__(self):
+                raise AssertionError('cannot format exception')
+
+        for error in (UnprintableError(), ValueError('\ud800')):
+            with self.subTest(error_type=type(error)):
+                await self._test_copy_records_generator_error(
+                    error, asynchronous=False, transaction=True)
+
+    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
+                await asyncio.sleep(0)
+
+        async def copy():
+            try:
+                await con.copy_records_to_table(
+                    'copytab',
+                    records=async_records() if asynchronous else records())
+            except ValueError as raised:
+                self.assertIs(raised, error)
+                if transaction:
+                    with self.assertRaises(
+                            asyncpg.InFailedSQLTransactionError):
+                        await con.fetchval('SELECT 1')
+                raise
+
+        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_no_binary_codec(self):
         await self.con.execute('''
             CREATE TABLE copytab(a uuid);