diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py index dc12d1b84609..73581f456be8 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py @@ -15,6 +15,7 @@ """Database cursor for Google Cloud Spanner DB API.""" from collections import namedtuple +from dataclasses import replace import sqlparse from google.api_core.exceptions import ( @@ -60,6 +61,22 @@ ColumnDetails = namedtuple("column_details", ["null_ok", "spanner_type"]) +def _ensure_where_clause(parsed_statement: ParsedStatement) -> ParsedStatement: + """Ensure a dummy WHERE clause is present on UPDATE/DELETE statements. + + TODO: Remove this once we no longer automatically add a dummy WHERE clause + to DML statements. + """ + if parsed_statement.statement_type == StatementType.UPDATE: + sql = parse_utils.ensure_where_clause(parsed_statement.statement.sql) + if sql != parsed_statement.statement.sql: + return replace( + parsed_statement, + statement=replace(parsed_statement.statement, sql=sql), + ) + return parsed_statement + + def check_not_closed(function): """`Cursor` class methods decorator. @@ -337,6 +354,7 @@ def _execute(self, sql, args=None, call_from_execute_many=False): else: self._itr = PeekIterator(self._result_set) elif self.connection._batch_mode == BatchMode.DML: + self._parsed_statement = _ensure_where_clause(self._parsed_statement) self.connection.execute_batch_dml_statement(self._parsed_statement) elif ( self._parsed_statement.statement_type == StatementType.QUERY @@ -353,12 +371,13 @@ def _execute(self, sql, args=None, call_from_execute_many=False): if not self.connection._client_transaction_started: self.connection.run_prior_DDL_statements() elif ( - self.connection.autocommit_dml_mode + not self.connection._client_transaction_started + and self.connection.autocommit_dml_mode is AutocommitDmlMode.PARTITIONED_NON_ATOMIC ): self._row_count = self.connection.database.execute_partitioned_dml( - sql, - params=args, + self._parsed_statement.statement.sql, + params=self._parsed_statement.statement.params, param_types=self._parsed_statement.statement.param_types, request_options=self.request_options, ) @@ -401,6 +420,7 @@ def _execute_in_rw_transaction(self): # For every other operation, we've got to ensure that # any prior DDL statements were run. self.connection.run_prior_DDL_statements() + self._parsed_statement = _ensure_where_clause(self._parsed_statement) statement = self._parsed_statement.statement if self.connection._client_transaction_started: while True: @@ -461,6 +481,9 @@ def executemany(self, operation, seq_of_params): + ", with executemany() method is not allowed." ) + self._parsed_statement = _ensure_where_clause(self._parsed_statement) + sql = self._parsed_statement.statement.sql + # For every operation, we've got to ensure that any prior DDL # statements were run. self.connection.run_prior_DDL_statements() @@ -473,10 +496,12 @@ def executemany(self, operation, seq_of_params): ): statements = [] for params in seq_of_params: - sql, params = parse_utils.sql_pyformat_args_to_spanner( - operation, params + statement_sql, params = parse_utils.sql_pyformat_args_to_spanner( + sql, params + ) + statements.append( + Statement(statement_sql, params, get_param_types(params)) ) - statements.append(Statement(sql, params, get_param_types(params))) many_result_set = batch_dml_executor.run_batch_dml(self, statements) else: many_result_set = StreamedManyResultSets() diff --git a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/parse_utils.py b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/parse_utils.py index 553aa1e1e052..6f4d841543c6 100644 --- a/packages/google-cloud-spanner/google/cloud/spanner_dbapi/parse_utils.py +++ b/packages/google-cloud-spanner/google/cloud/spanner_dbapi/parse_utils.py @@ -259,12 +259,11 @@ def classify_statement(query, args=None): args, get_param_types(args or None), ) - statement_type = _get_statement_type(statement) + statement_type = _get_statement_type(query) return ParsedStatement(statement_type, statement) -def _get_statement_type(statement): - query = statement.sql +def _get_statement_type(query): if RE_DDL.match(query): return StatementType.DDL if RE_IS_INSERT.match(query): @@ -275,9 +274,6 @@ def _get_statement_type(statement): return StatementType.QUERY if RE_IS_UPDATE.match(query) or RE_IS_DELETE.match(query): - # TODO: Remove this? It makes more sense to have this in SQLAlchemy and - # Django than here. - statement.sql = ensure_where_clause(query) return StatementType.UPDATE return StatementType.UNKNOWN diff --git a/packages/google-cloud-spanner/tests/mockserver_tests/test_basics.py b/packages/google-cloud-spanner/tests/mockserver_tests/test_basics.py index bc2120a406f9..af3b94f57fd0 100644 --- a/packages/google-cloud-spanner/tests/mockserver_tests/test_basics.py +++ b/packages/google-cloud-spanner/tests/mockserver_tests/test_basics.py @@ -112,6 +112,110 @@ def test_dbapi_partitioned_dml(self): TransactionOptions(dict(partitioned_dml={})), begin_request.options ) + def test_dbapi_partitioned_dml_with_parameters(self): + sql = "UPDATE singers SET foo = %s WHERE active = %s" + expected_sql = "UPDATE singers SET foo = @a0 WHERE active = @a1" + add_update_count(expected_sql, 42, AutocommitDmlMode.PARTITIONED_NON_ATOMIC) + connection = Connection(self.instance, self.database) + connection.autocommit = True + connection.set_autocommit_dml_mode(AutocommitDmlMode.PARTITIONED_NON_ATOMIC) + with connection.cursor() as cursor: + cursor.execute(sql, ["bar", True]) + self.assertEqual(42, cursor.rowcount) + requests = self.spanner_service.requests + self.assert_requests_sequence( + requests, + [BeginTransactionRequest, ExecuteSqlRequest], + TransactionType.PARTITIONED, + allow_multiple_batch_create=True, + ) + execute_requests = [ + request for request in requests if isinstance(request, ExecuteSqlRequest) + ] + self.assertEqual(len(execute_requests), 1) + self.assertEqual(execute_requests[0].sql, expected_sql) + self.assertEqual(execute_requests[0].params["a0"], "bar") + self.assertEqual(execute_requests[0].params["a1"], True) + self.assertEqual(execute_requests[0].param_types["a0"].code, TypeCode.STRING) + self.assertEqual(execute_requests[0].param_types["a1"].code, TypeCode.BOOL) + + def test_dbapi_partitioned_dml_does_not_add_dummy_where(self): + sql = "UPDATE singers SET foo = %s" + expected_sql = "UPDATE singers SET foo = @a0" + add_update_count(expected_sql, 15, AutocommitDmlMode.PARTITIONED_NON_ATOMIC) + connection = Connection(self.instance, self.database) + connection.autocommit = True + connection.set_autocommit_dml_mode(AutocommitDmlMode.PARTITIONED_NON_ATOMIC) + with connection.cursor() as cursor: + cursor.execute(sql, ["baz"]) + self.assertEqual(15, cursor.rowcount) + requests = self.spanner_service.requests + execute_requests = [ + request for request in requests if isinstance(request, ExecuteSqlRequest) + ] + self.assertEqual(len(execute_requests), 1) + self.assertEqual(execute_requests[0].sql, expected_sql) + self.assertNotIn("WHERE", execute_requests[0].sql) + + def test_dbapi_transactional_dml_adds_dummy_where(self): + sql = "UPDATE singers SET foo = %s" + expected_sql = "UPDATE singers SET foo = @a0 WHERE 1=1" + add_update_count(expected_sql, 1) + connection = Connection(self.instance, self.database) + connection.autocommit = True + with connection.cursor() as cursor: + cursor.execute(sql, ["baz"]) + self.assertEqual(1, cursor.rowcount) + requests = self.spanner_service.requests + execute_requests = [ + request for request in requests if isinstance(request, ExecuteSqlRequest) + ] + self.assertEqual(len(execute_requests), 1) + self.assertEqual(execute_requests[0].sql, expected_sql) + + def test_dbapi_executemany_unqualified_update_adds_dummy_where(self): + sql = "UPDATE singers SET foo = %s" + expected_sql = "UPDATE singers SET foo = @a0 WHERE 1=1" + add_update_count(expected_sql, 1) + connection = Connection(self.instance, self.database) + connection.autocommit = True + with connection.cursor() as cursor: + cursor.executemany(sql, [("bar",), ("baz",)]) + self.assertEqual(2, cursor.rowcount) + requests = self.spanner_service.requests + batch_requests = [ + request + for request in requests + if isinstance(request, ExecuteBatchDmlRequest) + ] + self.assertEqual(len(batch_requests), 1) + self.assertEqual(len(batch_requests[0].statements), 2) + self.assertEqual(batch_requests[0].statements[0].sql, expected_sql) + self.assertEqual(batch_requests[0].statements[1].sql, expected_sql) + + def test_dbapi_batch_dml_unqualified_update_adds_dummy_where(self): + sql = "UPDATE singers SET foo = 'bar'" + expected_sql = "UPDATE singers SET foo = 'bar' WHERE 1=1" + add_update_count(expected_sql, 1) + connection = Connection(self.instance, self.database) + connection.autocommit = True + with connection.cursor() as cursor: + cursor.execute("start batch dml") + cursor.execute(sql) + cursor.execute(sql) + cursor.execute("run batch") + self.assertEqual(2, cursor.rowcount) + requests = self.spanner_service.requests + batch_requests = [ + request + for request in requests + if isinstance(request, ExecuteBatchDmlRequest) + ] + self.assertEqual(len(batch_requests), 1) + self.assertEqual(len(batch_requests[0].statements), 2) + self.assertEqual(batch_requests[0].statements[0].sql, expected_sql) + self.assertEqual(batch_requests[0].statements[1].sql, expected_sql) + def test_batch_create_sessions_unavailable(self): add_select1_result() add_error(SpannerServicer.BatchCreateSessions.__name__, unavailable_status()) diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_cursor.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_cursor.py index 8d7b33fe4a98..567c617092c5 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_cursor.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_cursor.py @@ -25,6 +25,7 @@ from google.cloud import spanner from google.cloud.spanner_dbapi.connection import connect from google.cloud.spanner_dbapi.parsed_statement import ( + AutocommitDmlMode, ParsedStatement, Statement, StatementType, @@ -314,7 +315,7 @@ def test_execute_statement(self): mock_db.run_in_transaction = mock_run_in = mock.MagicMock() cursor.execute(sql="sql") mock_run_in.assert_called_once_with( - cursor._do_execute_update_in_autocommit, "sql", None + cursor._do_execute_update_in_autocommit, "sql WHERE 1=1", None ) def test_execute_statement_with_cursor_not_in_retry_mode(self): @@ -1644,3 +1645,271 @@ def test_execute_dml_then_ddl_then_rollback_discards_ddl(self, mock_peek_iterato cursor.execute(sql="SELECT 1") connection.commit() mock_database.update_ddl.assert_not_called() + + def test_execute_partitioned_dml_passes_statement_params(self): + from google.cloud.spanner_v1 import Type, TypeCode + + mock_database = mock.MagicMock() + mock_database.execute_partitioned_dml.return_value = 10 + connection = self._make_connection(self.INSTANCE, mock_database) + connection.autocommit = True + connection.set_autocommit_dml_mode(AutocommitDmlMode.PARTITIONED_NON_ATOMIC) + cursor = self._make_one(connection) + + # 1. Pyformat positional parameters (list) + cursor.execute("DELETE FROM users WHERE id = %s", [42]) + self.assertEqual(cursor.rowcount, 10) + mock_database.execute_partitioned_dml.assert_called_once_with( + "DELETE FROM users WHERE id = @a0", + params={"a0": 42}, + param_types={"a0": Type(code=TypeCode.INT64)}, + request_options=None, + ) + + # 2. Pyformat positional parameters (tuple) + mock_database.reset_mock() + cursor.execute("DELETE FROM users WHERE id = %s", (42,)) + self.assertEqual(cursor.rowcount, 10) + mock_database.execute_partitioned_dml.assert_called_once_with( + "DELETE FROM users WHERE id = @a0", + params={"a0": 42}, + param_types={"a0": Type(code=TypeCode.INT64)}, + request_options=None, + ) + + # 3. Named Spanner parameters + mock_database.reset_mock() + mock_database.execute_partitioned_dml.return_value = 5 + cursor.execute( + "UPDATE users SET active = @status WHERE id = @user_id", + {"status": False, "user_id": 100}, + ) + self.assertEqual(cursor.rowcount, 5) + mock_database.execute_partitioned_dml.assert_called_once_with( + "UPDATE users SET active = @status WHERE id = @user_id", + params={"status": False, "user_id": 100}, + param_types={ + "status": Type(code=TypeCode.BOOL), + "user_id": Type(code=TypeCode.INT64), + }, + request_options=None, + ) + + # 4. Pyformat named parameters (%(key)s) + mock_database.reset_mock() + mock_database.execute_partitioned_dml.return_value = 7 + cursor.execute( + "UPDATE users SET active = %(status)s WHERE id = %(user_id)s", + {"status": True, "user_id": 200}, + ) + self.assertEqual(cursor.rowcount, 7) + mock_database.execute_partitioned_dml.assert_called_once_with( + "UPDATE users SET active = @a0 WHERE id = @a1", + params={"a0": "True", "a1": "200"}, + param_types={ + "a0": Type(code=TypeCode.STRING), + "a1": Type(code=TypeCode.STRING), + }, + request_options=None, + ) + + # 5. Parameterless query with explicit WHERE 1=1 + mock_database.reset_mock() + mock_database.execute_partitioned_dml.return_value = 100 + cursor.execute("DELETE FROM users WHERE 1=1") + self.assertEqual(cursor.rowcount, 100) + mock_database.execute_partitioned_dml.assert_called_once_with( + "DELETE FROM users WHERE 1=1", + params=None, + param_types=None, + request_options=None, + ) + + def test_execute_partitioned_dml_does_not_add_dummy_where(self): + """Partitioned DML statements must NOT automatically receive a dummy WHERE 1=1.""" + from google.cloud.spanner_v1 import Type, TypeCode + + mock_database = mock.MagicMock() + mock_database.execute_partitioned_dml.return_value = 50 + connection = self._make_connection(self.INSTANCE, mock_database) + connection.autocommit = True + connection.set_autocommit_dml_mode(AutocommitDmlMode.PARTITIONED_NON_ATOMIC) + cursor = self._make_one(connection) + + # Unqualified UPDATE with parameters must NOT have WHERE 1=1 appended + cursor.execute("UPDATE users SET active = %s", [True]) + self.assertEqual(cursor.rowcount, 50) + mock_database.execute_partitioned_dml.assert_called_once_with( + "UPDATE users SET active = @a0", + params={"a0": True}, + param_types={"a0": Type(code=TypeCode.BOOL)}, + request_options=None, + ) + + # Unqualified DELETE must NOT have WHERE 1=1 appended + mock_database.reset_mock() + cursor.execute("DELETE FROM users") + self.assertEqual(cursor.rowcount, 50) + mock_database.execute_partitioned_dml.assert_called_once_with( + "DELETE FROM users", + params=None, + param_types=None, + request_options=None, + ) + + def test_execute_partitioned_dml_skipped_when_in_transaction(self): + """Statements executed in an explicit transaction must not run as Partitioned DML.""" + mock_database = mock.MagicMock() + connection = self._make_connection(self.INSTANCE, mock_database) + connection.autocommit = True + connection.set_autocommit_dml_mode(AutocommitDmlMode.PARTITIONED_NON_ATOMIC) + # Transaction started at client side + connection.autocommit = False + cursor = self._make_one(connection) + connection.run_statement = mock.MagicMock() + + cursor.execute("UPDATE users SET active = true WHERE id = 1") + mock_database.execute_partitioned_dml.assert_not_called() + + def test_executemany_unqualified_update_adds_dummy_where(self): + """Unqualified updates/deletes in executemany should receive a dummy WHERE 1=1.""" + connection = self._make_connection(self.INSTANCE, mock.MagicMock()) + cursor = self._make_one(connection) + + with mock.patch( + "google.cloud.spanner_dbapi.batch_dml_executor.run_batch_dml" + ) as mock_batch: + cursor.executemany("UPDATE users SET active = %s", [(True,), (False,)]) + statements = mock_batch.call_args[0][1] + self.assertEqual(len(statements), 2) + self.assertEqual( + statements[0].sql, "UPDATE users SET active = @a0 WHERE 1=1" + ) + self.assertEqual( + statements[1].sql, "UPDATE users SET active = @a0 WHERE 1=1" + ) + + def test_execute_unqualified_update_in_transaction_adds_dummy_where(self): + """Unqualified UPDATE in an explicit transaction must receive a dummy WHERE 1=1.""" + connection = self._make_connection(self.INSTANCE, mock.MagicMock()) + connection.autocommit = False + cursor = self._make_one(connection) + connection.run_statement = mock.MagicMock() + + cursor.execute("UPDATE users SET active = true") + connection.run_statement.assert_called_once() + statement = connection.run_statement.call_args[0][0] + self.assertEqual(statement.sql, "UPDATE users SET active = true WHERE 1=1") + + def test_execute_batch_mode_dml_unqualified_update_adds_dummy_where(self): + """Unqualified UPDATE in batch DML mode must receive a dummy WHERE 1=1.""" + from google.cloud.spanner_dbapi.connection import BatchMode + + connection = self._make_connection(self.INSTANCE, mock.MagicMock()) + connection._batch_mode = BatchMode.DML + connection.execute_batch_dml_statement = mock.MagicMock() + cursor = self._make_one(connection) + + cursor.execute("UPDATE users SET active = true") + connection.execute_batch_dml_statement.assert_called_once() + parsed_statement = connection.execute_batch_dml_statement.call_args[0][0] + self.assertEqual( + parsed_statement.statement.sql, "UPDATE users SET active = true WHERE 1=1" + ) + + def test_transaction_retry_unqualified_dml_does_not_duplicate_where_clause(self): + """Transaction retry of an unqualified UPDATE must not add WHERE 1=1 a second time.""" + connection = self._make_connection(self.INSTANCE, mock.MagicMock()) + connection.autocommit = False + cursor = self._make_one(connection) + + update_1_sql = "UPDATE table_name SET col = 1" + update_2_sql = "UPDATE table_name SET col = 2" + + metadata_mock = mock.Mock() + metadata_mock.trailing_metadata.return_value = {} + + mock_result_set = mock.MagicMock() + mock_result_set.metadata.row_type = None + + connection.run_statement = mock.Mock( + side_effect=[ + mock_result_set, # update_1 (attempt 1 -> succeeds) + Aborted( + "Aborted", errors=[metadata_mock] + ), # update_2 (attempt 1 -> aborts) + mock_result_set, # update_1 (replayed by retry_transaction) + mock_result_set, # update_2 (attempt 2 -> succeeds) + ] + ) + + cursor.execute(sql=update_1_sql) + cursor.execute(sql=update_2_sql) + + self.assertEqual(connection.run_statement.call_count, 4) + # Attempt 1 of update_1 + statement_attempt_1 = connection.run_statement.call_args_list[0][0][0] + self.assertEqual( + statement_attempt_1.sql, "UPDATE table_name SET col = 1 WHERE 1=1" + ) + + # Retry replay of update_1 + statement_retry = connection.run_statement.call_args_list[2][0][0] + self.assertEqual(statement_retry.sql, "UPDATE table_name SET col = 1 WHERE 1=1") + self.assertNotIn("WHERE 1=1 WHERE 1=1", statement_retry.sql) + + def test_transaction_retry_unqualified_executemany_does_not_duplicate_where_clause( + self, + ): + """Transaction retry of an unqualified executemany must not add WHERE 1=1 a second time.""" + connection = self._make_connection(self.INSTANCE, mock.MagicMock()) + connection.autocommit = False + cursor = self._make_one(connection) + + operation = "UPDATE table_name SET col = %s" + params_seq = [(1,), (2,)] + + metadata_mock = mock.Mock() + metadata_mock.trailing_metadata.return_value = {} + + mock_result_set = mock.MagicMock() + + # Step 1: executemany succeeds + with mock.patch( + "google.cloud.spanner_dbapi.batch_dml_executor.run_batch_dml", + return_value=mock_result_set, + ) as mock_batch: + cursor.executemany(operation, params_seq) + first_statements = mock_batch.call_args[0][1] + self.assertEqual( + first_statements[0].sql, "UPDATE table_name SET col = @a0 WHERE 1=1" + ) + self.assertEqual( + first_statements[1].sql, "UPDATE table_name SET col = @a0 WHERE 1=1" + ) + + # Step 2: a subsequent execute aborts, triggering retry of executemany and the subsequent execute + connection.run_statement = mock.Mock( + side_effect=[ + Aborted("Aborted", errors=[metadata_mock]), # aborts + mock.MagicMock(), # retry succeeds + ] + ) + + with mock.patch( + "google.cloud.spanner_dbapi.batch_dml_executor.run_batch_dml", + return_value=mock_result_set, + ) as mock_batch_retry: + cursor.execute("UPDATE other_table SET col = 3") + + # Verify that executemany was replayed by retry_transaction + mock_batch_retry.assert_called_once() + retried_statements = mock_batch_retry.call_args[0][1] + self.assertEqual( + retried_statements[0].sql, "UPDATE table_name SET col = @a0 WHERE 1=1" + ) + self.assertEqual( + retried_statements[1].sql, "UPDATE table_name SET col = @a0 WHERE 1=1" + ) + self.assertNotIn("WHERE 1=1 WHERE 1=1", retried_statements[0].sql) + self.assertNotIn("WHERE 1=1 WHERE 1=1", retried_statements[1].sql) diff --git a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py index c876b0b7644e..4587e2f0fa83 100644 --- a/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py +++ b/packages/google-cloud-spanner/tests/unit/spanner_dbapi/test_parse_utils.py @@ -693,3 +693,26 @@ def test_classify_statement_preserves_unicode_and_predicates(self): "SELECT `tätä`, id FROM users WHERE id = @a0", ) self.assertEqual(parsed.statement.params, {"a0": 1}) + + def test_classify_statement_does_not_add_dummy_where(self): + """classify_statement must classify statements without mutating them or adding WHERE 1=1.""" + query = "UPDATE users SET active = true" + parsed = classify_statement(query) + self.assertEqual(parsed.statement.sql, query) + self.assertEqual(parsed.statement_type, StatementType.UPDATE) + + delete_query = "DELETE FROM users" + parsed = classify_statement(delete_query) + self.assertEqual(parsed.statement.sql, delete_query) + self.assertEqual(parsed.statement_type, StatementType.UPDATE) + + def test_ensure_where_clause_is_idempotent(self): + """ensure_where_clause must not append duplicate WHERE clauses if called repeatedly.""" + from google.cloud.spanner_dbapi.parse_utils import ensure_where_clause + + sql = "UPDATE users SET active = true" + with_where = ensure_where_clause(sql) + self.assertEqual(with_where, "UPDATE users SET active = true WHERE 1=1") + # Second invocation must be a no-op + self.assertEqual(ensure_where_clause(with_where), with_where) + self.assertNotIn("WHERE 1=1 WHERE 1=1", ensure_where_clause(with_where))