From a86e2a2a033823dd95bb4a8511f86a9b39425b3d Mon Sep 17 00:00:00 2001 From: Vikash Singh <3116482+vi3k6i5@users.noreply.github.com> Date: Wed, 2 Feb 2022 16:17:14 +0530 Subject: [PATCH 1/6] fix: add support for row_count --- google/cloud/spanner_dbapi/cursor.py | 29 +++++++++++++++++++--------- 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/google/cloud/spanner_dbapi/cursor.py b/google/cloud/spanner_dbapi/cursor.py index 84b35292f0..91bbd42e67 100644 --- a/google/cloud/spanner_dbapi/cursor.py +++ b/google/cloud/spanner_dbapi/cursor.py @@ -14,6 +14,7 @@ """Database cursor for Google Cloud Spanner DB API.""" +import warnings from collections import namedtuple import sqlparse @@ -44,6 +45,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 +83,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 +138,12 @@ def description(self): @property def rowcount(self): - """The number of rows produced by the last `execute()` call. + """The number of rows produced by the last `.execute()`. - 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 produced by the last .execute*(). """ - return -1 + return self._row_count @check_not_closed def callproc(self, procname, args=None): @@ -170,7 +173,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) @@ -250,9 +257,10 @@ def execute(self, sql, args=None): class_ == parse_utils.STMT_INSERT, ) - (self._result_set, self._checksum,) = self.connection.run_statement( - statement - ) + ( + self._result_set, + self._checksum, + ) = self.connection.run_statement(statement) while True: try: self._itr = PeekIterator(self._result_set) @@ -414,6 +422,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) From a708206c1a8a47757af87b3fbb8d56e81b00d7df Mon Sep 17 00:00:00 2001 From: Vikash Singh <3116482+vi3k6i5@users.noreply.github.com> Date: Wed, 2 Feb 2022 16:37:57 +0530 Subject: [PATCH 2/6] docs: update rowcount property doc --- google/cloud/spanner_dbapi/cursor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/spanner_dbapi/cursor.py b/google/cloud/spanner_dbapi/cursor.py index 91bbd42e67..0434826b1a 100644 --- a/google/cloud/spanner_dbapi/cursor.py +++ b/google/cloud/spanner_dbapi/cursor.py @@ -138,7 +138,7 @@ def description(self): @property def rowcount(self): - """The number of rows produced by the last `.execute()`. + """The number of rows produced by the last `.execute()` call. :rtype: int :returns: The number of rows produced by the last .execute*(). From f1ca77a9312e64d76df51c86430fa391434a49ee Mon Sep 17 00:00:00 2001 From: Vikash Singh <3116482+vi3k6i5@users.noreply.github.com> Date: Thu, 3 Feb 2022 11:37:55 +0530 Subject: [PATCH 3/6] fix: updated tests for cursor to check row_count --- google/cloud/spanner_dbapi/cursor.py | 7 +++---- tests/unit/spanner_dbapi/test_cursor.py | 19 ++++++++++++------- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/google/cloud/spanner_dbapi/cursor.py b/google/cloud/spanner_dbapi/cursor.py index 0434826b1a..332dc14293 100644 --- a/google/cloud/spanner_dbapi/cursor.py +++ b/google/cloud/spanner_dbapi/cursor.py @@ -257,10 +257,9 @@ def execute(self, sql, args=None): class_ == parse_utils.STMT_INSERT, ) - ( - self._result_set, - self._checksum, - ) = self.connection.run_statement(statement) + (self._result_set, self._checksum,) = self.connection.run_statement( + statement + ) while True: try: self._itr = PeekIterator(self._result_set) diff --git a/tests/unit/spanner_dbapi/test_cursor.py b/tests/unit/spanner_dbapi/test_cursor.py index f7607b79bd..fab57e8b1b 100644 --- a/tests/unit/spanner_dbapi/test_cursor.py +++ b/tests/unit/spanner_dbapi/test_cursor.py @@ -62,10 +62,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 +95,26 @@ 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_execute_programming_error(self): from google.cloud.spanner_dbapi.exceptions import ProgrammingError @@ -704,6 +707,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 +719,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) From 8766fae7728c5ed4acb767975bfe55c96ab43e9f Mon Sep 17 00:00:00 2001 From: Vikash Singh <3116482+vi3k6i5@users.noreply.github.com> Date: Thu, 3 Feb 2022 11:46:10 +0530 Subject: [PATCH 4/6] refactor: lint fixes --- google/cloud/spanner_dbapi/cursor.py | 1 - 1 file changed, 1 deletion(-) diff --git a/google/cloud/spanner_dbapi/cursor.py b/google/cloud/spanner_dbapi/cursor.py index 332dc14293..b12228e52f 100644 --- a/google/cloud/spanner_dbapi/cursor.py +++ b/google/cloud/spanner_dbapi/cursor.py @@ -14,7 +14,6 @@ """Database cursor for Google Cloud Spanner DB API.""" -import warnings from collections import namedtuple import sqlparse From 747325ee9ccd8c83a7735ad7e7aeb009d9b1854c Mon Sep 17 00:00:00 2001 From: Vikash Singh <3116482+vi3k6i5@users.noreply.github.com> Date: Fri, 4 Feb 2022 12:57:52 +0530 Subject: [PATCH 5/6] test: add test for do_batch_update --- google/cloud/spanner_dbapi/cursor.py | 8 ++++-- tests/unit/spanner_dbapi/test_cursor.py | 38 +++++++++++++++++++++++-- 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/google/cloud/spanner_dbapi/cursor.py b/google/cloud/spanner_dbapi/cursor.py index b12228e52f..7c8c5bdbc5 100644 --- a/google/cloud/spanner_dbapi/cursor.py +++ b/google/cloud/spanner_dbapi/cursor.py @@ -137,11 +137,13 @@ 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. :rtype: int - :returns: The number of rows produced by the last .execute*(). + :returns: The number of rows updated by the last UPDATE, DELETE request's .execute*() call. """ + return self._row_count @check_not_closed @@ -187,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 diff --git a/tests/unit/spanner_dbapi/test_cursor.py b/tests/unit/spanner_dbapi/test_cursor.py index fab57e8b1b..51732bc1b0 100644 --- a/tests/unit/spanner_dbapi/test_cursor.py +++ b/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): @@ -116,6 +118,38 @@ def run_helper(ret_value): 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 From 1a9634aaff5536fc4dd8dcda9e81ce53eae75cc3 Mon Sep 17 00:00:00 2001 From: Vikash Singh <3116482+vi3k6i5@users.noreply.github.com> Date: Fri, 4 Feb 2022 14:48:43 +0530 Subject: [PATCH 6/6] refactor: Empty commit