fix: add support for row_count in cursor. (#675)

* fix: add support for row_count

* docs: update rowcount property doc

* fix: updated tests for cursor to check row_count

* refactor: lint fixes

* test: add test for do_batch_update

* refactor: Empty commit

google/cloud/spanner_dbapi/cursor.py

@@ -44,6 +44,8 @@
 
 from google.rpc.code_pb2 import ABORTED, OK
 
+_UNSET_COUNT = -1
+
 ColumnDetails = namedtuple("column_details", ["null_ok", "spanner_type"])
 Statement = namedtuple("Statement", "sql, params, param_types, checksum, is_insert")
 
@@ -80,6 +82,7 @@ class Cursor(object):
     def __init__(self, connection):
         self._itr = None
         self._result_set = None
+        self._row_count = _UNSET_COUNT
         self.lastrowid = None
         self.connection = connection
         self._is_closed = False
@@ -134,13 +137,14 @@ def description(self):
 
     @property
     def rowcount(self):
-        """The number of rows produced by the last `execute()` call.
+        """The number of rows updated by the last UPDATE, DELETE request's `execute()` call.
+        For SELECT requests the rowcount returns -1.
 
-        The property is non-operational and always returns -1. Request
-        resulting rows are streamed by the `fetch*()` methods and
-        can't be counted before they are all streamed.
+        :rtype: int
+        :returns: The number of rows updated by the last UPDATE, DELETE request's .execute*() call.
         """
-        return -1
+
+        return self._row_count
 
     @check_not_closed
     def callproc(self, procname, args=None):
@@ -170,7 +174,11 @@ def _do_execute_update(self, transaction, sql, params):
         result = transaction.execute_update(
             sql, params=params, param_types=get_param_types(params)
         )
-        self._itr = iter([result])
+        self._itr = None
+        if type(result) == int:
+            self._row_count = result
+
+        return result
 
     def _do_batch_update(self, transaction, statements, many_result_set):
         status, res = transaction.batch_update(statements)
@@ -181,6 +189,8 @@ def _do_batch_update(self, transaction, statements, many_result_set):
         elif status.code != OK:
             raise OperationalError(status.message)
 
+        self._row_count = sum([max(val, 0) for val in res])
+
     def _batch_DDLs(self, sql):
         """
         Check that the given operation contains only DDL
@@ -414,6 +424,9 @@ def _handle_DQL_with_snapshot(self, snapshot, sql, params):
         # Read the first element so that the StreamedResultSet can
         # return the metadata after a DQL statement.
         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):
         sql, params = parse_utils.sql_pyformat_args_to_spanner(sql, params)

tests/unit/spanner_dbapi/test_cursor.py

@@ -37,11 +37,13 @@ def _make_connection(self, *args, **kwargs):
 
         return Connection(*args, **kwargs)
 
-    def _transaction_mock(self):
+    def _transaction_mock(self, mock_response=[]):
         from google.rpc.code_pb2 import OK
 
         transaction = mock.Mock(committed=False, rolled_back=False)
-        transaction.batch_update = mock.Mock(return_value=[mock.Mock(code=OK), []])
+        transaction.batch_update = mock.Mock(
+            return_value=[mock.Mock(code=OK), mock_response]
+        )
         return transaction
 
     def test_property_connection(self):
@@ -62,10 +64,12 @@ def test_property_description(self):
         self.assertIsInstance(cursor.description[0], ColumnInfo)
 
     def test_property_rowcount(self):
+        from google.cloud.spanner_dbapi.cursor import _UNSET_COUNT
+
         connection = self._make_connection(self.INSTANCE, self.DATABASE)
         cursor = self._make_one(connection)
 
-        assert cursor.rowcount == -1
+        self.assertEqual(cursor.rowcount, _UNSET_COUNT)
 
     def test_callproc(self):
         from google.cloud.spanner_dbapi.exceptions import InterfaceError
@@ -93,25 +97,58 @@ def test_close(self, mock_client):
             cursor.execute("SELECT * FROM database")
 
     def test_do_execute_update(self):
-        from google.cloud.spanner_dbapi.checksum import ResultsChecksum
+        from google.cloud.spanner_dbapi.cursor import _UNSET_COUNT
 
         connection = self._make_connection(self.INSTANCE, self.DATABASE)
         cursor = self._make_one(connection)
-        cursor._checksum = ResultsChecksum()
         transaction = mock.MagicMock()
 
         def run_helper(ret_value):
             transaction.execute_update.return_value = ret_value
-            cursor._do_execute_update(
+            res = cursor._do_execute_update(
                 transaction=transaction, sql="SELECT * WHERE true", params={},
             )
-            return cursor.fetchall()
+            return res
 
         expected = "good"
-        self.assertEqual(run_helper(expected), [expected])
+        self.assertEqual(run_helper(expected), expected)
+        self.assertEqual(cursor._row_count, _UNSET_COUNT)
 
         expected = 1234
-        self.assertEqual(run_helper(expected), [expected])
+        self.assertEqual(run_helper(expected), expected)
+        self.assertEqual(cursor._row_count, expected)
+
+    def test_do_batch_update(self):
+        from google.cloud.spanner_dbapi import connect
+        from google.cloud.spanner_v1.param_types import INT64
+        from google.cloud.spanner_v1.types.spanner import Session
+
+        sql = "DELETE FROM table WHERE col1 = %s"
+
+        connection = connect("test-instance", "test-database")
+
+        connection.autocommit = True
+        transaction = self._transaction_mock(mock_response=[1, 1, 1])
+        cursor = connection.cursor()
+
+        with mock.patch(
+            "google.cloud.spanner_v1.services.spanner.client.SpannerClient.create_session",
+            return_value=Session(),
+        ):
+            with mock.patch(
+                "google.cloud.spanner_v1.session.Session.transaction",
+                return_value=transaction,
+            ):
+                cursor.executemany(sql, [(1,), (2,), (3,)])
+
+        transaction.batch_update.assert_called_once_with(
+            [
+                ("DELETE FROM table WHERE col1 = @a0", {"a0": 1}, {"a0": INT64}),
+                ("DELETE FROM table WHERE col1 = @a0", {"a0": 2}, {"a0": INT64}),
+                ("DELETE FROM table WHERE col1 = @a0", {"a0": 3}, {"a0": INT64}),
+            ]
+        )
+        self.assertEqual(cursor._row_count, 3)
 
     def test_execute_programming_error(self):
         from google.cloud.spanner_dbapi.exceptions import ProgrammingError
@@ -704,6 +741,7 @@ def test_setoutputsize(self):
 
     def test_handle_dql(self):
         from google.cloud.spanner_dbapi import utils
+        from google.cloud.spanner_dbapi.cursor import _UNSET_COUNT
 
         connection = self._make_connection(self.INSTANCE, mock.MagicMock())
         connection.database.snapshot.return_value.__enter__.return_value = (
@@ -715,6 +753,7 @@ def test_handle_dql(self):
         cursor._handle_DQL("sql", params=None)
         self.assertEqual(cursor._result_set, ["0"])
         self.assertIsInstance(cursor._itr, utils.PeekIterator)
+        self.assertEqual(cursor._row_count, _UNSET_COUNT)
 
     def test_context(self):
         connection = self._make_connection(self.INSTANCE, self.DATABASE)