def test_retry_transaction(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)
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_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_equal(self):
from google.cloud.spanner_dbapi.checksum import _compare_checksums
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
original = ResultsChecksum()
original.consume_result(5)
retried = ResultsChecksum()
retried.consume_result(5)
self.assertIsNone(_compare_checksums(original, retried))
def test_less_results(self):
from google.cloud.spanner_dbapi.checksum import _compare_checksums
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.exceptions import RetryAborted
original = ResultsChecksum()
original.consume_result(5)
retried = ResultsChecksum()
with self.assertRaises(RetryAborted):
_compare_checksums(original, retried)
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 _rerun_previous_statements(self):
"""
Helper to run all the remembered statements
from the last transaction.
"""
for statement in self._statements:
if isinstance(statement, list):
statements, checksum = statement
transaction = self.transaction_checkout()
status, res = transaction.batch_update(statements)
if status.code == ABORTED:
self.connection._transaction = None
raise Aborted(status.details)
retried_checksum = ResultsChecksum()
retried_checksum.consume_result(res)
retried_checksum.consume_result(status.code)
_compare_checksums(checksum, retried_checksum)
else:
res_iter, retried_checksum = self.run_statement(statement,
retried=True)
# executing all the completed statements
if statement != self._statements[-1]:
for res in res_iter:
retried_checksum.consume_result(res)
_compare_checksums(statement.checksum, retried_checksum)
# executing the failed statement
else:
# streaming up to the failed result or
# to the end of the streaming iterator
while len(retried_checksum) < len(statement.checksum):
try:
res = next(iter(res_iter))
retried_checksum.consume_result(res)
except StopIteration:
break
_compare_checksums(statement.checksum, retried_checksum)
def executemany(self, operation, seq_of_params):
"""Execute the given SQL with every parameters set
from the given sequence of parameters.
:type operation: str
:param operation: SQL code to execute.
:type seq_of_params: list
:param seq_of_params: Sequence of additional parameters to run
the query with.
"""
self._raise_if_closed()
classification = parse_utils.classify_stmt(operation)
if classification == parse_utils.STMT_DDL:
raise ProgrammingError(
"Executing DDL statements with executemany() method is not allowed."
)
many_result_set = StreamedManyResultSets()
if classification in (parse_utils.STMT_INSERT, parse_utils.STMT_UPDATING):
statements = []
for params in seq_of_params:
sql, params = parse_utils.sql_pyformat_args_to_spanner(
operation, params
)
statements.append((sql, params, get_param_types(params)))
if self.connection.autocommit:
self.connection.database.run_in_transaction(
self._do_batch_update, statements, many_result_set
)
else:
retried = False
while True:
try:
transaction = self.connection.transaction_checkout()
res_checksum = ResultsChecksum()
if not retried:
self.connection._statements.append(
(statements, res_checksum)
)
status, res = transaction.batch_update(statements)
many_result_set.add_iter(res)
res_checksum.consume_result(res)
res_checksum.consume_result(status.code)
if status.code == ABORTED:
self.connection._transaction = None
raise Aborted(status.details)
elif status.code != OK:
raise OperationalError(status.details)
break
except Aborted:
self.connection.retry_transaction()
retried = True
else:
for params in seq_of_params:
self.execute(operation, params)
many_result_set.add_iter(self._itr)
self._result_set = many_result_set
self._itr = many_result_set