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