From 433ba9060b5fafa949761f8c3339b511d05c5d6a Mon Sep 17 00:00:00 2001 From: larkee Date: Wed, 15 Sep 2021 21:15:33 +1000 Subject: [PATCH 01/13] 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 83172a3f51..02a379fec1 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 85927c9ffdce8c598bac43c29beb7ce2b0d0d303 Mon Sep 17 00:00:00 2001 From: larkee Date: Thu, 16 Sep 2021 10:13:16 +1000 Subject: [PATCH 02/13] 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 02a379fec1..bcd7a6bfa8 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 007f8b130bf13cbdcc770d365e8813a3d4f06a6e Mon Sep 17 00:00:00 2001 From: larkee Date: Thu, 16 Sep 2021 10:13:44 +1000 Subject: [PATCH 03/13] 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 0eea3eaf5b..2f0ce131c3 100644 --- a/tests/unit/spanner_dbapi/test_connection.py +++ b/tests/unit/spanner_dbapi/test_connection.py @@ -388,13 +388,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) @@ -405,13 +409,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 068289a1769fbf17a78d8fb995891a3ecaaf336f Mon Sep 17 00:00:00 2001 From: larkee Date: Thu, 16 Sep 2021 14:49:22 +1000 Subject: [PATCH 04/13] refactor: remove redundant mutations code --- google/cloud/spanner_dbapi/_helpers.py | 38 +------ google/cloud/spanner_dbapi/connection.py | 25 ++--- google/cloud/spanner_dbapi/parse_utils.py | 48 ++++----- tests/unit/spanner_dbapi/test__helpers.py | 11 -- tests/unit/spanner_dbapi/test_cursor.py | 3 - tests/unit/spanner_dbapi/test_parse_utils.py | 102 ++++++++----------- 6 files changed, 73 insertions(+), 154 deletions(-) diff --git a/google/cloud/spanner_dbapi/_helpers.py b/google/cloud/spanner_dbapi/_helpers.py index bcd7a6bfa8..d57e021551 100644 --- a/google/cloud/spanner_dbapi/_helpers.py +++ b/google/cloud/spanner_dbapi/_helpers.py @@ -66,41 +66,11 @@ def _execute_insert_heterogenous(transaction, sql_params_list): raise OperationalError(status.message) -def _execute_insert_homogenous(transaction, parts): - # Perform an insert in one shot. - table = parts.get("table") - columns = parts.get("columns") - values = parts.get("values") - return transaction.insert(table, columns, values) - - def handle_insert(connection, sql, params): - parts = parse_insert(sql, params) - - # The split between the two styles exists because: - # in the common case of multiple values being passed - # with simple pyformat arguments, - # SQL: INSERT INTO T (f1, f2) VALUES (%s, %s, %s) - # Params: [(1, 2, 3, 4, 5, 6, 7, 8, 9, 10,)] - # we can take advantage of a single RPC with: - # transaction.insert(table, columns, values) - # instead of invoking: - # with transaction: - # for sql, params in sql_params_list: - # transaction.execute_sql(sql, params, param_types) - # which invokes more RPCs and is more costly. - - if parts.get("homogenous"): - # The common case of multiple values being passed in - # non-complex pyformat args and need to be uploaded in one RPC. - return connection.database.run_in_transaction(_execute_insert_homogenous, parts) - else: - # All the other cases that are esoteric and need - # transaction.execute_sql - sql_params_list = parts.get("sql_params_list") - return connection.database.run_in_transaction( - _execute_insert_heterogenous, sql_params_list - ) + sql_params_list = parse_insert(sql, params) + return connection.database.run_in_transaction( + _execute_insert_heterogenous, sql_params_list + ) class ColumnInfo: diff --git a/google/cloud/spanner_dbapi/connection.py b/google/cloud/spanner_dbapi/connection.py index e6d1d64db1..26bab9fe63 100644 --- a/google/cloud/spanner_dbapi/connection.py +++ b/google/cloud/spanner_dbapi/connection.py @@ -24,7 +24,6 @@ from google.cloud.spanner_v1.snapshot import Snapshot from google.cloud.spanner_dbapi._helpers import _execute_insert_heterogenous -from google.cloud.spanner_dbapi._helpers import _execute_insert_homogenous from google.cloud.spanner_dbapi._helpers import parse_insert from google.cloud.spanner_dbapi.checksum import _compare_checksums from google.cloud.spanner_dbapi.checksum import ResultsChecksum @@ -430,22 +429,14 @@ def run_statement(self, statement, retried=False): self._statements.append(statement) if statement.is_insert: - parts = parse_insert(statement.sql, statement.params) - - if parts.get("homogenous"): - _execute_insert_homogenous(transaction, parts) - return ( - iter(()), - ResultsChecksum() if retried else statement.checksum, - ) - else: - _execute_insert_heterogenous( - transaction, parts.get("sql_params_list"), - ) - return ( - iter(()), - ResultsChecksum() if retried else statement.checksum, - ) + sql_params_list = parse_insert(statement.sql, statement.params) + _execute_insert_heterogenous( + transaction, sql_params_list, + ) + return ( + iter(()), + ResultsChecksum() if retried else statement.checksum, + ) return ( transaction.execute_sql( diff --git a/google/cloud/spanner_dbapi/parse_utils.py b/google/cloud/spanner_dbapi/parse_utils.py index 4f55a7b2c4..031fb1128c 100644 --- a/google/cloud/spanner_dbapi/parse_utils.py +++ b/google/cloud/spanner_dbapi/parse_utils.py @@ -220,42 +220,34 @@ def parse_insert(insert_sql, params): Case a) SQL: INSERT INTO T (f1, f2) VALUES (1, 2) it produces: - { - 'sql_params_list': [ - ('INSERT INTO T (f1, f2) VALUES (1, 2)', None), - ], - } + [ + ('INSERT INTO T (f1, f2) VALUES (1, 2)', None), + ] Case b) SQL: 'INSERT INTO T (s, c) SELECT st, zc FROM cus WHERE col IN (%s, %s)', it produces: - { - 'sql_params_list': [ - ('INSERT INTO T (s, c) SELECT st, zc FROM cus ORDER BY fn, ln', ('a', 'b')), - ] - } + [ + ('INSERT INTO T (s, c) SELECT st, zc FROM cus ORDER BY fn, ln', ('a', 'b')), + ] Case c) SQL: INSERT INTO T (f1, f2) VALUES (%s, %s), (%s, %s) Params: ['a', 'b', 'c', 'd'] it produces: - { - 'sql_params_list': [ - ('INSERT INTO T (f1, f2) VALUES (%s, %s)', ('a', 'b')), - ('INSERT INTO T (f1, f2) VALUES (%s, %s)', ('c', 'd')) - ], - } + [ + ('INSERT INTO T (f1, f2) VALUES (%s, %s)', ('a', 'b')), + ('INSERT INTO T (f1, f2) VALUES (%s, %s)', ('c', 'd')) + ] Case d) SQL: INSERT INTO T (f1, f2) VALUES (%s, LOWER(%s)), (UPPER(%s), %s) Params: ['a', 'b', 'c', 'd'] it produces: - { - 'sql_params_list': [ - ('INSERT INTO T (f1, f2) VALUES (%s, LOWER(%s))', ('a', 'b',)), - ('INSERT INTO T (f1, f2) VALUES (UPPER(%s), %s)', ('c', 'd',)) - ], - } + [ + ('INSERT INTO T (f1, f2) VALUES (%s, LOWER(%s))', ('a', 'b',)), + ('INSERT INTO T (f1, f2) VALUES (UPPER(%s), %s)', ('c', 'd',)) + ] :type insert_sql: str :param insert_sql: A SQL insert request. @@ -264,9 +256,7 @@ def parse_insert(insert_sql, params): :param params: A list of parameters. :rtype: dict - :returns: A dictionary that maps `sql_params_list` to the list of - parameters in cases a), b), d) or the dictionary with information - about the resulting table in case c). + :returns: A list of (sql, param) tuples """ # noqa match = RE_INSERT.search(insert_sql) @@ -279,7 +269,7 @@ def parse_insert(insert_sql, params): if not after_values_sql: # Case b) insert_sql = sanitize_literals_for_upload(insert_sql) - return {"sql_params_list": [(insert_sql, params)]} + return [(insert_sql, params)] if not params: # Case a) perhaps? @@ -300,7 +290,7 @@ def parse_insert(insert_sql, params): # Confirmed case of: # SQL: INSERT INTO T (a1, a2) VALUES (1, 2) # Params: None - return {"sql_params_list": [(insert_sql, None)]} + return [(insert_sql, None)] values_str = after_values_sql[0] _, values = parse_values(values_str) @@ -321,7 +311,7 @@ def parse_insert(insert_sql, params): for row in rows_list: sql_params_list.append((insert_sql_preamble, row)) - return {"sql_params_list": sql_params_list} + return sql_params_list # Case d) # insert_sql is of the form: @@ -349,7 +339,7 @@ def parse_insert(insert_sql, params): ) sql_param_tuples.append((row_sql, row_params)) - return {"sql_params_list": sql_param_tuples} + return sql_param_tuples def rows_for_insert_or_update(columns, params, pyformat_args=None): diff --git a/tests/unit/spanner_dbapi/test__helpers.py b/tests/unit/spanner_dbapi/test__helpers.py index 586ebbd912..abab05d3d6 100644 --- a/tests/unit/spanner_dbapi/test__helpers.py +++ b/tests/unit/spanner_dbapi/test__helpers.py @@ -68,17 +68,6 @@ def test__execute_insert_heterogenous_error(self): 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 - - transaction = mock.MagicMock() - transaction.insert = mock.MagicMock() - parts = mock.MagicMock() - parts.get = mock.MagicMock(return_value=0) - - _helpers._execute_insert_homogenous(transaction, parts) - transaction.insert.assert_called_once_with(0, 0, 0) - def test_handle_insert(self): from google.cloud.spanner_dbapi import _helpers diff --git a/tests/unit/spanner_dbapi/test_cursor.py b/tests/unit/spanner_dbapi/test_cursor.py index 90d07eb3db..97297b2b59 100644 --- a/tests/unit/spanner_dbapi/test_cursor.py +++ b/tests/unit/spanner_dbapi/test_cursor.py @@ -706,9 +706,6 @@ def test_setoutputsize(self): # # def test_do_execute_insert_heterogenous(self): # pass - # - # def test_do_execute_insert_homogenous(self): - # pass def test_handle_dql(self): from google.cloud.spanner_dbapi import utils diff --git a/tests/unit/spanner_dbapi/test_parse_utils.py b/tests/unit/spanner_dbapi/test_parse_utils.py index 994b02d615..2d52a6a82b 100644 --- a/tests/unit/spanner_dbapi/test_parse_utils.py +++ b/tests/unit/spanner_dbapi/test_parse_utils.py @@ -73,92 +73,74 @@ def test_parse_insert(self): ( "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", [1, 2, 3, 4, 5, 6], - { - "sql_params_list": [ - ( - "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", - (1, 2, 3), - ), - ( - "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", - (4, 5, 6), - ), - ] - }, + [ + ( + "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", + (1, 2, 3), + ), + ( + "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", + (4, 5, 6), + ), + ], ), ( "INSERT INTO django_migrations(app, name, applied) VALUES (%s, %s, %s)", [1, 2, 3, 4, 5, 6], - { - "sql_params_list": [ - ( - "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", - (1, 2, 3), - ), - ( - "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", - (4, 5, 6), - ), - ] - }, + [ + ("INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", (1, 2, 3)), + ("INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", (4, 5, 6)), + ], ), ( "INSERT INTO sales.addresses (street, city, state, zip_code) " "SELECT street, city, state, zip_code FROM sales.customers" "ORDER BY first_name, last_name", None, - { - "sql_params_list": [ - ( - "INSERT INTO sales.addresses (street, city, state, zip_code) " - "SELECT street, city, state, zip_code FROM sales.customers" - "ORDER BY first_name, last_name", - None, - ) - ] - }, + [ + ( + "INSERT INTO sales.addresses (street, city, state, zip_code) " + "SELECT street, city, state, zip_code FROM sales.customers" + "ORDER BY first_name, last_name", + None, + ) + ], ), ( "INSERT INTO ap (n, ct, cn) " "VALUES (%s, %s, %s), (%s, %s, %s), (%s, %s, %s),(%s, %s, %s)", (1, 2, 3, 4, 5, 6, 7, 8, 9), - { - "sql_params_list": [ - ("INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s)", (1, 2, 3)), - ("INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s)", (4, 5, 6)), - ("INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s)", (7, 8, 9)), - ] - }, + [ + ("INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s)", (1, 2, 3)), + ("INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s)", (4, 5, 6)), + ("INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s)", (7, 8, 9)), + ], ), ( "INSERT INTO `no` (`yes`) VALUES (%s)", (1, 4, 5), - { - "sql_params_list": [ - ("INSERT INTO `no` (`yes`) VALUES (%s)", (1,)), - ("INSERT INTO `no` (`yes`) VALUES (%s)", (4,)), - ("INSERT INTO `no` (`yes`) VALUES (%s)", (5,)), - ] - }, + [ + ("INSERT INTO `no` (`yes`) VALUES (%s)", (1,)), + ("INSERT INTO `no` (`yes`) VALUES (%s)", (4,)), + ("INSERT INTO `no` (`yes`) VALUES (%s)", (5,)), + ], ), ( "INSERT INTO T (f1, f2) VALUES (1, 2)", None, - {"sql_params_list": [("INSERT INTO T (f1, f2) VALUES (1, 2)", None)]}, + [("INSERT INTO T (f1, f2) VALUES (1, 2)", None)], ), ( "INSERT INTO `no` (`yes`, tiff) VALUES (%s, LOWER(%s)), (%s, %s), (%s, %s)", (1, "FOO", 5, 10, 11, 29), - { - "sql_params_list": [ - ( - "INSERT INTO `no` (`yes`, tiff) VALUES(%s, LOWER(%s))", - (1, "FOO"), - ), - ("INSERT INTO `no` (`yes`, tiff) VALUES(%s, %s)", (5, 10)), - ("INSERT INTO `no` (`yes`, tiff) VALUES(%s, %s)", (11, 29)), - ] - }, + [ + ( + "INSERT INTO `no` (`yes`, tiff) VALUES(%s, LOWER(%s))", + (1, "FOO"), + ), + ("INSERT INTO `no` (`yes`, tiff) VALUES(%s, %s)", (5, 10)), + ("INSERT INTO `no` (`yes`, tiff) VALUES(%s, %s)", (11, 29)), + ], ), ] @@ -425,5 +407,5 @@ def test_insert_from_select(self): ARGS = [5, "data2", "data3"] self.assertEqual( - parse_insert(SQL, ARGS), {"sql_params_list": [(SQL, ARGS)]}, + parse_insert(SQL, ARGS), [(SQL, ARGS)], ) From f6125c5495ae9e75b87fd031c496d5292af0b1bc Mon Sep 17 00:00:00 2001 From: larkee Date: Fri, 15 Oct 2021 11:58:44 +1100 Subject: [PATCH 05/13] perf: generate fewer insert statements for execute calls --- google/cloud/spanner_dbapi/parse_utils.py | 20 ++++++------- tests/unit/spanner_dbapi/test_parse_utils.py | 31 +++++++++----------- 2 files changed, 24 insertions(+), 27 deletions(-) diff --git a/google/cloud/spanner_dbapi/parse_utils.py b/google/cloud/spanner_dbapi/parse_utils.py index 031fb1128c..7519fdb88b 100644 --- a/google/cloud/spanner_dbapi/parse_utils.py +++ b/google/cloud/spanner_dbapi/parse_utils.py @@ -299,19 +299,19 @@ def parse_insert(insert_sql, params): # Case c) columns = [mi.strip(" `") for mi in match.group("columns").split(",")] - sql_params_list = [] - insert_sql_preamble = "INSERT INTO %s (%s) VALUES %s" % ( + values_pyformat = [str(arg) for arg in values.argv] + rows_list = rows_for_insert_or_update(columns, params, values_pyformat) + values_template = ', '.join([str(values.argv[0])] * len(rows_list)) + insert_sql = "INSERT INTO {} ({}) VALUES {}".format( match.group("table_name"), match.group("columns"), - values.argv[0], + values_template, ) - values_pyformat = [str(arg) for arg in values.argv] - rows_list = rows_for_insert_or_update(columns, params, values_pyformat) - insert_sql_preamble = sanitize_literals_for_upload(insert_sql_preamble) + insert_sql = sanitize_literals_for_upload(insert_sql) + flat_params = [] for row in rows_list: - sql_params_list.append((insert_sql_preamble, row)) - - return sql_params_list + flat_params.extend(row) + return [(insert_sql, tuple(flat_params))] # Case d) # insert_sql is of the form: @@ -331,7 +331,7 @@ def parse_insert(insert_sql, params): sql_param_tuples = [] for token_arg in values.argv: - row_sql = before_values_sql + " VALUES%s" % token_arg + row_sql = before_values_sql + " VALUES %s" % token_arg row_sql = sanitize_literals_for_upload(row_sql) row_params, params = ( tuple(params[0 : len(token_arg)]), diff --git a/tests/unit/spanner_dbapi/test_parse_utils.py b/tests/unit/spanner_dbapi/test_parse_utils.py index 2d52a6a82b..e454d962dc 100644 --- a/tests/unit/spanner_dbapi/test_parse_utils.py +++ b/tests/unit/spanner_dbapi/test_parse_utils.py @@ -75,12 +75,8 @@ def test_parse_insert(self): [1, 2, 3, 4, 5, 6], [ ( - "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", - (1, 2, 3), - ), - ( - "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", - (4, 5, 6), + "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s), (%s, %s, %s)", + (1, 2, 3, 4, 5, 6), ), ], ), @@ -88,8 +84,10 @@ def test_parse_insert(self): "INSERT INTO django_migrations(app, name, applied) VALUES (%s, %s, %s)", [1, 2, 3, 4, 5, 6], [ - ("INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", (1, 2, 3)), - ("INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", (4, 5, 6)), + ( + "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s), (%s, %s, %s)", + (1, 2, 3, 4, 5, 6), + ), ], ), ( @@ -111,18 +109,17 @@ def test_parse_insert(self): "VALUES (%s, %s, %s), (%s, %s, %s), (%s, %s, %s),(%s, %s, %s)", (1, 2, 3, 4, 5, 6, 7, 8, 9), [ - ("INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s)", (1, 2, 3)), - ("INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s)", (4, 5, 6)), - ("INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s)", (7, 8, 9)), + ( + "INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s), (%s, %s, %s), (%s, %s, %s)", + (1, 2, 3, 4, 5, 6, 7, 8, 9) + ), ], ), ( "INSERT INTO `no` (`yes`) VALUES (%s)", (1, 4, 5), [ - ("INSERT INTO `no` (`yes`) VALUES (%s)", (1,)), - ("INSERT INTO `no` (`yes`) VALUES (%s)", (4,)), - ("INSERT INTO `no` (`yes`) VALUES (%s)", (5,)), + ("INSERT INTO `no` (`yes`) VALUES (%s), (%s), (%s)", (1, 4, 5)) ], ), ( @@ -135,11 +132,11 @@ def test_parse_insert(self): (1, "FOO", 5, 10, 11, 29), [ ( - "INSERT INTO `no` (`yes`, tiff) VALUES(%s, LOWER(%s))", + "INSERT INTO `no` (`yes`, tiff) VALUES (%s, LOWER(%s))", (1, "FOO"), ), - ("INSERT INTO `no` (`yes`, tiff) VALUES(%s, %s)", (5, 10)), - ("INSERT INTO `no` (`yes`, tiff) VALUES(%s, %s)", (11, 29)), + ("INSERT INTO `no` (`yes`, tiff) VALUES (%s, %s)", (5, 10)), + ("INSERT INTO `no` (`yes`, tiff) VALUES (%s, %s)", (11, 29)), ], ), ] From 7a7aae2a7934fe0f91833860be16396fdf361a46 Mon Sep 17 00:00:00 2001 From: larkee Date: Fri, 15 Oct 2021 11:59:11 +1100 Subject: [PATCH 06/13] perf: generate fewer insert statements for executemany calls --- google/cloud/spanner_dbapi/cursor.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/google/cloud/spanner_dbapi/cursor.py b/google/cloud/spanner_dbapi/cursor.py index 112fcda291..c2a3e34669 100644 --- a/google/cloud/spanner_dbapi/cursor.py +++ b/google/cloud/spanner_dbapi/cursor.py @@ -270,7 +270,13 @@ def executemany(self, operation, seq_of_params): many_result_set = StreamedManyResultSets() - if classification in (parse_utils.STMT_INSERT, parse_utils.STMT_UPDATING): + if classification == parse_utils.STMT_INSERT: + flat_params = [] + for params in seq_of_params: + flat_params.extend(params) + self.execute(operation, flat_params) + + elif classification == parse_utils.STMT_UPDATING: statements = [] for params in seq_of_params: From 501a625bba3c11e21d96c4e92811af2a6895aa53 Mon Sep 17 00:00:00 2001 From: larkee Date: Fri, 15 Oct 2021 17:39:00 +1100 Subject: [PATCH 07/13] test: fix executemany insert tests --- tests/unit/spanner_dbapi/test_cursor.py | 65 ++++++++----------------- 1 file changed, 21 insertions(+), 44 deletions(-) diff --git a/tests/unit/spanner_dbapi/test_cursor.py b/tests/unit/spanner_dbapi/test_cursor.py index 97297b2b59..f8351282d0 100644 --- a/tests/unit/spanner_dbapi/test_cursor.py +++ b/tests/unit/spanner_dbapi/test_cursor.py @@ -437,14 +437,9 @@ def test_executemany_insert_batch_non_autocommit(self): transaction.batch_update.assert_called_once_with( [ ( - """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""", - {"a0": 1, "a1": 2, "a2": 3, "a3": 4}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64}, - ), - ( - """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""", - {"a0": 5, "a1": 6, "a2": 7, "a3": 8}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64}, + """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3), (@a4, @a5, @a6, @a7)""", + {"a0": 1, "a1": 2, "a2": 3, "a3": 4, "a4": 5, "a5": 6, "a6": 7, "a7": 8}, + {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64, "a4": INT64, "a5": INT64, "a6": INT64, "a7": INT64}, ), ] ) @@ -477,14 +472,9 @@ def test_executemany_insert_batch_autocommit(self): transaction.batch_update.assert_called_once_with( [ ( - """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""", - {"a0": 1, "a1": 2, "a2": 3, "a3": 4}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64}, - ), - ( - """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""", - {"a0": 5, "a1": 6, "a2": 7, "a3": 8}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64}, + """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3), (@a4, @a5, @a6, @a7)""", + {"a0": 1, "a1": 2, "a2": 3, "a3": 4, "a4": 5, "a5": 6, "a6": 7, "a7": 8}, + {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64, "a4": INT64, "a5": INT64, "a6": INT64, "a7": INT64}, ), ] ) @@ -533,7 +523,7 @@ def test_executemany_insert_batch_aborted(self): transaction1 = mock.Mock(committed=False, rolled_back=False) transaction1.batch_update = mock.Mock( - side_effect=[(mock.Mock(code=ABORTED, details=err_details), [])] + side_effect=[(mock.Mock(code=ABORTED, message=err_details), [])] ) transaction2 = self._transaction_mock() @@ -549,28 +539,18 @@ def test_executemany_insert_batch_aborted(self): transaction1.batch_update.assert_called_with( [ ( - """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""", - {"a0": 1, "a1": 2, "a2": 3, "a3": 4}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64}, - ), - ( - """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""", - {"a0": 5, "a1": 6, "a2": 7, "a3": 8}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64}, + """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3), (@a4, @a5, @a6, @a7)""", + {"a0": 1, "a1": 2, "a2": 3, "a3": 4, "a4": 5, "a5": 6, "a6": 7, "a7": 8}, + {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64, "a4": INT64, "a5": INT64, "a6": INT64, "a7": INT64}, ), ] ) transaction2.batch_update.assert_called_with( [ ( - """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""", - {"a0": 1, "a1": 2, "a2": 3, "a3": 4}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64}, - ), - ( - """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""", - {"a0": 5, "a1": 6, "a2": 7, "a3": 8}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64}, + """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3), (@a4, @a5, @a6, @a7)""", + {"a0": 1, "a1": 2, "a2": 3, "a3": 4, "a4": 5, "a5": 6, "a6": 7, "a7": 8}, + {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64, "a4": INT64, "a5": INT64, "a6": INT64, "a7": INT64}, ), ] ) @@ -579,17 +559,14 @@ def test_executemany_insert_batch_aborted(self): self.assertEqual( connection._statements[0][0], [ - ( - """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""", - {"a0": 1, "a1": 2, "a2": 3, "a3": 4}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64}, - ), - ( - """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3)""", - {"a0": 5, "a1": 6, "a2": 7, "a3": 8}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64}, - ), - ], + [ + ( + """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3), (@a4, @a5, @a6, @a7)""", + {"a0": 1, "a1": 2, "a2": 3, "a3": 4, "a4": 5, "a5": 6, "a6": 7, "a7": 8}, + {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64, "a4": INT64, "a5": INT64, "a6": INT64, "a7": INT64}, + ), + ], + ] ) self.assertIsInstance(connection._statements[0][1], ResultsChecksum) From e67752675466e14a259fb979fb5b161eb3175676 Mon Sep 17 00:00:00 2001 From: larkee Date: Tue, 2 Nov 2021 10:03:41 +1100 Subject: [PATCH 08/13] fix: handle aborted errors for executemany inserts --- google/cloud/spanner_dbapi/_helpers.py | 5 ++++- google/cloud/spanner_dbapi/cursor.py | 23 +++++++++++++---------- tests/unit/spanner_dbapi/test_cursor.py | 12 +++++------- 3 files changed, 22 insertions(+), 18 deletions(-) diff --git a/google/cloud/spanner_dbapi/_helpers.py b/google/cloud/spanner_dbapi/_helpers.py index d57e021551..999c33c748 100644 --- a/google/cloud/spanner_dbapi/_helpers.py +++ b/google/cloud/spanner_dbapi/_helpers.py @@ -11,12 +11,13 @@ # 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.api_core.exceptions import Aborted 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 +from google.rpc.code_pb2 import OK, ABORTED SQL_LIST_TABLES = """ SELECT @@ -63,6 +64,8 @@ def _execute_insert_heterogenous(transaction, sql_params_list): statements.append((sql, params, param_types)) status, _ = transaction.batch_update(statements) if status.code != OK: + if status.code == ABORTED: + raise Aborted(status.details) raise OperationalError(status.message) diff --git a/google/cloud/spanner_dbapi/cursor.py b/google/cloud/spanner_dbapi/cursor.py index c2a3e34669..e055d0fc98 100644 --- a/google/cloud/spanner_dbapi/cursor.py +++ b/google/cloud/spanner_dbapi/cursor.py @@ -270,15 +270,15 @@ def executemany(self, operation, seq_of_params): many_result_set = StreamedManyResultSets() - if classification == parse_utils.STMT_INSERT: - flat_params = [] - for params in seq_of_params: - flat_params.extend(params) - self.execute(operation, flat_params) - - elif classification == parse_utils.STMT_UPDATING: + if classification in (parse_utils.STMT_INSERT, parse_utils.STMT_UPDATING): statements = [] + if classification == parse_utils.STMT_INSERT: + flat_params = [] + for params in seq_of_params: + flat_params.extend(params) + operation, params = parse_utils.parse_insert(operation, flat_params)[0] + seq_of_params = [params] for params in seq_of_params: sql, params = parse_utils.sql_pyformat_args_to_spanner( operation, params @@ -286,9 +286,12 @@ def executemany(self, operation, seq_of_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 - ) + if classification == parse_utils.STMT_INSERT: + self.execute(operation, flat_params) + else: + self.connection.database.run_in_transaction( + self._do_batch_update, statements, many_result_set + ) else: retried = False while True: diff --git a/tests/unit/spanner_dbapi/test_cursor.py b/tests/unit/spanner_dbapi/test_cursor.py index f8351282d0..0625f52362 100644 --- a/tests/unit/spanner_dbapi/test_cursor.py +++ b/tests/unit/spanner_dbapi/test_cursor.py @@ -559,13 +559,11 @@ def test_executemany_insert_batch_aborted(self): self.assertEqual( connection._statements[0][0], [ - [ - ( - """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3), (@a4, @a5, @a6, @a7)""", - {"a0": 1, "a1": 2, "a2": 3, "a3": 4, "a4": 5, "a5": 6, "a6": 7, "a7": 8}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64, "a4": INT64, "a5": INT64, "a6": INT64, "a7": INT64}, - ), - ], + ( + """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3), (@a4, @a5, @a6, @a7)""", + {"a0": 1, "a1": 2, "a2": 3, "a3": 4, "a4": 5, "a5": 6, "a6": 7, "a7": 8}, + {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64, "a4": INT64, "a5": INT64, "a6": INT64, "a7": INT64}, + ), ] ) self.assertIsInstance(connection._statements[0][1], ResultsChecksum) From a37768352a42b55813befa77f2ee09e8195ae6e0 Mon Sep 17 00:00:00 2001 From: larkee Date: Wed, 3 Nov 2021 11:15:14 +1100 Subject: [PATCH 09/13] style: fix lint --- google/cloud/spanner_dbapi/parse_utils.py | 6 +- tests/unit/spanner_dbapi/test__helpers.py | 14 ++- tests/unit/spanner_dbapi/test_cursor.py | 112 +++++++++++++++++-- tests/unit/spanner_dbapi/test_parse_utils.py | 6 +- 4 files changed, 114 insertions(+), 24 deletions(-) diff --git a/google/cloud/spanner_dbapi/parse_utils.py b/google/cloud/spanner_dbapi/parse_utils.py index 7519fdb88b..f3d094748c 100644 --- a/google/cloud/spanner_dbapi/parse_utils.py +++ b/google/cloud/spanner_dbapi/parse_utils.py @@ -301,11 +301,9 @@ def parse_insert(insert_sql, params): columns = [mi.strip(" `") for mi in match.group("columns").split(",")] values_pyformat = [str(arg) for arg in values.argv] rows_list = rows_for_insert_or_update(columns, params, values_pyformat) - values_template = ', '.join([str(values.argv[0])] * len(rows_list)) + values_template = ", ".join([str(values.argv[0])] * len(rows_list)) insert_sql = "INSERT INTO {} ({}) VALUES {}".format( - match.group("table_name"), - match.group("columns"), - values_template, + match.group("table_name"), match.group("columns"), values_template, ) insert_sql = sanitize_literals_for_upload(insert_sql) flat_params = [] diff --git a/tests/unit/spanner_dbapi/test__helpers.py b/tests/unit/spanner_dbapi/test__helpers.py index abab05d3d6..be7526a11a 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]) diff --git a/tests/unit/spanner_dbapi/test_cursor.py b/tests/unit/spanner_dbapi/test_cursor.py index 0625f52362..4c45cbfac0 100644 --- a/tests/unit/spanner_dbapi/test_cursor.py +++ b/tests/unit/spanner_dbapi/test_cursor.py @@ -438,8 +438,26 @@ def test_executemany_insert_batch_non_autocommit(self): [ ( """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3), (@a4, @a5, @a6, @a7)""", - {"a0": 1, "a1": 2, "a2": 3, "a3": 4, "a4": 5, "a5": 6, "a6": 7, "a7": 8}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64, "a4": INT64, "a5": INT64, "a6": INT64, "a7": INT64}, + { + "a0": 1, + "a1": 2, + "a2": 3, + "a3": 4, + "a4": 5, + "a5": 6, + "a6": 7, + "a7": 8, + }, + { + "a0": INT64, + "a1": INT64, + "a2": INT64, + "a3": INT64, + "a4": INT64, + "a5": INT64, + "a6": INT64, + "a7": INT64, + }, ), ] ) @@ -473,8 +491,26 @@ def test_executemany_insert_batch_autocommit(self): [ ( """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3), (@a4, @a5, @a6, @a7)""", - {"a0": 1, "a1": 2, "a2": 3, "a3": 4, "a4": 5, "a5": 6, "a6": 7, "a7": 8}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64, "a4": INT64, "a5": INT64, "a6": INT64, "a7": INT64}, + { + "a0": 1, + "a1": 2, + "a2": 3, + "a3": 4, + "a4": 5, + "a5": 6, + "a6": 7, + "a7": 8, + }, + { + "a0": INT64, + "a1": INT64, + "a2": INT64, + "a3": INT64, + "a4": INT64, + "a5": INT64, + "a6": INT64, + "a7": INT64, + }, ), ] ) @@ -540,8 +576,26 @@ def test_executemany_insert_batch_aborted(self): [ ( """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3), (@a4, @a5, @a6, @a7)""", - {"a0": 1, "a1": 2, "a2": 3, "a3": 4, "a4": 5, "a5": 6, "a6": 7, "a7": 8}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64, "a4": INT64, "a5": INT64, "a6": INT64, "a7": INT64}, + { + "a0": 1, + "a1": 2, + "a2": 3, + "a3": 4, + "a4": 5, + "a5": 6, + "a6": 7, + "a7": 8, + }, + { + "a0": INT64, + "a1": INT64, + "a2": INT64, + "a3": INT64, + "a4": INT64, + "a5": INT64, + "a6": INT64, + "a7": INT64, + }, ), ] ) @@ -549,8 +603,26 @@ def test_executemany_insert_batch_aborted(self): [ ( """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3), (@a4, @a5, @a6, @a7)""", - {"a0": 1, "a1": 2, "a2": 3, "a3": 4, "a4": 5, "a5": 6, "a6": 7, "a7": 8}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64, "a4": INT64, "a5": INT64, "a6": INT64, "a7": INT64}, + { + "a0": 1, + "a1": 2, + "a2": 3, + "a3": 4, + "a4": 5, + "a5": 6, + "a6": 7, + "a7": 8, + }, + { + "a0": INT64, + "a1": INT64, + "a2": INT64, + "a3": INT64, + "a4": INT64, + "a5": INT64, + "a6": INT64, + "a7": INT64, + }, ), ] ) @@ -561,10 +633,28 @@ def test_executemany_insert_batch_aborted(self): [ ( """INSERT INTO table (col1, "col2", `col3`, `"col4"`) VALUES (@a0, @a1, @a2, @a3), (@a4, @a5, @a6, @a7)""", - {"a0": 1, "a1": 2, "a2": 3, "a3": 4, "a4": 5, "a5": 6, "a6": 7, "a7": 8}, - {"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64, "a4": INT64, "a5": INT64, "a6": INT64, "a7": INT64}, + { + "a0": 1, + "a1": 2, + "a2": 3, + "a3": 4, + "a4": 5, + "a5": 6, + "a6": 7, + "a7": 8, + }, + { + "a0": INT64, + "a1": INT64, + "a2": INT64, + "a3": INT64, + "a4": INT64, + "a5": INT64, + "a6": INT64, + "a7": INT64, + }, ), - ] + ], ) self.assertIsInstance(connection._statements[0][1], ResultsChecksum) diff --git a/tests/unit/spanner_dbapi/test_parse_utils.py b/tests/unit/spanner_dbapi/test_parse_utils.py index e454d962dc..ed826fcf81 100644 --- a/tests/unit/spanner_dbapi/test_parse_utils.py +++ b/tests/unit/spanner_dbapi/test_parse_utils.py @@ -111,16 +111,14 @@ def test_parse_insert(self): [ ( "INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s), (%s, %s, %s), (%s, %s, %s)", - (1, 2, 3, 4, 5, 6, 7, 8, 9) + (1, 2, 3, 4, 5, 6, 7, 8, 9), ), ], ), ( "INSERT INTO `no` (`yes`) VALUES (%s)", (1, 4, 5), - [ - ("INSERT INTO `no` (`yes`) VALUES (%s), (%s), (%s)", (1, 4, 5)) - ], + [("INSERT INTO `no` (`yes`) VALUES (%s), (%s), (%s)", (1, 4, 5))], ), ( "INSERT INTO T (f1, f2) VALUES (1, 2)", From 02bca0c6f0dcf8c7a99687e753c7a50ec256492c Mon Sep 17 00:00:00 2001 From: larkee Date: Wed, 10 Nov 2021 11:53:58 +1100 Subject: [PATCH 10/13] refactor: stop adding WHERE clause --- google/cloud/spanner_dbapi/cursor.py | 4 ---- google/cloud/spanner_dbapi/parse_utils.py | 14 ------------- tests/unit/spanner_dbapi/test_parse_utils.py | 21 -------------------- 3 files changed, 39 deletions(-) diff --git a/google/cloud/spanner_dbapi/cursor.py b/google/cloud/spanner_dbapi/cursor.py index e055d0fc98..3336ae9693 100644 --- a/google/cloud/spanner_dbapi/cursor.py +++ b/google/cloud/spanner_dbapi/cursor.py @@ -147,7 +147,6 @@ def close(self): self._is_closed = True def _do_execute_update(self, transaction, sql, params): - sql = parse_utils.ensure_where_clause(sql) sql, params = parse_utils.sql_pyformat_args_to_spanner(sql, params) result = transaction.execute_update( @@ -208,9 +207,6 @@ def execute(self, sql, args=None): self.connection.run_prior_DDL_statements() if not self.connection.autocommit: - if classification == parse_utils.STMT_UPDATING: - sql = parse_utils.ensure_where_clause(sql) - if classification != parse_utils.STMT_INSERT: sql, args = sql_pyformat_args_to_spanner(sql, args or None) diff --git a/google/cloud/spanner_dbapi/parse_utils.py b/google/cloud/spanner_dbapi/parse_utils.py index f3d094748c..5d2c20c18e 100644 --- a/google/cloud/spanner_dbapi/parse_utils.py +++ b/google/cloud/spanner_dbapi/parse_utils.py @@ -526,20 +526,6 @@ def get_param_types(params): return param_types -def ensure_where_clause(sql): - """ - Cloud Spanner requires a WHERE clause on UPDATE and DELETE statements. - Add a dummy WHERE clause if non detected. - - :type sql: str - :param sql: SQL code to check. - """ - if any(isinstance(token, sqlparse.sql.Where) for token in sqlparse.parse(sql)[0]): - return sql - - return sql + " WHERE 1=1" - - def escape_name(name): """ Apply backticks to the name that either contain '-' or diff --git a/tests/unit/spanner_dbapi/test_parse_utils.py b/tests/unit/spanner_dbapi/test_parse_utils.py index ed826fcf81..7683dd2b94 100644 --- a/tests/unit/spanner_dbapi/test_parse_utils.py +++ b/tests/unit/spanner_dbapi/test_parse_utils.py @@ -352,27 +352,6 @@ def test_get_param_types_none(self): self.assertEqual(get_param_types(None), None) - @unittest.skipIf(skip_condition, skip_message) - def test_ensure_where_clause(self): - from google.cloud.spanner_dbapi.parse_utils import ensure_where_clause - - cases = ( - "UPDATE a SET a.b=10 FROM articles a JOIN d c ON a.ai = c.ai WHERE c.ci = 1", - "UPDATE T SET A = 1 WHERE C1 = 1 AND C2 = 2", - "UPDATE T SET r=r*0.9 WHERE id IN (SELECT id FROM items WHERE r / w >= 1.3 AND q > 100)", - ) - err_cases = ( - "UPDATE (SELECT * FROM A JOIN c ON ai.id = c.id WHERE cl.ci = 1) SET d=5", - "DELETE * FROM TABLE", - ) - for sql in cases: - with self.subTest(sql=sql): - ensure_where_clause(sql) - - for sql in err_cases: - with self.subTest(sql=sql): - self.assertEqual(ensure_where_clause(sql), sql + " WHERE 1=1") - @unittest.skipIf(skip_condition, skip_message) def test_escape_name(self): from google.cloud.spanner_dbapi.parse_utils import escape_name From 307f60cd41dfd2b1a7d68ce691cf737e44c67c16 Mon Sep 17 00:00:00 2001 From: larkee Date: Thu, 11 Nov 2021 21:46:43 +1100 Subject: [PATCH 11/13] refactor: remove unnecessary parsing --- google/cloud/spanner_dbapi/_helpers.py | 21 +- google/cloud/spanner_dbapi/connection.py | 4 +- google/cloud/spanner_dbapi/parse_utils.py | 242 +------------------ tests/unit/spanner_dbapi/test__helpers.py | 46 ++-- tests/unit/spanner_dbapi/test_parse_utils.py | 186 -------------- 5 files changed, 22 insertions(+), 477 deletions(-) diff --git a/google/cloud/spanner_dbapi/_helpers.py b/google/cloud/spanner_dbapi/_helpers.py index 999c33c748..b869fc50d5 100644 --- a/google/cloud/spanner_dbapi/_helpers.py +++ b/google/cloud/spanner_dbapi/_helpers.py @@ -11,13 +11,9 @@ # 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.api_core.exceptions import Aborted -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, ABORTED SQL_LIST_TABLES = """ SELECT @@ -56,23 +52,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) - statements.append((sql, params, param_types)) - status, _ = transaction.batch_update(statements) - if status.code != OK: - if status.code == ABORTED: - raise Aborted(status.details) - raise OperationalError(status.message) +def _execute_insert_heterogenous(transaction, sql, params): + sql, params = sql_pyformat_args_to_spanner(sql, params) + transaction.execute_update(sql, params, get_param_types(params)) def handle_insert(connection, sql, params): - sql_params_list = parse_insert(sql, params) return connection.database.run_in_transaction( - _execute_insert_heterogenous, sql_params_list + _execute_insert_heterogenous, sql, params ) diff --git a/google/cloud/spanner_dbapi/connection.py b/google/cloud/spanner_dbapi/connection.py index 26bab9fe63..6c39693fda 100644 --- a/google/cloud/spanner_dbapi/connection.py +++ b/google/cloud/spanner_dbapi/connection.py @@ -24,7 +24,6 @@ from google.cloud.spanner_v1.snapshot import Snapshot from google.cloud.spanner_dbapi._helpers import _execute_insert_heterogenous -from google.cloud.spanner_dbapi._helpers import parse_insert from google.cloud.spanner_dbapi.checksum import _compare_checksums from google.cloud.spanner_dbapi.checksum import ResultsChecksum from google.cloud.spanner_dbapi.cursor import Cursor @@ -429,9 +428,8 @@ def run_statement(self, statement, retried=False): self._statements.append(statement) if statement.is_insert: - sql_params_list = parse_insert(statement.sql, statement.params) _execute_insert_heterogenous( - transaction, sql_params_list, + transaction, statement.sql, statement.params, ) return ( iter(()), diff --git a/google/cloud/spanner_dbapi/parse_utils.py b/google/cloud/spanner_dbapi/parse_utils.py index 5d2c20c18e..36d3f5b197 100644 --- a/google/cloud/spanner_dbapi/parse_utils.py +++ b/google/cloud/spanner_dbapi/parse_utils.py @@ -17,14 +17,12 @@ import datetime import decimal import re -from functools import reduce import sqlparse from google.cloud import spanner_v1 as spanner from google.cloud.spanner_v1 import JsonObject from .exceptions import Error, ProgrammingError -from .parser import parse_values from .types import DateStr, TimestampStr from .utils import sanitize_literals_for_upload @@ -185,6 +183,8 @@ def classify_stmt(query): :rtype: str :returns: The query type name. """ + query = sqlparse.format(query, strip_comments=True).strip() + if RE_DDL.match(query): return STMT_DDL @@ -199,244 +199,6 @@ def classify_stmt(query): return STMT_UPDATING -def parse_insert(insert_sql, params): - """ - Parse an INSERT statement and generate a list of tuples of the form: - [ - (SQL, params_per_row1), - (SQL, params_per_row2), - (SQL, params_per_row3), - ... - ] - - There are 4 variants of an INSERT statement: - a) INSERT INTO (columns...) VALUES (): no params - b) INSERT INTO
(columns...) SELECT_STMT: no params - c) INSERT INTO
(columns...) VALUES (%s,...): with params - d) INSERT INTO
(columns...) VALUES (%s,.....) with params and expressions - - Thus given each of the forms, it will produce a dictionary describing - how to upload the contents to Cloud Spanner: - Case a) - SQL: INSERT INTO T (f1, f2) VALUES (1, 2) - it produces: - [ - ('INSERT INTO T (f1, f2) VALUES (1, 2)', None), - ] - - Case b) - SQL: 'INSERT INTO T (s, c) SELECT st, zc FROM cus WHERE col IN (%s, %s)', - it produces: - [ - ('INSERT INTO T (s, c) SELECT st, zc FROM cus ORDER BY fn, ln', ('a', 'b')), - ] - - Case c) - SQL: INSERT INTO T (f1, f2) VALUES (%s, %s), (%s, %s) - Params: ['a', 'b', 'c', 'd'] - it produces: - [ - ('INSERT INTO T (f1, f2) VALUES (%s, %s)', ('a', 'b')), - ('INSERT INTO T (f1, f2) VALUES (%s, %s)', ('c', 'd')) - ] - - Case d) - SQL: INSERT INTO T (f1, f2) VALUES (%s, LOWER(%s)), (UPPER(%s), %s) - Params: ['a', 'b', 'c', 'd'] - it produces: - [ - ('INSERT INTO T (f1, f2) VALUES (%s, LOWER(%s))', ('a', 'b',)), - ('INSERT INTO T (f1, f2) VALUES (UPPER(%s), %s)', ('c', 'd',)) - ] - - :type insert_sql: str - :param insert_sql: A SQL insert request. - - :type params: list - :param params: A list of parameters. - - :rtype: dict - :returns: A list of (sql, param) tuples - """ # noqa - match = RE_INSERT.search(insert_sql) - - if not match: - raise ProgrammingError( - "Could not parse an INSERT statement from %s" % insert_sql - ) - - after_values_sql = RE_VALUES_TILL_END.findall(insert_sql) - if not after_values_sql: - # Case b) - insert_sql = sanitize_literals_for_upload(insert_sql) - return [(insert_sql, params)] - - if not params: - # Case a) perhaps? - # Check if any %s exists. - - # pyformat_str_count = after_values_sql.count("%s") - # if pyformat_str_count > 0: - # raise ProgrammingError( - # 'no params yet there are %d "%%s" tokens' % pyformat_str_count - # ) - for item in after_values_sql: - if item.count("%s") > 0: - raise ProgrammingError( - 'no params yet there are %d "%%s" tokens' % item.count("%s") - ) - - insert_sql = sanitize_literals_for_upload(insert_sql) - # Confirmed case of: - # SQL: INSERT INTO T (a1, a2) VALUES (1, 2) - # Params: None - return [(insert_sql, None)] - - values_str = after_values_sql[0] - _, values = parse_values(values_str) - - if values.homogenous(): - # Case c) - - columns = [mi.strip(" `") for mi in match.group("columns").split(",")] - values_pyformat = [str(arg) for arg in values.argv] - rows_list = rows_for_insert_or_update(columns, params, values_pyformat) - values_template = ", ".join([str(values.argv[0])] * len(rows_list)) - insert_sql = "INSERT INTO {} ({}) VALUES {}".format( - match.group("table_name"), match.group("columns"), values_template, - ) - insert_sql = sanitize_literals_for_upload(insert_sql) - flat_params = [] - for row in rows_list: - flat_params.extend(row) - return [(insert_sql, tuple(flat_params))] - - # Case d) - # insert_sql is of the form: - # INSERT INTO T(c1, c2) VALUES (%s, %s), (%s, LOWER(%s)) - - # Sanity check: - # length(all_args) == len(params) - args_len = reduce(lambda a, b: a + b, [len(arg) for arg in values.argv]) - if args_len != len(params): - raise ProgrammingError( - "Invalid length: VALUES(...) len: %d != len(params): %d" - % (args_len, len(params)) - ) - - trim_index = insert_sql.find(values_str) - before_values_sql = insert_sql[:trim_index] - - sql_param_tuples = [] - for token_arg in values.argv: - row_sql = before_values_sql + " VALUES %s" % token_arg - row_sql = sanitize_literals_for_upload(row_sql) - row_params, params = ( - tuple(params[0 : len(token_arg)]), - params[len(token_arg) :], - ) - sql_param_tuples.append((row_sql, row_params)) - - return sql_param_tuples - - -def rows_for_insert_or_update(columns, params, pyformat_args=None): - """ - Create a tupled list of params to be used as a single value per - value that inserted from a statement such as - SQL: 'INSERT INTO t (f1, f2, f3) VALUES (%s, %s, %s), (%s, %s, %s), (%s, %s, %s)' - Params A: [(1, 2, 3), (4, 5, 6), (7, 8, 9)] - Params B: [1, 2, 3, 4, 5, 6, 7, 8, 9] - - We'll have to convert both params types into: - Params: [(1, 2, 3,), (4, 5, 6,), (7, 8, 9,)] - - :type columns: list - :param columns: A list of the columns of the table. - - :type params: list - :param params: A list of parameters. - - :rtype: list - :returns: A properly restructured list of the parameters. - """ # noqa - if not pyformat_args: - # This is the case where we have for example: - # SQL: 'INSERT INTO t (f1, f2, f3)' - # Params A: [(1, 2, 3), (4, 5, 6), (7, 8, 9)] - # Params B: [1, 2, 3, 4, 5, 6, 7, 8, 9] - # - # We'll have to convert both params types into: - # [(1, 2, 3,), (4, 5, 6,), (7, 8, 9,)] - contains_all_list_or_tuples = True - for param in params: - if not (isinstance(param, list) or isinstance(param, tuple)): - contains_all_list_or_tuples = False - break - - if contains_all_list_or_tuples: - # The case with Params A: [(1, 2, 3), (4, 5, 6)] - # Ensure that each param's length == len(columns) - columns_len = len(columns) - for param in params: - if columns_len != len(param): - raise Error( - "\nlen(`%s`)=%d\n!=\ncolum_len(`%s`)=%d" - % (param, len(param), columns, columns_len) - ) - return params - else: - # The case with Params B: [1, 2, 3] - # Insert statements' params are only passed as tuples or lists, - # yet for do_execute_update, we've got to pass in list of list. - # https://googleapis.dev/python/spanner/latest/transaction-api.html\ - # #google.cloud.spanner_v1.transaction.Transaction.insert - n_stride = len(columns) - else: - # This is the case where we have for example: - # SQL: 'INSERT INTO t (f1, f2, f3) VALUES (%s, %s, %s), - # (%s, %s, %s), (%s, %s, %s)' - # Params: [1, 2, 3, 4, 5, 6, 7, 8, 9] - # which should become - # Columns: (f1, f2, f3) - # new_params: [(1, 2, 3,), (4, 5, 6,), (7, 8, 9,)] - - # Sanity check 1: all the pyformat_values should have the exact same - # length. - first, rest = pyformat_args[0], pyformat_args[1:] - n_stride = first.count("%s") - for pyfmt_value in rest: - n = pyfmt_value.count("%s") - if n_stride != n: - raise Error( - "\nlen(`%s`)=%d\n!=\nlen(`%s`)=%d" - % (first, n_stride, pyfmt_value, n) - ) - - # Sanity check 2: len(params) MUST be a multiple of n_stride aka - # len(count of %s). - # so that we can properly group for example: - # Given pyformat args: - # (%s, %s, %s) - # Params: - # [1, 2, 3, 4, 5, 6, 7, 8, 9] - # into - # [(1, 2, 3), (4, 5, 6), (7, 8, 9)] - if (len(params) % n_stride) != 0: - raise ProgrammingError( - "Invalid length: len(params)=%d MUST be a multiple of " - "len(pyformat_args)=%d" % (len(params), n_stride) - ) - - # Now chop up the strides. - strides = [] - for step in range(0, len(params), n_stride): - stride = tuple(params[step : step + n_stride :]) - strides.append(stride) - - return strides - - def sql_pyformat_args_to_spanner(sql, params): """ Transform pyformat set SQL to named arguments for Cloud Spanner. diff --git a/tests/unit/spanner_dbapi/test__helpers.py b/tests/unit/spanner_dbapi/test__helpers.py index be7526a11a..eecdc8545b 100644 --- a/tests/unit/spanner_dbapi/test__helpers.py +++ b/tests/unit/spanner_dbapi/test__helpers.py @@ -21,8 +21,6 @@ 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) @@ -34,21 +32,16 @@ def test__execute_insert_heterogenous(self): "google.cloud.spanner_dbapi._helpers.get_param_types", return_value=None ) as mock_param_types: transaction = mock.MagicMock() - status = Status(code=OK) - transaction.batch_update = mock_batch = mock.MagicMock( - return_value=(status, 1) - ) - _helpers._execute_insert_heterogenous(transaction, [params]) + transaction.execute_update = mock_update = 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_batch.assert_called_once_with([(sql, None, None)]) + mock_update.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 + from google.api_core.exceptions import Unknown sql = "sql" params = (sql, None) @@ -60,17 +53,14 @@ def test__execute_insert_heterogenous_error(self): "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.execute_update = mock_update = mock.MagicMock(side_effect=Unknown("Unknown")) - with self.assertRaises(OperationalError): - _helpers._execute_insert_heterogenous(transaction, [params]) + with self.assertRaises(Unknown): + _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)]) + mock_update.assert_called_once_with(sql, None, None) def test_handle_insert(self): from google.cloud.spanner_dbapi import _helpers @@ -78,19 +68,13 @@ def test_handle_insert(self): connection = mock.MagicMock() connection.database.run_in_transaction = mock_run_in = mock.MagicMock() sql = "sql" - parts = mock.MagicMock() - with mock.patch( - "google.cloud.spanner_dbapi._helpers.parse_insert", return_value=parts - ): - parts.get = mock.MagicMock(return_value=True) - mock_run_in.return_value = 0 - result = _helpers.handle_insert(connection, sql, None) - self.assertEqual(result, 0) - - parts.get = mock.MagicMock(return_value=False) - mock_run_in.return_value = 1 - result = _helpers.handle_insert(connection, sql, None) - self.assertEqual(result, 1) + mock_run_in.return_value = 0 + result = _helpers.handle_insert(connection, sql, None) + self.assertEqual(result, 0) + + mock_run_in.return_value = 1 + result = _helpers.handle_insert(connection, sql, None) + self.assertEqual(result, 1) class TestColumnInfo(unittest.TestCase): diff --git a/tests/unit/spanner_dbapi/test_parse_utils.py b/tests/unit/spanner_dbapi/test_parse_utils.py index 7683dd2b94..9c51287412 100644 --- a/tests/unit/spanner_dbapi/test_parse_utils.py +++ b/tests/unit/spanner_dbapi/test_parse_utils.py @@ -61,176 +61,6 @@ def test_classify_stmt(self): for query, want_class in cases: self.assertEqual(classify_stmt(query), want_class) - @unittest.skipIf(skip_condition, skip_message) - def test_parse_insert(self): - from google.cloud.spanner_dbapi.parse_utils import parse_insert - from google.cloud.spanner_dbapi.exceptions import ProgrammingError - - with self.assertRaises(ProgrammingError): - parse_insert("bad-sql", None) - - cases = [ - ( - "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)", - [1, 2, 3, 4, 5, 6], - [ - ( - "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s), (%s, %s, %s)", - (1, 2, 3, 4, 5, 6), - ), - ], - ), - ( - "INSERT INTO django_migrations(app, name, applied) VALUES (%s, %s, %s)", - [1, 2, 3, 4, 5, 6], - [ - ( - "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s), (%s, %s, %s)", - (1, 2, 3, 4, 5, 6), - ), - ], - ), - ( - "INSERT INTO sales.addresses (street, city, state, zip_code) " - "SELECT street, city, state, zip_code FROM sales.customers" - "ORDER BY first_name, last_name", - None, - [ - ( - "INSERT INTO sales.addresses (street, city, state, zip_code) " - "SELECT street, city, state, zip_code FROM sales.customers" - "ORDER BY first_name, last_name", - None, - ) - ], - ), - ( - "INSERT INTO ap (n, ct, cn) " - "VALUES (%s, %s, %s), (%s, %s, %s), (%s, %s, %s),(%s, %s, %s)", - (1, 2, 3, 4, 5, 6, 7, 8, 9), - [ - ( - "INSERT INTO ap (n, ct, cn) VALUES (%s, %s, %s), (%s, %s, %s), (%s, %s, %s)", - (1, 2, 3, 4, 5, 6, 7, 8, 9), - ), - ], - ), - ( - "INSERT INTO `no` (`yes`) VALUES (%s)", - (1, 4, 5), - [("INSERT INTO `no` (`yes`) VALUES (%s), (%s), (%s)", (1, 4, 5))], - ), - ( - "INSERT INTO T (f1, f2) VALUES (1, 2)", - None, - [("INSERT INTO T (f1, f2) VALUES (1, 2)", None)], - ), - ( - "INSERT INTO `no` (`yes`, tiff) VALUES (%s, LOWER(%s)), (%s, %s), (%s, %s)", - (1, "FOO", 5, 10, 11, 29), - [ - ( - "INSERT INTO `no` (`yes`, tiff) VALUES (%s, LOWER(%s))", - (1, "FOO"), - ), - ("INSERT INTO `no` (`yes`, tiff) VALUES (%s, %s)", (5, 10)), - ("INSERT INTO `no` (`yes`, tiff) VALUES (%s, %s)", (11, 29)), - ], - ), - ] - - sql = "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s)" - with self.assertRaises(ProgrammingError): - parse_insert(sql, None) - - for sql, params, want in cases: - with self.subTest(sql=sql): - got = parse_insert(sql, params) - self.assertEqual(got, want, "Mismatch with parse_insert of `%s`" % sql) - - @unittest.skipIf(skip_condition, skip_message) - def test_parse_insert_invalid(self): - from google.cloud.spanner_dbapi import exceptions - from google.cloud.spanner_dbapi.parse_utils import parse_insert - - cases = [ - ( - "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s), (%s, %s, %s)", - [1, 2, 3, 4, 5, 6, 7], - "len\\(params\\)=7 MUST be a multiple of len\\(pyformat_args\\)=3", - ), - ( - "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s), (%s, %s, LOWER(%s))", - [1, 2, 3, 4, 5, 6, 7], - "Invalid length: VALUES\\(...\\) len: 6 != len\\(params\\): 7", - ), - ( - "INSERT INTO django_migrations (app, name, applied) VALUES (%s, %s, %s), (%s, %s, LOWER(%s)))", - [1, 2, 3, 4, 5, 6], - "VALUES: expected `,` got \\) in \\)", - ), - ] - - for sql, params, wantException in cases: - with self.subTest(sql=sql): - self.assertRaisesRegex( - exceptions.ProgrammingError, - wantException, - lambda: parse_insert(sql, params), - ) - - @unittest.skipIf(skip_condition, skip_message) - def test_rows_for_insert_or_update(self): - from google.cloud.spanner_dbapi.parse_utils import rows_for_insert_or_update - from google.cloud.spanner_dbapi.exceptions import Error - - with self.assertRaises(Error): - rows_for_insert_or_update([0], [[]]) - - with self.assertRaises(Error): - rows_for_insert_or_update([0], None, ["0", "%s"]) - - cases = [ - ( - ["id", "app", "name"], - [(5, "ap", "n"), (6, "bp", "m")], - None, - [(5, "ap", "n"), (6, "bp", "m")], - ), - ( - ["app", "name"], - [("ap", "n"), ("bp", "m")], - None, - [("ap", "n"), ("bp", "m")], - ), - ( - ["app", "name", "fn"], - ["ap", "n", "f1", "bp", "m", "f2", "cp", "o", "f3"], - ["(%s, %s, %s)", "(%s, %s, %s)", "(%s, %s, %s)"], - [("ap", "n", "f1"), ("bp", "m", "f2"), ("cp", "o", "f3")], - ), - ( - ["app", "name", "fn", "ln"], - [ - ("ap", "n", (45, "nested"), "ll"), - ("bp", "m", "f2", "mt"), - ("fp", "cp", "o", "f3"), - ], - None, - [ - ("ap", "n", (45, "nested"), "ll"), - ("bp", "m", "f2", "mt"), - ("fp", "cp", "o", "f3"), - ], - ), - (["app", "name", "fn"], ["ap", "n", "f1"], None, [("ap", "n", "f1")]), - ] - - for i, (columns, params, pyformat_args, want) in enumerate(cases): - with self.subTest(i=i): - got = rows_for_insert_or_update(columns, params, pyformat_args) - self.assertEqual(got, want) - @unittest.skipIf(skip_condition, skip_message) def test_sql_pyformat_args_to_spanner(self): from google.cloud.spanner_dbapi.parse_utils import sql_pyformat_args_to_spanner @@ -367,19 +197,3 @@ def test_escape_name(self): with self.subTest(name=name): got = escape_name(name) self.assertEqual(got, want) - - def test_insert_from_select(self): - """Check that INSERT from SELECT clause can be executed with arguments.""" - from google.cloud.spanner_dbapi.parse_utils import parse_insert - - SQL = """ -INSERT INTO tab_name (id, data) -SELECT tab_name.id + %s AS anon_1, tab_name.data -FROM tab_name -WHERE tab_name.data IN (%s, %s) -""" - ARGS = [5, "data2", "data3"] - - self.assertEqual( - parse_insert(SQL, ARGS), [(SQL, ARGS)], - ) From 8d771a14f8297aa8df03cb0a068072812608e390 Mon Sep 17 00:00:00 2001 From: larkee Date: Mon, 15 Nov 2021 15:49:10 +1100 Subject: [PATCH 12/13] refactor: revert executemany changes --- google/cloud/spanner_dbapi/cursor.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/google/cloud/spanner_dbapi/cursor.py b/google/cloud/spanner_dbapi/cursor.py index 3336ae9693..5761c5fef0 100644 --- a/google/cloud/spanner_dbapi/cursor.py +++ b/google/cloud/spanner_dbapi/cursor.py @@ -269,12 +269,6 @@ def executemany(self, operation, seq_of_params): if classification in (parse_utils.STMT_INSERT, parse_utils.STMT_UPDATING): statements = [] - if classification == parse_utils.STMT_INSERT: - flat_params = [] - for params in seq_of_params: - flat_params.extend(params) - operation, params = parse_utils.parse_insert(operation, flat_params)[0] - seq_of_params = [params] for params in seq_of_params: sql, params = parse_utils.sql_pyformat_args_to_spanner( operation, params @@ -282,12 +276,9 @@ def executemany(self, operation, seq_of_params): statements.append((sql, params, get_param_types(params))) if self.connection.autocommit: - if classification == parse_utils.STMT_INSERT: - self.execute(operation, flat_params) - else: - self.connection.database.run_in_transaction( - self._do_batch_update, statements, many_result_set - ) + self.connection.database.run_in_transaction( + self._do_batch_update, statements, many_result_set + ) else: retried = False while True: From 6d479c0d3330c04224c6f841ff723411bf50b515 Mon Sep 17 00:00:00 2001 From: larkee Date: Tue, 16 Nov 2021 21:59:34 +1100 Subject: [PATCH 13/13] refactor: revert where clause removal --- google/cloud/spanner_dbapi/_helpers.py | 2 ++ google/cloud/spanner_dbapi/cursor.py | 4 ++-- google/cloud/spanner_dbapi/parse_utils.py | 14 +++++++++++++ tests/unit/spanner_dbapi/test_parse_utils.py | 21 ++++++++++++++++++++ 4 files changed, 39 insertions(+), 2 deletions(-) diff --git a/google/cloud/spanner_dbapi/_helpers.py b/google/cloud/spanner_dbapi/_helpers.py index b869fc50d5..1fdd221080 100644 --- a/google/cloud/spanner_dbapi/_helpers.py +++ b/google/cloud/spanner_dbapi/_helpers.py @@ -11,10 +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.parse_utils import get_param_types from google.cloud.spanner_dbapi.parse_utils import sql_pyformat_args_to_spanner from google.cloud.spanner_v1 import param_types + SQL_LIST_TABLES = """ SELECT t.table_name diff --git a/google/cloud/spanner_dbapi/cursor.py b/google/cloud/spanner_dbapi/cursor.py index 5761c5fef0..de6912f880 100644 --- a/google/cloud/spanner_dbapi/cursor.py +++ b/google/cloud/spanner_dbapi/cursor.py @@ -147,6 +147,7 @@ def close(self): self._is_closed = True def _do_execute_update(self, transaction, sql, params): + sql = parse_utils.ensure_where_clause(sql) sql, params = parse_utils.sql_pyformat_args_to_spanner(sql, params) result = transaction.execute_update( @@ -207,8 +208,7 @@ def execute(self, sql, args=None): self.connection.run_prior_DDL_statements() if not self.connection.autocommit: - if classification != parse_utils.STMT_INSERT: - sql, args = sql_pyformat_args_to_spanner(sql, args or None) + sql, args = sql_pyformat_args_to_spanner(sql, args or None) statement = Statement( sql, diff --git a/google/cloud/spanner_dbapi/parse_utils.py b/google/cloud/spanner_dbapi/parse_utils.py index 36d3f5b197..199991bc5b 100644 --- a/google/cloud/spanner_dbapi/parse_utils.py +++ b/google/cloud/spanner_dbapi/parse_utils.py @@ -288,6 +288,20 @@ def get_param_types(params): return param_types +def ensure_where_clause(sql): + """ + Cloud Spanner requires a WHERE clause on UPDATE and DELETE statements. + Add a dummy WHERE clause if non detected. + + :type sql: str + :param sql: SQL code to check. + """ + if any(isinstance(token, sqlparse.sql.Where) for token in sqlparse.parse(sql)[0]): + return sql + + return sql + " WHERE 1=1" + + def escape_name(name): """ Apply backticks to the name that either contain '-' or diff --git a/tests/unit/spanner_dbapi/test_parse_utils.py b/tests/unit/spanner_dbapi/test_parse_utils.py index 9c51287412..511ad838cf 100644 --- a/tests/unit/spanner_dbapi/test_parse_utils.py +++ b/tests/unit/spanner_dbapi/test_parse_utils.py @@ -182,6 +182,27 @@ def test_get_param_types_none(self): self.assertEqual(get_param_types(None), None) + @unittest.skipIf(skip_condition, skip_message) + def test_ensure_where_clause(self): + from google.cloud.spanner_dbapi.parse_utils import ensure_where_clause + + cases = ( + "UPDATE a SET a.b=10 FROM articles a JOIN d c ON a.ai = c.ai WHERE c.ci = 1", + "UPDATE T SET A = 1 WHERE C1 = 1 AND C2 = 2", + "UPDATE T SET r=r*0.9 WHERE id IN (SELECT id FROM items WHERE r / w >= 1.3 AND q > 100)", + ) + err_cases = ( + "UPDATE (SELECT * FROM A JOIN c ON ai.id = c.id WHERE cl.ci = 1) SET d=5", + "DELETE * FROM TABLE", + ) + for sql in cases: + with self.subTest(sql=sql): + ensure_where_clause(sql) + + for sql in err_cases: + with self.subTest(sql=sql): + self.assertEqual(ensure_where_clause(sql), sql + " WHERE 1=1") + @unittest.skipIf(skip_condition, skip_message) def test_escape_name(self): from google.cloud.spanner_dbapi.parse_utils import escape_name