diff --git a/google/cloud/spanner_dbapi/_helpers.py b/google/cloud/spanner_dbapi/_helpers.py index 2fcdd59137..4f8871c7dd 100644 --- a/google/cloud/spanner_dbapi/_helpers.py +++ b/google/cloud/spanner_dbapi/_helpers.py @@ -11,12 +11,12 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - +from google.cloud.spanner_dbapi.exceptions import OperationalError from google.cloud.spanner_dbapi.parse_utils import get_param_types from google.cloud.spanner_dbapi.parse_utils import parse_insert from google.cloud.spanner_dbapi.parse_utils import sql_pyformat_args_to_spanner from google.cloud.spanner_v1 import param_types - +from google.rpc.code_pb2 import OK SQL_LIST_TABLES = """ SELECT @@ -56,10 +56,14 @@ def _execute_insert_heterogenous(transaction, sql_params_list): + statements = [] for sql, params in sql_params_list: sql, params = sql_pyformat_args_to_spanner(sql, params) param_types = get_param_types(params) - transaction.execute_update(sql, params=params, param_types=param_types) + statements.append((sql, params, param_types)) + status, _ = transaction.batch_update(statements) + if status.code != OK: + raise OperationalError(status.message) def _execute_insert_homogenous(transaction, parts): diff --git a/tests/unit/spanner_dbapi/test__helpers.py b/tests/unit/spanner_dbapi/test__helpers.py index 84d6b3e323..297cc9d679 100644 --- a/tests/unit/spanner_dbapi/test__helpers.py +++ b/tests/unit/spanner_dbapi/test__helpers.py @@ -21,6 +21,8 @@ class TestHelpers(unittest.TestCase): def test__execute_insert_heterogenous(self): from google.cloud.spanner_dbapi import _helpers + from google.rpc.status_pb2 import Status + from google.rpc.code_pb2 import OK sql = "sql" params = (sql, None) @@ -32,12 +34,43 @@ def test__execute_insert_heterogenous(self): "google.cloud.spanner_dbapi._helpers.get_param_types", return_value=None ) as mock_param_types: transaction = mock.MagicMock() - transaction.execute_update = mock_execute = mock.MagicMock() + status = Status(code=OK) + transaction.batch_update = mock_batch = mock.MagicMock( + return_value=(status, 1) + ) _helpers._execute_insert_heterogenous(transaction, [params]) mock_pyformat.assert_called_once_with(params[0], params[1]) mock_param_types.assert_called_once_with(None) - mock_execute.assert_called_once_with(sql, params=None, param_types=None) + mock_batch.assert_called_once_with([(sql, None, None)]) + + def test__execute_insert_heterogenous_error(self): + from google.cloud.spanner_dbapi import _helpers + from google.cloud.spanner_dbapi import OperationalError + from google.rpc.status_pb2 import Status + from google.rpc.code_pb2 import UNKNOWN + + sql = "sql" + params = (sql, None) + with mock.patch( + "google.cloud.spanner_dbapi._helpers.sql_pyformat_args_to_spanner", + return_value=params, + ) as mock_pyformat: + with mock.patch( + "google.cloud.spanner_dbapi._helpers.get_param_types", return_value=None + ) as mock_param_types: + transaction = mock.MagicMock() + status = Status(code=UNKNOWN) + transaction.batch_update = mock_batch = mock.MagicMock( + return_value=(status, 0) + ) + + with self.assertRaises(OperationalError): + _helpers._execute_insert_heterogenous(transaction, [params]) + + mock_pyformat.assert_called_once_with(params[0], params[1]) + mock_param_types.assert_called_once_with(None) + mock_batch.assert_called_once_with([(sql, None, None)]) def test__execute_insert_homogenous(self): from google.cloud.spanner_dbapi import _helpers diff --git a/tests/unit/spanner_dbapi/test_connection.py b/tests/unit/spanner_dbapi/test_connection.py index 34e50255f9..2b0938637f 100644 --- a/tests/unit/spanner_dbapi/test_connection.py +++ b/tests/unit/spanner_dbapi/test_connection.py @@ -387,13 +387,17 @@ 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() - connection.transaction_checkout = mock.Mock() + 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) @@ -404,13 +408,17 @@ 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 + from google.rpc.status_pb2 import Status + from google.rpc.code_pb2 import OK 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() + 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)