]> 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 07:26:58 +0000 (09:26 +0200)
psycopg/psycopg/cursor.py
psycopg/psycopg/cursor_async.py

index fec09a032762dc615c494df22b0edd9be2a70721..c64eab2bd4441006d26a65e2a02b5f8f45559ddd 100644 (file)
@@ -575,14 +575,17 @@ class Cursor(BaseCursor["Connection[Any]", Row]):
         """
         Iterate row-by-row on a result from the database.
         """
-        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 830a20b0d4e6f51517df57b932707cdfbc37802d..0b665b43534b34a708b294fe2067dac6e26b6771 100644 (file)
@@ -93,14 +93,19 @@ class AsyncCursor(BaseCursor["AsyncConnection[Any]", Row]):
         *,
         binary: Optional[bool] = None,
     ) -> AsyncIterator[Row]:
-        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]:
         self._check_result_for_fetch()