From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Mon, 12 Oct 2020 14:16:25 +0000 (-0400)
Subject: Ensure all cursor self.handle_exception() calls are covered
X-Git-Tag: rel_1_4_0b1~40
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=286fe9b0b165e46099066cc16552b93a0b853a7e;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git

Ensure all cursor self.handle_exception() calls are covered

Fixes: #5642
Change-Id: I07a77483e6e2ec593d87d3d3467a4339c5f77a26
---

diff --git a/lib/sqlalchemy/engine/cursor.py b/lib/sqlalchemy/engine/cursor.py
index 6c4a756c9a..b06c622e06 100644
--- a/lib/sqlalchemy/engine/cursor.py
+++ b/lib/sqlalchemy/engine/cursor.py
@@ -1079,7 +1079,7 @@ class BufferedRowCursorFetchStrategy(CursorFetchStrategy):
                 try:
                     result._soft_close(hard=hard_close)
                 except BaseException as e:
-                    self.handle_exception(result, e)
+                    self.handle_exception(result, dbapi_cursor, e)
                 return None
         return self._rowbuffer.popleft()
 
@@ -1093,7 +1093,7 @@ class BufferedRowCursorFetchStrategy(CursorFetchStrategy):
             try:
                 buf.extend(dbapi_cursor.fetchmany(size - lb))
             except BaseException as e:
-                self.handle_exception(result, e)
+                self.handle_exception(result, dbapi_cursor, e)
 
         result = buf[0:size]
         self._rowbuffer = collections.deque(buf[size:])
diff --git a/test/sql/test_resultset.py b/test/sql/test_resultset.py
index 7d62665417..2deec3f829 100644
--- a/test/sql/test_resultset.py
+++ b/test/sql/test_resultset.py
@@ -44,6 +44,7 @@ from sqlalchemy.testing import in_
 from sqlalchemy.testing import is_
 from sqlalchemy.testing import is_true
 from sqlalchemy.testing import le_
+from sqlalchemy.testing import mock
 from sqlalchemy.testing import ne_
 from sqlalchemy.testing import not_in
 from sqlalchemy.testing.mock import Mock
@@ -2486,6 +2487,73 @@ class AlternateCursorResultTest(fixtures.TablesTest):
         result.fetchmany(5)
         eq_(len(result.cursor_strategy._rowbuffer), 296)
 
+        self._test_result_processor(
+            _cursor.BufferedRowCursorFetchStrategy, False
+        )
+
+    @testing.combinations(
+        _cursor.CursorFetchStrategy,
+        _cursor.BufferedRowCursorFetchStrategy,
+        # does not handle error in fetch
+        # _cursor.FullyBufferedCursorFetchStrategy,
+        argnames="strategy_cls",
+    )
+    @testing.combinations(
+        "fetchone",
+        "fetchmany",
+        "fetchmany_w_num",
+        "fetchall",
+        argnames="method_name",
+    )
+    def test_handle_error_in_fetch(self, strategy_cls, method_name):
+        class cursor(object):
+            def raise_(self):
+                raise IOError("random non-DBAPI error during cursor operation")
+
+            def fetchone(self):
+                self.raise_()
+
+            def fetchmany(self, num=None):
+                self.raise_()
+
+            def fetchall(self):
+                self.raise_()
+
+            def close(self):
+                self.raise_()
+
+        with self._proxy_fixture(strategy_cls):
+            with self.engine.connect() as conn:
+                r = conn.execute(select(self.table))
+                assert isinstance(r.cursor_strategy, strategy_cls)
+                with mock.patch.object(r, "cursor", cursor()):
+
+                    with testing.expect_raises_message(
+                        IOError, "random non-DBAPI"
+                    ):
+                        if method_name == "fetchmany_w_num":
+                            r.fetchmany(10)
+                        else:
+                            getattr(r, method_name)()
+                            getattr(r, method_name)()
+
+                r.close()
+
+    def test_buffered_row_close_error_during_fetchone(self):
+        def raise_(**kw):
+            raise IOError("random non-DBAPI error during cursor operation")
+
+        with self._proxy_fixture(_cursor.BufferedRowCursorFetchStrategy):
+            with self.engine.connect() as conn:
+                r = conn.execute(select(self.table).limit(1))
+
+                r.fetchone()
+                with mock.patch.object(
+                    r, "_soft_close", raise_
+                ), testing.expect_raises_message(IOError, "random non-DBAPI"):
+                    r.first()
+                r.close()
+
 
 class MergeCursorResultTest(fixtures.TablesTest):
     __backend__ = True