From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Tue, 7 Jun 2022 04:46:56 +0000 (+0200)
Subject: fix: shorter tracebacks on stream() error
X-Git-Tag: 3.1~66^2~2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=8307a47cd6fe07e88d9108feae28a555de1c2e82;p=thirdparty%2Fpsycopg.git

fix: shorter tracebacks on stream() error
---

diff --git a/psycopg/psycopg/cursor.py b/psycopg/psycopg/cursor.py
index a2777e470..be391bb27 100644
--- a/psycopg/psycopg/cursor.py
+++ b/psycopg/psycopg/cursor.py
@@ -776,14 +776,17 @@ class Cursor(BaseCursor["Connection[Any]", Row]):
         if self._pgconn.pipeline_status:
             raise e.ProgrammingError("stream() cannot be used in pipeline mode")
 
-        with self._conn.lock:
-            self._conn.wait(self._stream_send_gen(query, params, binary=binary))
-            first = True
-            while self._conn.wait(self._stream_fetchone_gen(first)):
-                # We know that, if we got a result, it has a single row.
-                rec: Row = self._tx.load_row(0, self._make_row)  # type: ignore
-                yield rec
-                first = False
+        try:
+            with self._conn.lock:
+                self._conn.wait(self._stream_send_gen(query, params, binary=binary))
+                first = True
+                while self._conn.wait(self._stream_fetchone_gen(first)):
+                    # We know that, if we got a result, it has a single row.
+                    rec: Row = self._tx.load_row(0, self._make_row)  # type: ignore
+                    yield rec
+                    first = False
+        except e.Error as ex:
+            raise ex.with_traceback(None)
 
     def fetchone(self) -> Optional[Row]:
         """
diff --git a/psycopg/psycopg/cursor_async.py b/psycopg/psycopg/cursor_async.py
index 539678b94..4598bf434 100644
--- a/psycopg/psycopg/cursor_async.py
+++ b/psycopg/psycopg/cursor_async.py
@@ -132,14 +132,19 @@ class AsyncCursor(BaseCursor["AsyncConnection[Any]", Row]):
         if self._pgconn.pipeline_status:
             raise e.ProgrammingError("stream() cannot be used in pipeline mode")
 
-        async with self._conn.lock:
-            await self._conn.wait(self._stream_send_gen(query, params, binary=binary))
-            first = True
-            while await self._conn.wait(self._stream_fetchone_gen(first)):
-                # We know that, if we got a result, it has a single row.
-                rec: Row = self._tx.load_row(0, self._make_row)  # type: ignore
-                yield rec
-                first = False
+        try:
+            async with self._conn.lock:
+                await self._conn.wait(
+                    self._stream_send_gen(query, params, binary=binary)
+                )
+                first = True
+                while await self._conn.wait(self._stream_fetchone_gen(first)):
+                    # We know that, if we got a result, it has a single row.
+                    rec: Row = self._tx.load_row(0, self._make_row)  # type: ignore
+                    yield rec
+                    first = False
+        except e.Error as ex:
+            raise ex.with_traceback(None)
 
     async def fetchone(self) -> Optional[Row]:
         await self._fetch_pipeline()