def test_retry_transaction_w_multiple_statement(self):
"""Check retrying an aborted transaction."""
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.cursor import Statement
row = ["field1", "field2"]
connection = self._make_connection()
checksum = ResultsChecksum()
checksum.consume_result(row)
retried_checkum = ResultsChecksum()
statement = Statement("SELECT 1", [], {}, checksum, False)
statement1 = Statement("SELECT 2", [], {}, checksum, False)
connection._statements.append(statement)
connection._statements.append(statement1)
run_mock = connection.run_statement = mock.Mock()
run_mock.return_value = ([row], retried_checkum)
with mock.patch(
"google.cloud.spanner_dbapi.connection._compare_checksums"
) as compare_mock:
connection.retry_transaction()
compare_mock.assert_called_with(checksum, retried_checkum)
run_mock.assert_called_with(statement1, retried=True)
def test_fetchmany_retry_aborted_statements_checksums_mismatch(
self, mock_client):
"""Check transaction retrying with underlying data being changed."""
from google.api_core.exceptions import Aborted
from google.cloud.spanner_dbapi.exceptions import RetryAborted
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.connection import connect
from google.cloud.spanner_dbapi.cursor import Statement
row = ["field1", "field2"]
row2 = ["updated_field1", "field2"]
connection = connect("test-instance", "test-database")
cursor = connection.cursor()
cursor._checksum = ResultsChecksum()
cursor._checksum.consume_result(row)
statement = Statement("SELECT 1", [], {}, cursor._checksum, False)
connection._statements.append(statement)
with mock.patch(
"google.cloud.spanner_dbapi.cursor.Cursor.__next__",
side_effect=(Aborted("Aborted"), None),
):
with mock.patch(
"google.cloud.spanner_dbapi.connection.Connection.run_statement",
return_value=([row2], ResultsChecksum()),
) as run_mock:
with self.assertRaises(RetryAborted):
cursor.fetchmany(len(row))
run_mock.assert_called_with(statement, retried=True)
def test_fetchall_retry_aborted_statements(self, mock_client):
"""Check that retried transaction executing the same statements."""
from google.api_core.exceptions import Aborted
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.connection import connect
from google.cloud.spanner_dbapi.cursor import Statement
row = ["field1", "field2"]
connection = connect("test-instance", "test-database")
cursor = connection.cursor()
cursor._checksum = ResultsChecksum()
cursor._checksum.consume_result(row)
statement = Statement("SELECT 1", [], {}, cursor._checksum, False)
connection._statements.append(statement)
with mock.patch(
"google.cloud.spanner_dbapi.cursor.Cursor.__iter__",
side_effect=(Aborted("Aborted"), iter(row)),
):
with mock.patch(
"google.cloud.spanner_dbapi.connection.Connection.run_statement",
return_value=([row], ResultsChecksum()),
) as run_mock:
cursor.fetchall()
run_mock.assert_called_with(statement, retried=True)
def test_commit_retry_aborted_statements(self):
"""Check that retried transaction executing the same statements."""
from google.api_core.exceptions import Aborted
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.connection import connect
from google.cloud.spanner_dbapi.cursor import Statement
row = ["field1", "field2"]
with mock.patch(
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True,
):
with mock.patch(
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
):
connection = connect("test-instance", "test-database")
cursor = connection.cursor()
cursor._checksum = ResultsChecksum()
cursor._checksum.consume_result(row)
statement = Statement("SELECT 1", [], {}, cursor._checksum, False)
connection._statements.append(statement)
connection._transaction = mock.Mock(rolled_back=False, committed=False)
with mock.patch.object(
connection._transaction, "commit", side_effect=(Aborted("Aborted"), None),
):
with mock.patch(
"google.cloud.spanner_dbapi.connection.Connection.run_statement",
return_value=([row], ResultsChecksum()),
) as run_mock:
connection.commit()
run_mock.assert_called_with(statement, retried=True)
def test_retry_transaction_checksum_mismatch(self):
"""
Check retrying an aborted transaction
with results checksums mismatch.
"""
from google.cloud.spanner_dbapi.exceptions import RetryAborted
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.cursor import Statement
row = ["field1", "field2"]
retried_row = ["field3", "field4"]
connection = self._make_connection()
checksum = ResultsChecksum()
checksum.consume_result(row)
retried_checkum = ResultsChecksum()
statement = Statement("SELECT 1", [], {}, checksum, False)
connection._statements.append(statement)
with mock.patch(
"google.cloud.spanner_dbapi.connection.Connection.run_statement",
return_value=([retried_row], retried_checkum),
):
with self.assertRaises(RetryAborted):
connection.retry_transaction()
def test_commit_retry_aborted_statements(self, mock_client):
"""Check that retried transaction executing the same statements."""
from google.api_core.exceptions import Aborted
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.connection import connect
from google.cloud.spanner_dbapi.cursor import Statement
row = ["field1", "field2"]
connection = connect("test-instance", "test-database")
cursor = connection.cursor()
cursor._checksum = ResultsChecksum()
cursor._checksum.consume_result(row)
statement = Statement("SELECT 1", [], {}, cursor._checksum, False)
connection._statements.append(statement)
mock_transaction = mock.Mock(rolled_back=False, committed=False)
connection._transaction = mock_transaction
mock_transaction.commit.side_effect = [Aborted("Aborted"), None]
run_mock = connection.run_statement = mock.Mock()
run_mock.return_value = ([row], ResultsChecksum())
connection.commit()
run_mock.assert_called_with(statement, retried=True)
def test_retry_aborted_retry_without_delay(self, mock_client):
"""
Check that in case of a retried transaction failed,
the connection will retry it once again.
"""
from google.api_core.exceptions import Aborted
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.connection import connect
from google.cloud.spanner_dbapi.cursor import Statement
row = ["field1", "field2"]
connection = connect("test-instance", "test-database")
cursor = connection.cursor()
cursor._checksum = ResultsChecksum()
cursor._checksum.consume_result(row)
statement = Statement("SELECT 1", [], {}, cursor._checksum, False)
connection._statements.append(statement)
metadata_mock = mock.Mock()
metadata_mock.trailing_metadata.return_value = {}
run_mock = connection.run_statement = mock.Mock()
run_mock.side_effect = [
Aborted("Aborted", errors=[metadata_mock]),
([row], ResultsChecksum()),
]
connection._get_retry_delay = mock.Mock(return_value=False)
connection.retry_transaction()
run_mock.assert_has_calls(
(mock.call(statement, retried=True), mock.call(statement, retried=True),)
)
def test_retry_transaction_w_empty_response(self):
"""Check retrying an aborted transaction."""
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.cursor import Statement
row = []
connection = self._make_connection()
checksum = ResultsChecksum()
checksum.count = 1
retried_checkum = ResultsChecksum()
statement = Statement("SELECT 1", [], {}, checksum, False)
connection._statements.append(statement)
with mock.patch(
"google.cloud.spanner_dbapi.connection.Connection.run_statement",
return_value=(row, retried_checkum),
) as run_mock:
with mock.patch(
"google.cloud.spanner_dbapi.connection._compare_checksums"
) as compare_mock:
connection.retry_transaction()
compare_mock.assert_called_with(checksum, retried_checkum)
run_mock.assert_called_with(statement, retried=True)
def test_run_statement_w_retried(self):
"""Check that Connection doesn't remember re-executed statements."""
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.cursor import Statement
sql = """SELECT 23 FROM table WHERE id = @a1"""
params = {"a1": "value"}
param_types = {"a1": str}
connection = self._make_connection()
connection.transaction_checkout = mock.Mock()
statement = Statement(sql, params, param_types, ResultsChecksum(), False)
connection.run_statement(statement, retried=True)
self.assertEqual(len(connection._statements), 0)
def test_run_statement_w_homogeneous_insert_statements(self):
"""Check that Connection executed homogeneous insert statements."""
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.cursor import Statement
sql = "INSERT INTO T (f1, f2) VALUES (%s, %s), (%s, %s)"
params = ["a", "b", "c", "d"]
param_types = {"f1": str, "f2": str}
connection = self._make_connection()
connection.transaction_checkout = mock.Mock()
statement = Statement(sql, params, param_types, ResultsChecksum(), True)
connection.run_statement(statement, retried=True)
self.assertEqual(len(connection._statements), 0)
def test_retry_aborted_retry(self):
"""
Check that in case of a retried transaction failed,
the connection will retry it once again.
"""
from google.api_core.exceptions import Aborted
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.connection import connect
from google.cloud.spanner_dbapi.cursor import Statement
row = ["field1", "field2"]
with mock.patch(
"google.cloud.spanner_v1.instance.Instance.exists", return_value=True,
):
with mock.patch(
"google.cloud.spanner_v1.database.Database.exists", return_value=True,
):
connection = connect("test-instance", "test-database")
cursor = connection.cursor()
cursor._checksum = ResultsChecksum()
cursor._checksum.consume_result(row)
statement = Statement("SELECT 1", [], {}, cursor._checksum, False)
connection._statements.append(statement)
metadata_mock = mock.Mock()
metadata_mock.trailing_metadata.return_value = {}
with mock.patch.object(
connection,
"run_statement",
side_effect=(
Aborted("Aborted", errors=[metadata_mock]),
([row], ResultsChecksum()),
),
) as retry_mock:
connection.retry_transaction()
retry_mock.assert_has_calls(
(
mock.call(statement, retried=True),
mock.call(statement, retried=True),
)
)
def test_run_statement_wo_retried(self):
"""Check that Connection remembers executed statements."""
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.cursor import Statement
sql = """SELECT 23 FROM table WHERE id = @a1"""
params = {"a1": "value"}
param_types = {"a1": str}
connection = self._make_connection()
connection.transaction_checkout = mock.Mock()
statement = Statement(sql, params, param_types, ResultsChecksum(), False)
connection.run_statement(statement)
self.assertEqual(connection._statements[0].sql, sql)
self.assertEqual(connection._statements[0].params, params)
self.assertEqual(connection._statements[0].param_types, param_types)
self.assertIsInstance(connection._statements[0].checksum, ResultsChecksum)
def test_run_statement_w_heterogenous_insert_statements(self):
"""Check that Connection executed heterogenous insert statements."""
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.cursor import Statement
sql = "INSERT INTO T (f1, f2) VALUES (1, 2)"
params = None
param_types = None
connection = self._make_connection()
statement = Statement(sql, params, param_types, ResultsChecksum(),
True)
with mock.patch(
"google.cloud.spanner_dbapi.connection.Connection.transaction_checkout"
):
connection.run_statement(statement, retried=True)
self.assertEqual(len(connection._statements), 0)
def test_retry_transaction_raise_max_internal_retries(self):
"""Check retrying raise an error of max internal retries."""
from google.cloud.spanner_dbapi import connection as conn
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.cursor import Statement
conn.MAX_INTERNAL_RETRIES = 0
row = ["field1", "field2"]
connection = self._make_connection()
checksum = ResultsChecksum()
checksum.consume_result(row)
statement = Statement("SELECT 1", [], {}, checksum, False)
connection._statements.append(statement)
with self.assertRaises(Exception):
connection.retry_transaction()
conn.MAX_INTERNAL_RETRIES = 50
def test_run_statement_w_heterogenous_insert_statements(self):
"""Check that Connection executed heterogenous insert statements."""
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.cursor import Statement
from google.rpc.status_pb2 import Status
from google.rpc.code_pb2 import OK
sql = "INSERT INTO T (f1, f2) VALUES (1, 2)"
params = None
param_types = None
connection = self._make_connection()
transaction = mock.MagicMock()
connection.transaction_checkout = mock.Mock(return_value=transaction)
transaction.batch_update = mock.Mock(return_value=(Status(code=OK), 1))
statement = Statement(sql, params, param_types, ResultsChecksum(), True)
connection.run_statement(statement, retried=True)
self.assertEqual(len(connection._statements), 0)
