refactor

Author: IlyaFaer

google/cloud/spanner_dbapi/cursor.py

@@ -334,8 +334,9 @@ def fetchone(self):
         except StopIteration:
             return
         except Aborted:
-            self.connection.retry_transaction()
-            return self.fetchone()
+            if not self.connection.read_only:
+                self.connection.retry_transaction()
+                return self.fetchone()
 
     def fetchall(self):
         """Fetch all (remaining) rows of a query result, returning them as
@@ -350,8 +351,9 @@ def fetchall(self):
                     self._checksum.consume_result(row)
                 res.append(row)
         except Aborted:
-            self.connection.retry_transaction()
-            return self.fetchall()
+            if not self.connection.read_only:
+                self.connection.retry_transaction()
+                return self.fetchall()
 
         return res
 
@@ -381,8 +383,9 @@ def fetchmany(self, size=None):
             except StopIteration:
                 break
             except Aborted:
-                self.connection.retry_transaction()
-                return self.fetchmany(size)
+                if not self.connection.read_only:
+                    self.connection.retry_transaction()
+                    return self.fetchmany(size)
 
         return items
 
@@ -405,30 +408,21 @@ def _handle_DQL_with_snapshot(self, snapshot, sql, params):
         res = snapshot.execute_sql(
             sql, params=params, param_types=get_param_types(params)
         )
-        if type(res) == int:
-            self._row_count = res
-            self._itr = None
-        else:
-            # Immediately using:
-            #   iter(response)
-            # here, because this Spanner API doesn't provide
-            # easy mechanisms to detect when only a single item
-            # is returned or many, yet mixing results that
-            # are for .fetchone() with those that would result in
-            # many items returns a RuntimeError if .fetchone() is
-            # invoked and vice versa.
-            self._result_set = res
-            # Read the first element so that the StreamedResultSet can
-            # return the metadata after a DQL statement. See issue #155.
-            while True:
-                try:
-                    self._itr = PeekIterator(self._result_set)
-                    break
-                except Aborted:
-                    self.connection.retry_transaction()
-            # Unfortunately, Spanner doesn't seem to send back
-            # information about the number of rows available.
-            self._row_count = _UNSET_COUNT
+        # Immediately using:
+        #   iter(response)
+        # here, because this Spanner API doesn't provide
+        # easy mechanisms to detect when only a single item
+        # is returned or many, yet mixing results that
+        # are for .fetchone() with those that would result in
+        # many items returns a RuntimeError if .fetchone() is
+        # invoked and vice versa.
+        self._result_set = res
+        # Read the first element so that the StreamedResultSet can
+        # return the metadata after a DQL statement. See issue #155.
+        self._itr = PeekIterator(self._result_set)
+        # Unfortunately, Spanner doesn't seem to send back
+        # information about the number of rows available.
+        self._row_count = _UNSET_COUNT
 
     def _handle_DQL(self, sql, params):
         if self.connection.read_only and not self.connection.autocommit:

tests/unit/spanner_dbapi/test_cursor.py

@@ -707,14 +707,9 @@ def test_handle_dql(self):
         ) = mock.MagicMock()
         cursor = self._make_one(connection)
 
-        mock_snapshot.execute_sql.return_value = int(0)
+        mock_snapshot.execute_sql.return_value = ["0"]
         cursor._handle_DQL("sql", params=None)
-        self.assertEqual(cursor._row_count, 0)
-        self.assertIsNone(cursor._itr)
-
-        mock_snapshot.execute_sql.return_value = "0"
-        cursor._handle_DQL("sql", params=None)
-        self.assertEqual(cursor._result_set, "0")
+        self.assertEqual(cursor._result_set, ["0"])
         self.assertIsInstance(cursor._itr, utils.PeekIterator)
         self.assertEqual(cursor._row_count, _UNSET_COUNT)
 
@@ -831,37 +826,6 @@ def test_peek_iterator_aborted(self, mock_client):
 
                 retry_mock.assert_called_with()
 
-    @mock.patch("google.cloud.spanner_v1.Client")
-    def test_peek_iterator_aborted_autocommit(self, mock_client):
-        """
-        Checking that an Aborted exception is retried in case it happened while
-        streaming the first element with a PeekIterator in autocommit mode.
-        """
-        from google.api_core.exceptions import Aborted
-        from google.cloud.spanner_dbapi.connection import connect
-
-        connection = connect("test-instance", "test-database")
-
-        connection.autocommit = True
-        cursor = connection.cursor()
-        with mock.patch(
-            "google.cloud.spanner_dbapi.utils.PeekIterator.__init__",
-            side_effect=(Aborted("Aborted"), None),
-        ):
-            with mock.patch(
-                "google.cloud.spanner_dbapi.connection.Connection.retry_transaction"
-            ) as retry_mock:
-                with mock.patch(
-                    "google.cloud.spanner_dbapi.connection.Connection.run_statement",
-                    return_value=((1, 2, 3), None),
-                ):
-                    with mock.patch(
-                        "google.cloud.spanner_v1.database.Database.snapshot"
-                    ):
-                        cursor.execute("SELECT * FROM table_name")
-
-                retry_mock.assert_called_with()
-
     @mock.patch("google.cloud.spanner_v1.Client")
     def test_fetchone_retry_aborted(self, mock_client):
         """Check that aborted fetch re-executing transaction."""