From fc1f666aa3565a42a05dd22fe69f98ffdd09d8c4 Mon Sep 17 00:00:00 2001 From: Daniele Varrazzo Date: Fri, 20 May 2022 10:55:17 +0200 Subject: [PATCH] fix: shorter traceback even on COPY statement error The previous commit only addressed errors raised during transferring data, not the initial COPY statement. --- psycopg/psycopg/cursor.py | 6 +++--- psycopg/psycopg/cursor_async.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/psycopg/psycopg/cursor.py b/psycopg/psycopg/cursor.py index f6a626b04..a2777e470 100644 --- a/psycopg/psycopg/cursor.py +++ b/psycopg/psycopg/cursor.py @@ -870,10 +870,10 @@ class Cursor(BaseCursor["Connection[Any]", Row]): :rtype: Copy """ - with self._conn.lock: - self._conn.wait(self._start_copy_gen(statement, params)) - try: + with self._conn.lock: + self._conn.wait(self._start_copy_gen(statement, params)) + with Copy(self) as copy: yield copy except e.Error as ex: diff --git a/psycopg/psycopg/cursor_async.py b/psycopg/psycopg/cursor_async.py index 316ce21e2..539678b94 100644 --- a/psycopg/psycopg/cursor_async.py +++ b/psycopg/psycopg/cursor_async.py @@ -196,10 +196,10 @@ class AsyncCursor(BaseCursor["AsyncConnection[Any]", Row]): """ :rtype: AsyncCopy """ - async with self._conn.lock: - await self._conn.wait(self._start_copy_gen(statement, params)) - try: + async with self._conn.lock: + await self._conn.wait(self._start_copy_gen(statement, params)) + async with AsyncCopy(self) as copy: yield copy except e.Error as ex: -- 2.47.2