Skip to content

Commit 19ab6ef

Browse files
perf: add option for last_statement (#1313)
* perf: add option for last_statement Adds an option to indicate that a statement is the last statement in a read/write transaction. Setting this option allows Spanner to optimize the execution of the statement, and defer some validations until the Commit RPC that should follow directly after this statement. The last_statement option is automatically used by the dbapi driver when a statement is executed in autocommit mode. * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent 7a5afba commit 19ab6ef

File tree

7 files changed

+246
-8
lines changed

7 files changed

+246
-8
lines changed

google/cloud/spanner_dbapi/batch_dml_executor.py

+5-3
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,9 @@ def run_batch_dml(cursor: "Cursor", statements: List[Statement]):
8787
for statement in statements:
8888
statements_tuple.append(statement.get_tuple())
8989
if not connection._client_transaction_started:
90-
res = connection.database.run_in_transaction(_do_batch_update, statements_tuple)
90+
res = connection.database.run_in_transaction(
91+
_do_batch_update_autocommit, statements_tuple
92+
)
9193
many_result_set.add_iter(res)
9294
cursor._row_count = sum([max(val, 0) for val in res])
9395
else:
@@ -113,10 +115,10 @@ def run_batch_dml(cursor: "Cursor", statements: List[Statement]):
113115
connection._transaction_helper.retry_transaction()
114116

115117

116-
def _do_batch_update(transaction, statements):
118+
def _do_batch_update_autocommit(transaction, statements):
117119
from google.cloud.spanner_dbapi import OperationalError
118120

119-
status, res = transaction.batch_update(statements)
121+
status, res = transaction.batch_update(statements, last_statement=True)
120122
if status.code == ABORTED:
121123
raise Aborted(status.message)
122124
elif status.code != OK:

google/cloud/spanner_dbapi/cursor.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,10 @@ def _do_execute_update_in_autocommit(self, transaction, sql, params):
229229
self.connection._transaction = transaction
230230
self.connection._snapshot = None
231231
self._result_set = transaction.execute_sql(
232-
sql, params=params, param_types=get_param_types(params)
232+
sql,
233+
params=params,
234+
param_types=get_param_types(params),
235+
last_statement=True,
233236
)
234237
self._itr = PeekIterator(self._result_set)
235238
self._row_count = None

google/cloud/spanner_v1/snapshot.py

+15
Original file line numberDiff line numberDiff line change
@@ -389,6 +389,7 @@ def execute_sql(
389389
query_mode=None,
390390
query_options=None,
391391
request_options=None,
392+
last_statement=False,
392393
partition=None,
393394
retry=gapic_v1.method.DEFAULT,
394395
timeout=gapic_v1.method.DEFAULT,
@@ -432,6 +433,19 @@ def execute_sql(
432433
If a dict is provided, it must be of the same form as the protobuf
433434
message :class:`~google.cloud.spanner_v1.types.RequestOptions`.
434435
436+
:type last_statement: bool
437+
:param last_statement:
438+
If set to true, this option marks the end of the transaction. The
439+
transaction should be committed or aborted after this statement
440+
executes, and attempts to execute any other requests against this
441+
transaction (including reads and queries) will be rejected. Mixing
442+
mutations with statements that are marked as the last statement is
443+
not allowed.
444+
For DML statements, setting this option may cause some error
445+
reporting to be deferred until commit time (e.g. validation of
446+
unique constraints). Given this, successful execution of a DML
447+
statement should not be assumed until the transaction commits.
448+
435449
:type partition: bytes
436450
:param partition: (Optional) one of the partition tokens returned
437451
from :meth:`partition_query`.
@@ -536,6 +550,7 @@ def execute_sql(
536550
seqno=self._execute_sql_count,
537551
query_options=query_options,
538552
request_options=request_options,
553+
last_statement=last_statement,
539554
data_boost_enabled=data_boost_enabled,
540555
directed_read_options=directed_read_options,
541556
)

google/cloud/spanner_v1/transaction.py

+30
Original file line numberDiff line numberDiff line change
@@ -349,6 +349,7 @@ def execute_update(
349349
query_mode=None,
350350
query_options=None,
351351
request_options=None,
352+
last_statement=False,
352353
*,
353354
retry=gapic_v1.method.DEFAULT,
354355
timeout=gapic_v1.method.DEFAULT,
@@ -385,6 +386,19 @@ def execute_update(
385386
If a dict is provided, it must be of the same form as the protobuf
386387
message :class:`~google.cloud.spanner_v1.types.RequestOptions`.
387388
389+
:type last_statement: bool
390+
:param last_statement:
391+
If set to true, this option marks the end of the transaction. The
392+
transaction should be committed or aborted after this statement
393+
executes, and attempts to execute any other requests against this
394+
transaction (including reads and queries) will be rejected. Mixing
395+
mutations with statements that are marked as the last statement is
396+
not allowed.
397+
For DML statements, setting this option may cause some error
398+
reporting to be deferred until commit time (e.g. validation of
399+
unique constraints). Given this, successful execution of a DML
400+
statement should not be assumed until the transaction commits.
401+
388402
:type retry: :class:`~google.api_core.retry.Retry`
389403
:param retry: (Optional) The retry settings for this request.
390404
@@ -433,6 +447,7 @@ def execute_update(
433447
query_options=query_options,
434448
seqno=seqno,
435449
request_options=request_options,
450+
last_statement=last_statement,
436451
)
437452

438453
method = functools.partial(
@@ -478,6 +493,7 @@ def batch_update(
478493
self,
479494
statements,
480495
request_options=None,
496+
last_statement=False,
481497
*,
482498
retry=gapic_v1.method.DEFAULT,
483499
timeout=gapic_v1.method.DEFAULT,
@@ -502,6 +518,19 @@ def batch_update(
502518
If a dict is provided, it must be of the same form as the protobuf
503519
message :class:`~google.cloud.spanner_v1.types.RequestOptions`.
504520
521+
:type last_statement: bool
522+
:param last_statement:
523+
If set to true, this option marks the end of the transaction. The
524+
transaction should be committed or aborted after this statement
525+
executes, and attempts to execute any other requests against this
526+
transaction (including reads and queries) will be rejected. Mixing
527+
mutations with statements that are marked as the last statement is
528+
not allowed.
529+
For DML statements, setting this option may cause some error
530+
reporting to be deferred until commit time (e.g. validation of
531+
unique constraints). Given this, successful execution of a DML
532+
statement should not be assumed until the transaction commits.
533+
505534
:type retry: :class:`~google.api_core.retry.Retry`
506535
:param retry: (Optional) The retry settings for this request.
507536
@@ -558,6 +587,7 @@ def batch_update(
558587
statements=parsed,
559588
seqno=seqno,
560589
request_options=request_options,
590+
last_statements=last_statement,
561591
)
562592

563593
method = functools.partial(

tests/mockserver_tests/test_basics.py

+57
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,10 @@
2020
ExecuteSqlRequest,
2121
BeginTransactionRequest,
2222
TransactionOptions,
23+
ExecuteBatchDmlRequest,
24+
TypeCode,
2325
)
26+
from google.cloud.spanner_v1.transaction import Transaction
2427
from google.cloud.spanner_v1.testing.mock_spanner import SpannerServicer
2528

2629
from tests.mockserver_tests.mock_server_test_base import (
@@ -29,6 +32,7 @@
2932
add_update_count,
3033
add_error,
3134
unavailable_status,
35+
add_single_result,
3236
)
3337

3438

@@ -107,3 +111,56 @@ def test_execute_streaming_sql_unavailable(self):
107111
# The ExecuteStreamingSql call should be retried.
108112
self.assertTrue(isinstance(requests[1], ExecuteSqlRequest))
109113
self.assertTrue(isinstance(requests[2], ExecuteSqlRequest))
114+
115+
def test_last_statement_update(self):
116+
sql = "update my_table set my_col=1 where id=2"
117+
add_update_count(sql, 1)
118+
self.database.run_in_transaction(
119+
lambda transaction: transaction.execute_update(sql, last_statement=True)
120+
)
121+
requests = list(
122+
filter(
123+
lambda msg: isinstance(msg, ExecuteSqlRequest),
124+
self.spanner_service.requests,
125+
)
126+
)
127+
self.assertEqual(1, len(requests), msg=requests)
128+
self.assertTrue(requests[0].last_statement, requests[0])
129+
130+
def test_last_statement_batch_update(self):
131+
sql = "update my_table set my_col=1 where id=2"
132+
add_update_count(sql, 1)
133+
self.database.run_in_transaction(
134+
lambda transaction: transaction.batch_update(
135+
[sql, sql], last_statement=True
136+
)
137+
)
138+
requests = list(
139+
filter(
140+
lambda msg: isinstance(msg, ExecuteBatchDmlRequest),
141+
self.spanner_service.requests,
142+
)
143+
)
144+
self.assertEqual(1, len(requests), msg=requests)
145+
self.assertTrue(requests[0].last_statements, requests[0])
146+
147+
def test_last_statement_query(self):
148+
sql = "insert into my_table (value) values ('One') then return id"
149+
add_single_result(sql, "c", TypeCode.INT64, [("1",)])
150+
self.database.run_in_transaction(
151+
lambda transaction: _execute_query(transaction, sql)
152+
)
153+
requests = list(
154+
filter(
155+
lambda msg: isinstance(msg, ExecuteSqlRequest),
156+
self.spanner_service.requests,
157+
)
158+
)
159+
self.assertEqual(1, len(requests), msg=requests)
160+
self.assertTrue(requests[0].last_statement, requests[0])
161+
162+
163+
def _execute_query(transaction: Transaction, sql: str):
164+
rows = transaction.execute_sql(sql, last_statement=True)
165+
for _ in rows:
166+
pass
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
# Copyright 2025 Google LLC All rights reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from google.cloud.spanner_dbapi import Connection
16+
from google.cloud.spanner_v1 import (
17+
ExecuteSqlRequest,
18+
TypeCode,
19+
CommitRequest,
20+
ExecuteBatchDmlRequest,
21+
)
22+
from tests.mockserver_tests.mock_server_test_base import (
23+
MockServerTestBase,
24+
add_single_result,
25+
add_update_count,
26+
)
27+
28+
29+
class TestDbapiAutoCommit(MockServerTestBase):
30+
@classmethod
31+
def setup_class(cls):
32+
super().setup_class()
33+
add_single_result(
34+
"select name from singers", "name", TypeCode.STRING, [("Some Singer",)]
35+
)
36+
add_update_count("insert into singers (id, name) values (1, 'Some Singer')", 1)
37+
38+
def test_select_autocommit(self):
39+
connection = Connection(self.instance, self.database)
40+
connection.autocommit = True
41+
with connection.cursor() as cursor:
42+
cursor.execute("select name from singers")
43+
result_list = cursor.fetchall()
44+
for _ in result_list:
45+
pass
46+
requests = list(
47+
filter(
48+
lambda msg: isinstance(msg, ExecuteSqlRequest),
49+
self.spanner_service.requests,
50+
)
51+
)
52+
self.assertEqual(1, len(requests))
53+
self.assertFalse(requests[0].last_statement, requests[0])
54+
self.assertIsNotNone(requests[0].transaction, requests[0])
55+
self.assertIsNotNone(requests[0].transaction.single_use, requests[0])
56+
self.assertTrue(requests[0].transaction.single_use.read_only, requests[0])
57+
58+
def test_dml_autocommit(self):
59+
connection = Connection(self.instance, self.database)
60+
connection.autocommit = True
61+
with connection.cursor() as cursor:
62+
cursor.execute("insert into singers (id, name) values (1, 'Some Singer')")
63+
self.assertEqual(1, cursor.rowcount)
64+
requests = list(
65+
filter(
66+
lambda msg: isinstance(msg, ExecuteSqlRequest),
67+
self.spanner_service.requests,
68+
)
69+
)
70+
self.assertEqual(1, len(requests))
71+
self.assertTrue(requests[0].last_statement, requests[0])
72+
commit_requests = list(
73+
filter(
74+
lambda msg: isinstance(msg, CommitRequest),
75+
self.spanner_service.requests,
76+
)
77+
)
78+
self.assertEqual(1, len(commit_requests))
79+
80+
def test_executemany_autocommit(self):
81+
connection = Connection(self.instance, self.database)
82+
connection.autocommit = True
83+
with connection.cursor() as cursor:
84+
cursor.executemany(
85+
"insert into singers (id, name) values (1, 'Some Singer')", [(), ()]
86+
)
87+
self.assertEqual(2, cursor.rowcount)
88+
requests = list(
89+
filter(
90+
lambda msg: isinstance(msg, ExecuteBatchDmlRequest),
91+
self.spanner_service.requests,
92+
)
93+
)
94+
self.assertEqual(1, len(requests))
95+
self.assertTrue(requests[0].last_statements, requests[0])
96+
commit_requests = list(
97+
filter(
98+
lambda msg: isinstance(msg, CommitRequest),
99+
self.spanner_service.requests,
100+
)
101+
)
102+
self.assertEqual(1, len(commit_requests))
103+
104+
def test_batch_dml_autocommit(self):
105+
connection = Connection(self.instance, self.database)
106+
connection.autocommit = True
107+
with connection.cursor() as cursor:
108+
cursor.execute("start batch dml")
109+
cursor.execute("insert into singers (id, name) values (1, 'Some Singer')")
110+
cursor.execute("insert into singers (id, name) values (1, 'Some Singer')")
111+
cursor.execute("run batch")
112+
self.assertEqual(2, cursor.rowcount)
113+
requests = list(
114+
filter(
115+
lambda msg: isinstance(msg, ExecuteBatchDmlRequest),
116+
self.spanner_service.requests,
117+
)
118+
)
119+
self.assertEqual(1, len(requests))
120+
self.assertTrue(requests[0].last_statements, requests[0])
121+
commit_requests = list(
122+
filter(
123+
lambda msg: isinstance(msg, CommitRequest),
124+
self.spanner_service.requests,
125+
)
126+
)
127+
self.assertEqual(1, len(commit_requests))

tests/unit/spanner_dbapi/test_cursor.py

+8-4
Original file line numberDiff line numberDiff line change
@@ -148,7 +148,8 @@ def test_do_batch_update(self):
148148
("DELETE FROM table WHERE col1 = @a0", {"a0": 1}, {"a0": INT64}),
149149
("DELETE FROM table WHERE col1 = @a0", {"a0": 2}, {"a0": INT64}),
150150
("DELETE FROM table WHERE col1 = @a0", {"a0": 3}, {"a0": INT64}),
151-
]
151+
],
152+
last_statement=True,
152153
)
153154
self.assertEqual(cursor._row_count, 3)
154155

@@ -539,7 +540,8 @@ def test_executemany_delete_batch_autocommit(self):
539540
("DELETE FROM table WHERE col1 = @a0", {"a0": 1}, {"a0": INT64}),
540541
("DELETE FROM table WHERE col1 = @a0", {"a0": 2}, {"a0": INT64}),
541542
("DELETE FROM table WHERE col1 = @a0", {"a0": 3}, {"a0": INT64}),
542-
]
543+
],
544+
last_statement=True,
543545
)
544546

545547
def test_executemany_update_batch_autocommit(self):
@@ -582,7 +584,8 @@ def test_executemany_update_batch_autocommit(self):
582584
{"a0": 3, "a1": "c"},
583585
{"a0": INT64, "a1": STRING},
584586
),
585-
]
587+
],
588+
last_statement=True,
586589
)
587590

588591
def test_executemany_insert_batch_non_autocommit(self):
@@ -659,7 +662,8 @@ def test_executemany_insert_batch_autocommit(self):
659662
{"a0": 5, "a1": 6, "a2": 7, "a3": 8},
660663
{"a0": INT64, "a1": INT64, "a2": INT64, "a3": INT64},
661664
),
662-
]
665+
],
666+
last_statement=True,
663667
)
664668
transaction.commit.assert_called_once()
665669

0 commit comments

Comments
 (0)