def test_fetchmany_retry_aborted(self):
"""Check that aborted fetch re-executing transaction."""
from google.api_core.exceptions import Aborted
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.connection import connect
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()
with mock.patch(
"google.cloud.spanner_dbapi.cursor.Cursor.__next__",
side_effect=(Aborted("Aborted"), None),
):
with mock.patch(
"google.cloud.spanner_dbapi.connection.Connection.retry_transaction"
) as retry_mock:
cursor.fetchmany()
retry_mock.assert_called_with()
async def test_other_error_details_present():
any1 = Any()
any1.Pack(RetryInfo())
any2 = Any()
any2.Pack(ErrorInfo(reason="RESET", domain="pubsublite.googleapis.com"))
status_pb = Status(code=10, details=[any1, any2])
assert is_reset_signal(Aborted("", response=make_call(status_pb)))
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_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_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_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_peek_iterator_aborted_autocommit(self, mock_client):
"""
Checking that an Aborted exception is retried in case it happened while
streaming the first element with a PeekIterator in autocommit mode.
"""
from google.api_core.exceptions import Aborted
from google.cloud.spanner_dbapi.connection import connect
connection = connect("test-instance", "test-database")
connection.autocommit = True
cursor = connection.cursor()
with mock.patch(
"google.cloud.spanner_dbapi.utils.PeekIterator.__init__",
side_effect=(Aborted("Aborted"), None),
):
with mock.patch(
"google.cloud.spanner_dbapi.connection.Connection.retry_transaction"
) as retry_mock:
with mock.patch(
"google.cloud.spanner_dbapi.connection.Connection.run_statement",
return_value=((1, 2, 3), None),
):
with mock.patch(
"google.cloud.spanner_v1.database.Database.snapshot"
):
cursor.execute("SELECT * FROM table_name")
retry_mock.assert_called_with()
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 _do_batch_update(self, transaction, statements, many_result_set):
status, res = transaction.batch_update(statements)
many_result_set.add_iter(res)
if status.code == ABORTED:
raise Aborted(status.details)
elif status.code != OK:
raise OperationalError(status.details)
def _do_batch_update(self, transaction, statements, many_result_set):
status, res = transaction.batch_update(statements)
many_result_set.add_iter(res)
if status.code == ABORTED:
raise Aborted(status.message)
elif status.code != OK:
raise OperationalError(status.message)
self._row_count = sum([max(val, 0) for val in res])
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 _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 test_read_only_not_retried(self):
"""
Testing the unlikely case of a read-only transaction
failed with Aborted exception. In this case the
transaction should not be automatically retried.
"""
from google.api_core.exceptions import Aborted
connection = self._make_connection(read_only=True)
connection.retry_transaction = mock.Mock()
cursor = connection.cursor()
cursor._itr = mock.Mock(
__next__=mock.Mock(side_effect=Aborted("Aborted"), ))
cursor.fetchone()
cursor.fetchall()
cursor.fetchmany(5)
connection.retry_transaction.assert_not_called()
def test_fetchall_retry_aborted(self, mock_client):
"""Check that aborted fetch re-executing transaction."""
from google.api_core.exceptions import Aborted
from google.cloud.spanner_dbapi.checksum import ResultsChecksum
from google.cloud.spanner_dbapi.connection import connect
connection = connect("test-instance", "test-database")
cursor = connection.cursor()
cursor._checksum = ResultsChecksum()
with mock.patch(
"google.cloud.spanner_dbapi.cursor.Cursor.__iter__",
side_effect=(Aborted("Aborted"), iter([])),
):
with mock.patch(
"google.cloud.spanner_dbapi.connection.Connection.retry_transaction"
) as retry_mock:
cursor.fetchall()
retry_mock.assert_called_with()
def make_reset_signal() -> GoogleAPICallError:
any = Any()
any.Pack(ErrorInfo(reason="RESET", domain="pubsublite.googleapis.com"))
status_pb = Status(code=10, details=[any])
return Aborted("", response=make_call(status_pb))
async def test_extracted_status_is_none():
status_pb = Status(code=10, details=[])
assert not is_reset_signal(
Aborted("", response=make_call_without_metadata(status_pb))
)
async def test_wrong_domain():
any = Any()
any.Pack(ErrorInfo(reason="RESET", domain="other.googleapis.com"))
status_pb = Status(code=10, details=[any])
assert not is_reset_signal(Aborted("", response=make_call(status_pb)))
async def test_wrong_error_detail():
any = Any()
any.Pack(RetryInfo())
status_pb = Status(code=10, details=[any])
assert not is_reset_signal(Aborted("", response=make_call(status_pb)))
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
async def test_missing_call():
assert not is_reset_signal(Aborted(""))
