Skip to content

Commit 4afea77

Browse files
HemangChothaniHemangChothani
authored
fix: connection attribute of connection class and include related unit tests (#228)
1 parent 1f80a39 commit 4afea77

File tree

2 files changed

+253
-2
lines changed

2 files changed

+253
-2
lines changed

‎google/cloud/spanner_dbapi/cursor.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -279,7 +279,7 @@ def fetchall(self):
279279
self._checksum.consume_result(row)
280280
res.append(row)
281281
except Aborted:
282-
self._connection.retry_transaction()
282+
self.connection.retry_transaction()
283283
return self.fetchall()
284284

285285
return res
@@ -310,7 +310,7 @@ def fetchmany(self, size=None):
310310
except StopIteration:
311311
break
312312
except Aborted:
313-
self._connection.retry_transaction()
313+
self.connection.retry_transaction()
314314
return self.fetchmany(size)
315315

316316
return items

‎tests/unit/spanner_dbapi/test_cursor.py‎

Lines changed: 251 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -315,6 +315,22 @@ def test_fetchone(self):
315315
self.assertEqual(cursor.fetchone(), lst[i])
316316
self.assertIsNone(cursor.fetchone())
317317

318+
@unittest.skipIf(
319+
sys.version_info[0] < 3, "Python 2 has an outdated iterator definition"
320+
)
321+
def test_fetchone_w_autocommit(self):
322+
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
323+
324+
connection = self._make_connection(self.INSTANCE, mock.MagicMock())
325+
connection.autocommit = True
326+
cursor = self._make_one(connection)
327+
cursor._checksum = ResultsChecksum()
328+
lst = [1, 2, 3]
329+
cursor._itr = iter(lst)
330+
for i in range(len(lst)):
331+
self.assertEqual(cursor.fetchone(), lst[i])
332+
self.assertIsNone(cursor.fetchone())
333+
318334
def test_fetchmany(self):
319335
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
320336

@@ -329,6 +345,21 @@ def test_fetchmany(self):
329345
result = cursor.fetchmany(len(lst))
330346
self.assertEqual(result, lst[1:])
331347

348+
def test_fetchmany_w_autocommit(self):
349+
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
350+
351+
connection = self._make_connection(self.INSTANCE, mock.MagicMock())
352+
connection.autocommit = True
353+
cursor = self._make_one(connection)
354+
cursor._checksum = ResultsChecksum()
355+
lst = [(1,), (2,), (3,)]
356+
cursor._itr = iter(lst)
357+
358+
self.assertEqual(cursor.fetchmany(), [lst[0]])
359+
360+
result = cursor.fetchmany(len(lst))
361+
self.assertEqual(result, lst[1:])
362+
332363
def test_fetchall(self):
333364
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
334365

@@ -339,6 +370,17 @@ def test_fetchall(self):
339370
cursor._itr = iter(lst)
340371
self.assertEqual(cursor.fetchall(), lst)
341372

373+
def test_fetchall_w_autocommit(self):
374+
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
375+
376+
connection = self._make_connection(self.INSTANCE, mock.MagicMock())
377+
connection.autocommit = True
378+
cursor = self._make_one(connection)
379+
cursor._checksum = ResultsChecksum()
380+
lst = [(1,), (2,), (3,)]
381+
cursor._itr = iter(lst)
382+
self.assertEqual(cursor.fetchall(), lst)
383+
342384
def test_nextset(self):
343385
from google.cloud.spanner_dbapi import exceptions
344386

@@ -586,3 +628,212 @@ def test_fetchone_retry_aborted_statements_checksums_mismatch(self):
586628
cursor.fetchone()
587629

588630
run_mock.assert_called_with(statement, retried=True)
631+
632+
def test_fetchall_retry_aborted(self):
633+
"""Check that aborted fetch re-executing transaction."""
634+
from google.api_core.exceptions import Aborted
635+
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
636+
from google.cloud.spanner_dbapi.connection import connect
637+
638+
with mock.patch(
639+
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True,
640+
):
641+
with mock.patch(
642+
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
643+
):
644+
connection = connect("test-instance", "test-database")
645+
646+
cursor = connection.cursor()
647+
cursor._checksum = ResultsChecksum()
648+
649+
with mock.patch(
650+
"google.cloud.spanner_dbapi.cursor.Cursor.__iter__",
651+
side_effect=(Aborted("Aborted"), iter([])),
652+
):
653+
with mock.patch(
654+
"google.cloud.spanner_dbapi.connection.Connection.retry_transaction"
655+
) as retry_mock:
656+
657+
cursor.fetchall()
658+
659+
retry_mock.assert_called_with()
660+
661+
def test_fetchall_retry_aborted_statements(self):
662+
"""Check that retried transaction executing the same statements."""
663+
from google.api_core.exceptions import Aborted
664+
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
665+
from google.cloud.spanner_dbapi.connection import connect
666+
from google.cloud.spanner_dbapi.cursor import Statement
667+
668+
row = ["field1", "field2"]
669+
with mock.patch(
670+
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True,
671+
):
672+
with mock.patch(
673+
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
674+
):
675+
connection = connect("test-instance", "test-database")
676+
677+
cursor = connection.cursor()
678+
cursor._checksum = ResultsChecksum()
679+
cursor._checksum.consume_result(row)
680+
681+
statement = Statement("SELECT 1", [], {}, cursor._checksum, False)
682+
connection._statements.append(statement)
683+
684+
with mock.patch(
685+
"google.cloud.spanner_dbapi.cursor.Cursor.__iter__",
686+
side_effect=(Aborted("Aborted"), iter(row)),
687+
):
688+
with mock.patch(
689+
"google.cloud.spanner_dbapi.connection.Connection.run_statement",
690+
return_value=([row], ResultsChecksum()),
691+
) as run_mock:
692+
cursor.fetchall()
693+
694+
run_mock.assert_called_with(statement, retried=True)
695+
696+
def test_fetchall_retry_aborted_statements_checksums_mismatch(self):
697+
"""Check transaction retrying with underlying data being changed."""
698+
from google.api_core.exceptions import Aborted
699+
from google.cloud.spanner_dbapi.exceptions import RetryAborted
700+
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
701+
from google.cloud.spanner_dbapi.connection import connect
702+
from google.cloud.spanner_dbapi.cursor import Statement
703+
704+
row = ["field1", "field2"]
705+
row2 = ["updated_field1", "field2"]
706+
707+
with mock.patch(
708+
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True,
709+
):
710+
with mock.patch(
711+
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
712+
):
713+
connection = connect("test-instance", "test-database")
714+
715+
cursor = connection.cursor()
716+
cursor._checksum = ResultsChecksum()
717+
cursor._checksum.consume_result(row)
718+
719+
statement = Statement("SELECT 1", [], {}, cursor._checksum, False)
720+
connection._statements.append(statement)
721+
722+
with mock.patch(
723+
"google.cloud.spanner_dbapi.cursor.Cursor.__iter__",
724+
side_effect=(Aborted("Aborted"), iter(row)),
725+
):
726+
with mock.patch(
727+
"google.cloud.spanner_dbapi.connection.Connection.run_statement",
728+
return_value=([row2], ResultsChecksum()),
729+
) as run_mock:
730+
731+
with self.assertRaises(RetryAborted):
732+
cursor.fetchall()
733+
734+
run_mock.assert_called_with(statement, retried=True)
735+
736+
def test_fetchmany_retry_aborted(self):
737+
"""Check that aborted fetch re-executing transaction."""
738+
from google.api_core.exceptions import Aborted
739+
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
740+
from google.cloud.spanner_dbapi.connection import connect
741+
742+
with mock.patch(
743+
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True,
744+
):
745+
with mock.patch(
746+
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
747+
):
748+
connection = connect("test-instance", "test-database")
749+
750+
cursor = connection.cursor()
751+
cursor._checksum = ResultsChecksum()
752+
753+
with mock.patch(
754+
"google.cloud.spanner_dbapi.cursor.Cursor.__next__",
755+
side_effect=(Aborted("Aborted"), None),
756+
):
757+
with mock.patch(
758+
"google.cloud.spanner_dbapi.connection.Connection.retry_transaction"
759+
) as retry_mock:
760+
761+
cursor.fetchmany()
762+
763+
retry_mock.assert_called_with()
764+
765+
def test_fetchmany_retry_aborted_statements(self):
766+
"""Check that retried transaction executing the same statements."""
767+
from google.api_core.exceptions import Aborted
768+
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
769+
from google.cloud.spanner_dbapi.connection import connect
770+
from google.cloud.spanner_dbapi.cursor import Statement
771+
772+
row = ["field1", "field2"]
773+
with mock.patch(
774+
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True,
775+
):
776+
with mock.patch(
777+
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
778+
):
779+
connection = connect("test-instance", "test-database")
780+
781+
cursor = connection.cursor()
782+
cursor._checksum = ResultsChecksum()
783+
cursor._checksum.consume_result(row)
784+
785+
statement = Statement("SELECT 1", [], {}, cursor._checksum, False)
786+
connection._statements.append(statement)
787+
788+
with mock.patch(
789+
"google.cloud.spanner_dbapi.cursor.Cursor.__next__",
790+
side_effect=(Aborted("Aborted"), None),
791+
):
792+
with mock.patch(
793+
"google.cloud.spanner_dbapi.connection.Connection.run_statement",
794+
return_value=([row], ResultsChecksum()),
795+
) as run_mock:
796+
797+
cursor.fetchmany(len(row))
798+
799+
run_mock.assert_called_with(statement, retried=True)
800+
801+
def test_fetchmany_retry_aborted_statements_checksums_mismatch(self):
802+
"""Check transaction retrying with underlying data being changed."""
803+
from google.api_core.exceptions import Aborted
804+
from google.cloud.spanner_dbapi.exceptions import RetryAborted
805+
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
806+
from google.cloud.spanner_dbapi.connection import connect
807+
from google.cloud.spanner_dbapi.cursor import Statement
808+
809+
row = ["field1", "field2"]
810+
row2 = ["updated_field1", "field2"]
811+
812+
with mock.patch(
813+
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True,
814+
):
815+
with mock.patch(
816+
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
817+
):
818+
connection = connect("test-instance", "test-database")
819+
820+
cursor = connection.cursor()
821+
cursor._checksum = ResultsChecksum()
822+
cursor._checksum.consume_result(row)
823+
824+
statement = Statement("SELECT 1", [], {}, cursor._checksum, False)
825+
connection._statements.append(statement)
826+
827+
with mock.patch(
828+
"google.cloud.spanner_dbapi.cursor.Cursor.__next__",
829+
side_effect=(Aborted("Aborted"), None),
830+
):
831+
with mock.patch(
832+
"google.cloud.spanner_dbapi.connection.Connection.run_statement",
833+
return_value=([row2], ResultsChecksum()),
834+
) as run_mock:
835+
836+
with self.assertRaises(RetryAborted):
837+
cursor.fetchmany(len(row))
838+
839+
run_mock.assert_called_with(statement, retried=True)

0 commit comments

Comments
 (0)