Skip to content

Commit ac64472

Browse files
olavloiteolavloite
authored
fix: support THEN RETURN for insert, update, delete (#503)
* fix: support THEN RETURN for insert, update, delete Support THEN RETURN clauses for INSERT, UPDATE, and DELETE statements. Fixes #498 * test: override insert auto-generated pk test
1 parent 142cbee commit ac64472

File tree

5 files changed

+201
-1
lines changed

5 files changed

+201
-1
lines changed

‎google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -611,6 +611,10 @@ class SpannerDialect(DefaultDialect):
611611
supports_native_decimal = True
612612
supports_statement_cache = True
613613

614+
insert_returning = True
615+
update_returning = True
616+
delete_returning = True
617+
614618
ddl_compiler = SpannerDDLCompiler
615619
preparer = SpannerIdentifierPreparer
616620
statement_compiler = SpannerSQLCompiler

‎noxfile.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -327,7 +327,7 @@ def mockserver(session):
327327
"9999",
328328
)
329329
session.run(
330-
"py.test", "--quiet", os.path.join("test/mockserver_tests"), *session.posargs
330+
"py.test", "--quiet", os.path.join("test", "mockserver_tests"), *session.posargs
331331
)
332332

333333

‎test/mockserver_tests/bit_reversed_sequence_model.py

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
# Copyright 2024 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 sqlalchemy import String, BigInteger, Sequence, TextClause
16+
from sqlalchemy.orm import DeclarativeBase
17+
from sqlalchemy.orm import Mapped
18+
from sqlalchemy.orm import mapped_column
19+
20+
21+
class Base(DeclarativeBase):
22+
pass
23+
24+
25+
class Singer(Base):
26+
__tablename__ = "singers"
27+
id: Mapped[int] = mapped_column(
28+
BigInteger,
29+
Sequence("singer_id"),
30+
server_default=TextClause("GET_NEXT_SEQUENCE_VALUE(SEQUENCE singer_id)"),
31+
primary_key=True,
32+
)
33+
name: Mapped[str] = mapped_column(String)

‎test/mockserver_tests/test_bit_reversed_sequence.py

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
# Copyright 2024 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 sqlalchemy import create_engine
16+
from sqlalchemy.orm import Session
17+
from sqlalchemy.testing import eq_, is_instance_of
18+
from google.cloud.spanner_v1 import (
19+
FixedSizePool,
20+
ResultSet,
21+
BatchCreateSessionsRequest,
22+
ExecuteSqlRequest,
23+
CommitRequest,
24+
GetSessionRequest,
25+
BeginTransactionRequest,
26+
)
27+
from test.mockserver_tests.mock_server_test_base import (
28+
MockServerTestBase,
29+
add_result,
30+
)
31+
from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest
32+
import google.cloud.spanner_v1.types.type as spanner_type
33+
import google.cloud.spanner_v1.types.result_set as result_set
34+
35+
36+
class TestBitReversedSequence(MockServerTestBase):
37+
def test_create_table(self):
38+
from test.mockserver_tests.bit_reversed_sequence_model import Base
39+
40+
add_result(
41+
"""SELECT true
42+
FROM INFORMATION_SCHEMA.TABLES
43+
WHERE TABLE_SCHEMA="" AND TABLE_NAME="singers"
44+
LIMIT 1
45+
""",
46+
ResultSet(),
47+
)
48+
add_result(
49+
"""SELECT true
50+
FROM INFORMATION_SCHEMA.SEQUENCES
51+
WHERE NAME="singer_id"
52+
AND SCHEMA=""
53+
LIMIT 1""",
54+
ResultSet(),
55+
)
56+
engine = create_engine(
57+
"spanner:///projects/p/instances/i/databases/d",
58+
connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
59+
)
60+
Base.metadata.create_all(engine)
61+
requests = self.database_admin_service.requests
62+
eq_(1, len(requests))
63+
is_instance_of(requests[0], UpdateDatabaseDdlRequest)
64+
eq_(2, len(requests[0].statements))
65+
eq_(
66+
"CREATE SEQUENCE singer_id OPTIONS "
67+
"(sequence_kind = 'bit_reversed_positive')",
68+
requests[0].statements[0],
69+
)
70+
eq_(
71+
"CREATE TABLE singers (\n"
72+
"\tid INT64 NOT NULL DEFAULT "
73+
"(GET_NEXT_SEQUENCE_VALUE(SEQUENCE singer_id)), \n"
74+
"\tname STRING(MAX) NOT NULL\n"
75+
") PRIMARY KEY (id)",
76+
requests[0].statements[1],
77+
)
78+
79+
def test_insert_row(self):
80+
from test.mockserver_tests.bit_reversed_sequence_model import Singer
81+
82+
result = result_set.ResultSet(
83+
dict(
84+
metadata=result_set.ResultSetMetadata(
85+
dict(
86+
row_type=spanner_type.StructType(
87+
dict(
88+
fields=[
89+
spanner_type.StructType.Field(
90+
dict(
91+
name="id",
92+
type=spanner_type.Type(
93+
dict(code=spanner_type.TypeCode.INT64)
94+
),
95+
)
96+
)
97+
]
98+
)
99+
)
100+
)
101+
),
102+
stats=result_set.ResultSetStats(
103+
dict(
104+
row_count_exact=1,
105+
)
106+
),
107+
)
108+
)
109+
result.rows.extend(["1"])
110+
111+
add_result(
112+
"INSERT INTO singers (id, name) "
113+
"VALUES ( GET_NEXT_SEQUENCE_VALUE(SEQUENCE singer_id), @a0) "
114+
"THEN RETURN singers.id",
115+
result,
116+
)
117+
engine = create_engine(
118+
"spanner:///projects/p/instances/i/databases/d",
119+
connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
120+
)
121+
122+
with Session(engine) as session:
123+
singer = Singer(name="Test")
124+
session.add(singer)
125+
# Flush the session to send the insert statement to the database.
126+
session.flush()
127+
eq_(1, singer.id)
128+
session.commit()
129+
# Verify the requests that we got.
130+
requests = self.spanner_service.requests
131+
eq_(5, len(requests))
132+
is_instance_of(requests[0], BatchCreateSessionsRequest)
133+
# We should get rid of this extra round-trip for GetSession....
134+
is_instance_of(requests[1], GetSessionRequest)
135+
is_instance_of(requests[2], BeginTransactionRequest)
136+
is_instance_of(requests[3], ExecuteSqlRequest)
137+
is_instance_of(requests[4], CommitRequest)

‎test/test_suite_20.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2171,6 +2171,32 @@ def test_autoclose_on_insert(self):
21712171
assert r.is_insert
21722172
assert not r.returns_rows
21732173

2174+
def test_autoclose_on_insert_implicit_returning(self, connection):
2175+
"""
2176+
SPANNER OVERRIDE:
2177+
2178+
Cloud Spanner doesn't support tables with an auto increment primary key,
2179+
following insertions will fail with `400 id must not be NULL in table
2180+
autoinc_pk`.
2181+
2182+
Overriding the tests and adding a manual primary key value to avoid the same
2183+
failures.
2184+
"""
2185+
r = connection.execute(
2186+
# return_defaults() ensures RETURNING will be used,
2187+
# new in 2.0 as sqlite/mariadb offer both RETURNING and
2188+
# cursor.lastrowid
2189+
self.tables.autoinc_pk.insert().return_defaults(),
2190+
dict(id=2, data="some data"),
2191+
)
2192+
assert r._soft_closed
2193+
assert not r.closed
2194+
assert r.is_insert
2195+
2196+
# Spanner does not return any rows in this case, because the primary key
2197+
# is not auto-generated.
2198+
assert not r.returns_rows
2199+
21742200

21752201
class BytesTest(_LiteralRoundTripFixture, fixtures.TestBase):
21762202
__backend__ = True

0 commit comments

Comments
 (0)