From ef18bbbdeb18251dd03f5e46b644c605bafc67e0 Mon Sep 17 00:00:00 2001 From: larkee Date: Wed, 15 Sep 2021 21:15:33 +1000 Subject: [PATCH 1/4] perf: use batch_update instead of sending a separate request for each statement --- google/cloud/spanner_dbapi/_helpers.py | 4 +++- tests/unit/spanner_dbapi/test__helpers.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/google/cloud/spanner_dbapi/_helpers.py b/google/cloud/spanner_dbapi/_helpers.py index 2fcdd59137..6d17b19629 100644 --- a/google/cloud/spanner_dbapi/_helpers.py +++ b/google/cloud/spanner_dbapi/_helpers.py @@ -56,10 +56,12 @@ 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)) + transaction.batch_update(statements) 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..b318360ac4 100644 --- a/tests/unit/spanner_dbapi/test__helpers.py +++ b/tests/unit/spanner_dbapi/test__helpers.py @@ -32,12 +32,12 @@ 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() + transaction.batch_update = mock_batch = mock.MagicMock() _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_homogenous(self): from google.cloud.spanner_dbapi import _helpers From 3b62eef72e89224e44e0579ca8e8ffa7873f3a7a Mon Sep 17 00:00:00 2001 From: larkee Date: Thu, 16 Sep 2021 10:13:16 +1000 Subject: [PATCH 2/4] fix: check status for errors --- google/cloud/spanner_dbapi/_helpers.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/google/cloud/spanner_dbapi/_helpers.py b/google/cloud/spanner_dbapi/_helpers.py index 6d17b19629..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 @@ -61,7 +61,9 @@ def _execute_insert_heterogenous(transaction, sql_params_list): sql, params = sql_pyformat_args_to_spanner(sql, params) param_types = get_param_types(params) statements.append((sql, params, param_types)) - transaction.batch_update(statements) + status, _ = transaction.batch_update(statements) + if status.code != OK: + raise OperationalError(status.message) def _execute_insert_homogenous(transaction, parts): From 858fb4aeb8d9b28464cb7cbe3a5b6ee36e05e840 Mon Sep 17 00:00:00 2001 From: larkee Date: Thu, 16 Sep 2021 10:13:44 +1000 Subject: [PATCH 3/4] test: add and update tests --- tests/unit/spanner_dbapi/test__helpers.py | 31 ++++++++++++++++++++- tests/unit/spanner_dbapi/test_connection.py | 12 ++++++-- 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/tests/unit/spanner_dbapi/test__helpers.py b/tests/unit/spanner_dbapi/test__helpers.py index b318360ac4..586ebbd912 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,13 +34,40 @@ 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.batch_update = mock_batch = 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_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) From 35adea9497906d6739de5e907fcfab21542a8981 Mon Sep 17 00:00:00 2001 From: Owl Bot Date: Wed, 13 Oct 2021 00:21:57 +0000 Subject: [PATCH 4/4] =?UTF-8?q?=F0=9F=A6=89=20Updates=20from=20OwlBot?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --- tests/unit/spanner_dbapi/test__helpers.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/tests/unit/spanner_dbapi/test__helpers.py b/tests/unit/spanner_dbapi/test__helpers.py index 586ebbd912..297cc9d679 100644 --- a/tests/unit/spanner_dbapi/test__helpers.py +++ b/tests/unit/spanner_dbapi/test__helpers.py @@ -35,7 +35,9 @@ def test__execute_insert_heterogenous(self): ) as mock_param_types: transaction = mock.MagicMock() status = Status(code=OK) - transaction.batch_update = mock_batch = mock.MagicMock(return_value=(status, 1)) + 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]) @@ -51,15 +53,17 @@ def test__execute_insert_heterogenous_error(self): sql = "sql" params = (sql, None) with mock.patch( - "google.cloud.spanner_dbapi._helpers.sql_pyformat_args_to_spanner", - return_value=params, + "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 + "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)) + transaction.batch_update = mock_batch = mock.MagicMock( + return_value=(status, 0) + ) with self.assertRaises(OperationalError): _helpers._execute_insert_heterogenous(transaction, [params])