Skip to content

Commit 1fe9f4b

Browse files
waltaskewolavloite
waltaskew
and
olavloite
authored
feat: enable SQLAlchemy 2.0's insertmany feature (#721)
* feat: enable SQLAlchemy 2.0's insertmany feature - Enable the use_insertmanyvalues flag on the dialect to support SQLAlchemy 2.0's insertmany feature, which allows multiple ORM objects to be inserted together in bulk, even if the table has server-side generated values which must be included in a THEN RETURN clause - Provide an example for using the feature with client-side supplied UUIDs and insert_sentinel columns. - Ensure that the feature is not enables for bit-reversed primary keys. In other dialects, an incrementing primary key can be used rather than a sentinel column. In Spanner, the bit-reversed integers do not meet the ordering requirement to be used as implicit sentinels https://docs.sqlalchemy.org/en/20/core/internals.html#sqlalchemy.engine.default.DefaultDialect.use_insertmanyvalues https://docs.sqlalchemy.org/en/20/core/connections.html#insert-many-values-behavior-for-insert-statements https://docs.sqlalchemy.org/en/20/core/connections.html#configuring-sentinel-columns Fixes: #720 * chore: update tests for multiplexed sessions * chore: make sample runnable * fix: exclude commit_timestamp col from THEN RETURN --------- Co-authored-by: Knut Olav Løite <koloite@gmail.com>
1 parent e6b6cf6 commit 1fe9f4b

File tree

6 files changed

+339
-1
lines changed

6 files changed

+339
-1
lines changed

‎google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -425,7 +425,14 @@ def returning_clause(self, stmt, returning_cols, **kw):
425425
self._label_select_column(
426426
None, c, True, False, {"spanner_is_returning": True}
427427
)
428-
for c in expression._select_iterables(returning_cols)
428+
for c in expression._select_iterables(
429+
filter(
430+
lambda col: not col.dialect_options.get("spanner", {}).get(
431+
"exclude_from_returning", False
432+
),
433+
returning_cols,
434+
)
435+
)
429436
]
430437

431438
return "THEN RETURN " + ", ".join(columns)
@@ -831,6 +838,7 @@ class SpannerDialect(DefaultDialect):
831838
update_returning = True
832839
delete_returning = True
833840
supports_multivalues_insert = True
841+
use_insertmanyvalues = True
834842

835843
ddl_compiler = SpannerDDLCompiler
836844
preparer = SpannerIdentifierPreparer

‎samples/insertmany_sample.py‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
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+
16+
from datetime import datetime
17+
import uuid
18+
from sqlalchemy import text, String, create_engine
19+
from sqlalchemy.orm import DeclarativeBase, Session
20+
from sqlalchemy.orm import Mapped
21+
from sqlalchemy.orm import mapped_column
22+
from sample_helper import run_sample
23+
24+
25+
class Base(DeclarativeBase):
26+
pass
27+
28+
29+
# To use SQLAlchemy 2.0's insertmany feature, models must have a
30+
# unique column marked as an "insert_sentinal" with client-side
31+
# generated values passed into it. This allows SQLAlchemy to perform a
32+
# single bulk insert, even if the table has columns with server-side
33+
# defaults which must be retrieved from a THEN RETURN clause, for
34+
# operations like:
35+
#
36+
# with Session.begin() as session:
37+
# session.add(Singer(name="a"))
38+
# session.add(Singer(name="b"))
39+
#
40+
# Read more in the SQLAlchemy documentation of this feature:
41+
# https://docs.sqlalchemy.org/en/20/core/connections.html#configuring-sentinel-columns
42+
43+
44+
class Singer(Base):
45+
__tablename__ = "singers_with_sentinel"
46+
id: Mapped[str] = mapped_column(
47+
String(36),
48+
primary_key=True,
49+
# Supply a unique UUID client-side
50+
default=lambda: str(uuid.uuid4()),
51+
# The column is unique and can be used as an insert_sentinel
52+
insert_sentinel=True,
53+
# Set a server-side default for write outside SQLAlchemy
54+
server_default=text("GENERATE_UUID()"),
55+
)
56+
name: Mapped[str]
57+
inserted_at: Mapped[datetime] = mapped_column(
58+
server_default=text("CURRENT_TIMESTAMP()")
59+
)
60+
61+
62+
# Shows how to insert data using SQLAlchemy, including relationships that are
63+
# defined both as foreign keys and as interleaved tables.
64+
def insertmany():
65+
engine = create_engine(
66+
"spanner:///projects/sample-project/"
67+
"instances/sample-instance/"
68+
"databases/sample-database",
69+
echo=True,
70+
)
71+
# Create the sample table.
72+
Base.metadata.create_all(engine)
73+
74+
# Insert two singers in one session. These two singers will be inserted using
75+
# a single INSERT statement with a THEN RETURN clause to return the generated
76+
# creation timestamp.
77+
with Session(engine) as session:
78+
session.add(Singer(name="John Smith"))
79+
session.add(Singer(name="Jane Smith"))
80+
session.commit()
81+
82+
83+
if __name__ == "__main__":
84+
run_sample(insertmany)

‎samples/noxfile.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,11 @@ def informational_fk(session):
9292
_sample(session)
9393

9494

95+
@nox.session()
96+
def insertmany(session):
97+
_sample(session)
98+
99+
95100
@nox.session()
96101
def _all_samples(session):
97102
_sample(session)

‎test/mockserver_tests/insertmany_model.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
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 datetime import datetime
16+
import uuid
17+
from sqlalchemy import text, String
18+
from sqlalchemy.orm import DeclarativeBase
19+
from sqlalchemy.orm import Mapped
20+
from sqlalchemy.orm import mapped_column
21+
22+
23+
class Base(DeclarativeBase):
24+
pass
25+
26+
27+
class SingerUUID(Base):
28+
__tablename__ = "singers_uuid"
29+
id: Mapped[str] = mapped_column(
30+
String(36),
31+
primary_key=True,
32+
server_default=text("GENERATE_UUID()"),
33+
default=lambda: str(uuid.uuid4()),
34+
insert_sentinel=True,
35+
)
36+
name: Mapped[str]
37+
inserted_at: Mapped[datetime] = mapped_column(
38+
server_default=text("CURRENT_TIMESTAMP()")
39+
)
40+
41+
42+
class SingerIntID(Base):
43+
__tablename__ = "singers_int_id"
44+
id: Mapped[int] = mapped_column(primary_key=True)
45+
name: Mapped[str] = mapped_column(String)
46+
inserted_at: Mapped[datetime] = mapped_column(
47+
server_default=text("CURRENT_TIMESTAMP()")
48+
)

‎test/mockserver_tests/test_insertmany.py‎

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
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+
import uuid
16+
from unittest import mock
17+
18+
import sqlalchemy
19+
from sqlalchemy.orm import Session
20+
from sqlalchemy.testing import eq_, is_instance_of
21+
from google.cloud.spanner_v1 import (
22+
ExecuteSqlRequest,
23+
CommitRequest,
24+
RollbackRequest,
25+
BeginTransactionRequest,
26+
CreateSessionRequest,
27+
)
28+
from test.mockserver_tests.mock_server_test_base import (
29+
MockServerTestBase,
30+
add_result,
31+
)
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 TestInsertmany(MockServerTestBase):
37+
@mock.patch.object(uuid, "uuid4", mock.MagicMock(side_effect=["a", "b"]))
38+
def test_insertmany_with_uuid_sentinels(self):
39+
"""Ensures one bulk insert for ORM objects distinguished by uuid."""
40+
from test.mockserver_tests.insertmany_model import SingerUUID
41+
42+
self.add_uuid_insert_result(
43+
"INSERT INTO singers_uuid (id, name) "
44+
"VALUES (@a0, @a1), (@a2, @a3) "
45+
"THEN RETURN inserted_at, id"
46+
)
47+
engine = self.create_engine()
48+
49+
with Session(engine) as session:
50+
session.add(SingerUUID(name="a"))
51+
session.add(SingerUUID(name="b"))
52+
session.commit()
53+
54+
# Verify the requests that we got.
55+
requests = self.spanner_service.requests
56+
eq_(4, len(requests))
57+
is_instance_of(requests[0], CreateSessionRequest)
58+
is_instance_of(requests[1], BeginTransactionRequest)
59+
is_instance_of(requests[2], ExecuteSqlRequest)
60+
is_instance_of(requests[3], CommitRequest)
61+
62+
def test_no_insertmany_with_bit_reversed_id(self):
63+
"""Ensures we don't try to bulk insert rows with bit-reversed PKs.
64+
65+
SQLAlchemy's insertmany support requires either incrementing
66+
PKs or client-side supplied sentinel values such as UUIDs.
67+
Spanner's bit-reversed integer PKs don't meet the ordering
68+
requirement, so we need to make sure we don't try to bulk
69+
insert with them.
70+
"""
71+
from test.mockserver_tests.insertmany_model import SingerIntID
72+
73+
self.add_int_id_insert_result(
74+
"INSERT INTO singers_int_id (name) "
75+
"VALUES (@a0) "
76+
"THEN RETURN id, inserted_at"
77+
)
78+
engine = self.create_engine()
79+
80+
with Session(engine) as session:
81+
session.add(SingerIntID(name="a"))
82+
session.add(SingerIntID(name="b"))
83+
try:
84+
session.commit()
85+
except sqlalchemy.exc.SAWarning:
86+
# This will fail because we're returning the same PK
87+
# for two rows. The mock server doesn't currently
88+
# support associating the same query with two
89+
# different results. For our purposes that's okay --
90+
# we just want to ensure we generate two INSERTs, not
91+
# one.
92+
pass
93+
94+
# Verify the requests that we got.
95+
requests = self.spanner_service.requests
96+
eq_(5, len(requests))
97+
is_instance_of(requests[0], CreateSessionRequest)
98+
is_instance_of(requests[1], BeginTransactionRequest)
99+
is_instance_of(requests[2], ExecuteSqlRequest)
100+
is_instance_of(requests[3], ExecuteSqlRequest)
101+
is_instance_of(requests[4], RollbackRequest)
102+
103+
def add_uuid_insert_result(self, sql):
104+
result = result_set.ResultSet(
105+
dict(
106+
metadata=result_set.ResultSetMetadata(
107+
dict(
108+
row_type=spanner_type.StructType(
109+
dict(
110+
fields=[
111+
spanner_type.StructType.Field(
112+
dict(
113+
name="inserted_at",
114+
type=spanner_type.Type(
115+
dict(
116+
code=spanner_type.TypeCode.TIMESTAMP
117+
)
118+
),
119+
)
120+
),
121+
spanner_type.StructType.Field(
122+
dict(
123+
name="id",
124+
type=spanner_type.Type(
125+
dict(code=spanner_type.TypeCode.STRING)
126+
),
127+
)
128+
),
129+
]
130+
)
131+
)
132+
)
133+
),
134+
)
135+
)
136+
result.rows.extend(
137+
[
138+
(
139+
"2020-06-02T23:58:40Z",
140+
"a",
141+
),
142+
(
143+
"2020-06-02T23:58:41Z",
144+
"b",
145+
),
146+
]
147+
)
148+
add_result(sql, result)
149+
150+
def add_int_id_insert_result(self, sql):
151+
result = result_set.ResultSet(
152+
dict(
153+
metadata=result_set.ResultSetMetadata(
154+
dict(
155+
row_type=spanner_type.StructType(
156+
dict(
157+
fields=[
158+
spanner_type.StructType.Field(
159+
dict(
160+
name="id",
161+
type=spanner_type.Type(
162+
dict(code=spanner_type.TypeCode.INT64)
163+
),
164+
)
165+
),
166+
spanner_type.StructType.Field(
167+
dict(
168+
name="inserted_at",
169+
type=spanner_type.Type(
170+
dict(
171+
code=spanner_type.TypeCode.TIMESTAMP
172+
)
173+
),
174+
)
175+
),
176+
]
177+
)
178+
)
179+
)
180+
),
181+
)
182+
)
183+
result.rows.extend(
184+
[
185+
(
186+
"1",
187+
"2020-06-02T23:58:40Z",
188+
),
189+
]
190+
)
191+
add_result(sql, result)

‎test/system/test_basics.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,8 @@ class TimestampUser(Base):
316316
updated_at: Mapped[datetime.datetime] = mapped_column(
317317
spanner_allow_commit_timestamp=True,
318318
default=text("PENDING_COMMIT_TIMESTAMP()"),
319+
# Make sure that this column is never part of a THEN RETURN clause.
320+
spanner_exclude_from_returning=True,
319321
)
320322

321323
@event.listens_for(TimestampUser, "before_update")

0 commit comments

Comments
 (0)