From 3275b6338737e01c3401080a271535c14ff6140b Mon Sep 17 00:00:00 2001 From: Riti Grover Date: Wed, 23 Sep 2026 20:06:05 +0530 Subject: [PATCH] Keep the server error when the connection closes mid-operation --- asyncpg/protocol/protocol.pyx | 16 +++++++++++++--- tests/test_execute.py | 26 ++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/asyncpg/protocol/protocol.pyx b/asyncpg/protocol/protocol.pyx index 4e1494e7..3a2d27b1 100644 --- a/asyncpg/protocol/protocol.pyx +++ b/asyncpg/protocol/protocol.pyx @@ -716,9 +716,19 @@ cdef class BaseProtocol(CoreProtocol): cdef _handle_waiter_on_connection_lost(self, cause): if self.waiter is not None and not self.waiter.done(): - exc = apg_exc.ConnectionDoesNotExistError( - 'connection was closed in the middle of ' - 'operation') + msg = 'connection was closed in the middle of operation' + if (self.result_type == RESULT_FAILED and + isinstance(self.result, dict)): + # The server sent an ErrorResponse and then closed the + # connection without a ReadyForQuery (a FATAL error, or + # pgbouncer's query_wait_timeout). Do not lose it. + server_exc = apg_exc_base.PostgresError.new( + self.result, query=self.last_query) + if cause is not None: + server_exc.__cause__ = cause + cause = server_exc + msg = '{}: {}'.format(msg, server_exc.args[0]) + exc = apg_exc.ConnectionDoesNotExistError(msg) if cause is not None: exc.__cause__ = cause self.waiter.set_exception(exc) diff --git a/tests/test_execute.py b/tests/test_execute.py index f8a0e43a..900d940e 100644 --- a/tests/test_execute.py +++ b/tests/test_execute.py @@ -98,6 +98,32 @@ async def test_execute_script_interrupted_terminate(self): self.con.terminate() + async def test_execute_script_interrupted_by_server(self): + # The server reports why it is closing the connection with an + # ErrorResponse and then closes it without a ReadyForQuery (this + # is also what pgbouncer does on query_wait_timeout). The reported + # error must not be lost. + pid = await self.con.fetchval('SELECT pg_backend_pid()') + fut = self.loop.create_task( + self.con.execute('''SELECT pg_sleep(10)''')) + + await asyncio.sleep(0.2) + + other = await self.connect() + try: + await other.execute('SELECT pg_terminate_backend($1)', pid) + finally: + await other.close() + + with self.assertRaisesRegex( + asyncpg.ConnectionDoesNotExistError, + 'closed in the middle of operation: ' + 'terminating connection'): + await fut + + exc = fut.exception() + self.assertIsInstance(exc.__cause__, exceptions.AdminShutdownError) + class TestExecuteMany(tb.ConnectedTestCase): def setUp(self):