Skip to content
This repository was archived by the owner on Mar 31, 2026. It is now read-only.

Commit a42423c

Browse files
authored
Merge branch 'main' into e2e_traces_
2 parents d16787a + 19ab6ef commit a42423c

File tree

7 files changed

+246
-8
lines changed

7 files changed

+246
-8
lines changed

google/cloud/spanner_dbapi/batch_dml_executor.py

Lines changed: 5 additions & 3 deletions
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

Lines changed: 4 additions & 1 deletion
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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -395,6 +395,7 @@ def execute_sql(
395395
query_mode=None,
396396
query_options=None,
397397
request_options=None,
398+
last_statement=False,
398399
partition=None,
399400
retry=gapic_v1.method.DEFAULT,
400401
timeout=gapic_v1.method.DEFAULT,
@@ -438,6 +439,19 @@ def execute_sql(
438439
If a dict is provided, it must be of the same form as the protobuf
439440
message :class:`~google.cloud.spanner_v1.types.RequestOptions`.
440441
442+
:type last_statement: bool
443+
:param last_statement:
444+
If set to true, this option marks the end of the transaction. The
445+
transaction should be committed or aborted after this statement
446+
executes, and attempts to execute any other requests against this
447+
transaction (including reads and queries) will be rejected. Mixing
448+
mutations with statements that are marked as the last statement is
449+
not allowed.
450+
For DML statements, setting this option may cause some error
451+
reporting to be deferred until commit time (e.g. validation of
452+
unique constraints). Given this, successful execution of a DML
453+
statement should not be assumed until the transaction commits.
454+
441455
:type partition: bytes
442456
:param partition: (Optional) one of the partition tokens returned
443457
from :meth:`partition_query`.
@@ -542,6 +556,7 @@ def execute_sql(
542556
seqno=self._execute_sql_count,
543557
query_options=query_options,
544558
request_options=request_options,
559+
last_statement=last_statement,
545560
data_boost_enabled=data_boost_enabled,
546561
directed_read_options=directed_read_options,
547562
)

google/cloud/spanner_v1/transaction.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -354,6 +354,7 @@ def execute_update(
354354
query_mode=None,
355355
query_options=None,
356356
request_options=None,
357+
last_statement=False,
357358
*,
358359
retry=gapic_v1.method.DEFAULT,
359360
timeout=gapic_v1.method.DEFAULT,
@@ -390,6 +391,19 @@ def execute_update(
390391
If a dict is provided, it must be of the same form as the protobuf
391392
message :class:`~google.cloud.spanner_v1.types.RequestOptions`.
392393
394+
:type last_statement: bool
395+
:param last_statement:
396+
If set to true, this option marks the end of the transaction. The
397+
transaction should be committed or aborted after this statement
398+
executes, and attempts to execute any other requests against this
399+
transaction (including reads and queries) will be rejected. Mixing
400+
mutations with statements that are marked as the last statement is
401+
not allowed.
402+
For DML statements, setting this option may cause some error
403+
reporting to be deferred until commit time (e.g. validation of
404+
unique constraints). Given this, successful execution of a DML
405+
statement should not be assumed until the transaction commits.
406+
393407
:type retry: :class:`~google.api_core.retry.Retry`
394408
:param retry: (Optional) The retry settings for this request.
395409
@@ -438,6 +452,7 @@ def execute_update(
438452
query_options=query_options,
439453
seqno=seqno,
440454
request_options=request_options,
455+
last_statement=last_statement,
441456
)
442457

443458
method = functools.partial(
@@ -485,6 +500,7 @@ def batch_update(
485500
self,
486501
statements,
487502
request_options=None,
503+
last_statement=False,
488504
*,
489505
retry=gapic_v1.method.DEFAULT,
490506
timeout=gapic_v1.method.DEFAULT,
@@ -509,6 +525,19 @@ def batch_update(
509525
If a dict is provided, it must be of the same form as the protobuf
510526
message :class:`~google.cloud.spanner_v1.types.RequestOptions`.
511527
528+
:type last_statement: bool
529+
:param last_statement:
530+
If set to true, this option marks the end of the transaction. The
531+
transaction should be committed or aborted after this statement
532+
executes, and attempts to execute any other requests against this
533+
transaction (including reads and queries) will be rejected. Mixing
534+
mutations with statements that are marked as the last statement is
535+
not allowed.
536+
For DML statements, setting this option may cause some error
537+
reporting to be deferred until commit time (e.g. validation of
538+
unique constraints). Given this, successful execution of a DML
539+
statement should not be assumed until the transaction commits.
540+
512541
:type retry: :class:`~google.api_core.retry.Retry`
513542
:param retry: (Optional) The retry settings for this request.
514543
@@ -565,6 +594,7 @@ def batch_update(
565594
statements=parsed,
566595
seqno=seqno,
567596
request_options=request_options,
597+
last_statements=last_statement,
568598
)
569599

570600
method = functools.partial(

tests/mockserver_tests/test_basics.py

Lines changed: 57 additions & 0 deletions
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
Lines changed: 127 additions & 0 deletions
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

Lines changed: 8 additions & 4 deletions
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)