]> git.ipfire.org Git - thirdparty/psycopg.git/commitdiff
fix: shorter tracebacks on stream() error
authorDaniele Varrazzo <daniele.varrazzo@gmail.com>
Tue, 7 Jun 2022 04:46:56 +0000 (06:46 +0200)
committerDaniele Varrazzo <daniele.varrazzo@gmail.com>
Tue, 7 Jun 2022 08:46:27 +0000 (10:46 +0200)
psycopg/psycopg/cursor.py
psycopg/psycopg/cursor_async.py

index a2777e470026d4c5da9ebbf5e204545fd08a749b..be391bb2778bfe8adca68a46ac8c025ebf6ef94e 100644 (file)
@@ -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]:
         """
index 539678b94918dd41d08006bf8a22f9c756c3c20b..4598bf434fa7aabccdab825580da1b4dee7ed7d5 100644 (file)
@@ -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()