Skip to content

fix(spanner): support parameters and preserve safety in partitioned DML #19035

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
olavloite wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(spanner): support parameters and preserve safety in partitioned DML #19035

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 31 additions & 6 deletions packages/google-cloud-spanner/google/cloud/spanner_dbapi/cursor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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
Expand All @@ -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,
)
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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()
Expand All @@ -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()
Expand Down
8 changes: 2 additions & 6 deletions packages/google-cloud-spanner/google/cloud/spanner_dbapi/parse_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -255,12 +255,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):
Expand All @@ -271,9 +270,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
Expand Down
104 changes: 104 additions & 0 deletions packages/google-cloud-spanner/tests/mockserver_tests/test_basics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down
Loading
Loading