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)