Skip to content

Commit 02bb0e2

Browse files
olavloiteolavloite
committed
feat: add SpannerPickleType
Adds a SpannerPickleType that can be used as the implementation for the standard SQLAlchemy PickleType. The SpannerPickleType ensures that the binary values are encoded/decoded to/from base64 strings, which is how Spanner stores binary values. Fixes #654
1 parent 16c87e4 commit 02bb0e2

File tree

5 files changed

+303
-1
lines changed

5 files changed

+303
-1
lines changed

‎google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
14+
import base64
1415

1516
import pkg_resources
1617
import re
@@ -27,7 +28,7 @@
2728
from google.cloud.spanner_v1 import Client
2829
from sqlalchemy.exc import NoSuchTableError
2930
from sqlalchemy.sql import elements
30-
from sqlalchemy import ForeignKeyConstraint, types
31+
from sqlalchemy import ForeignKeyConstraint, types, TypeDecorator, PickleType
3132
from sqlalchemy.engine.base import Engine
3233
from sqlalchemy.engine.default import DefaultDialect, DefaultExecutionContext
3334
from sqlalchemy.event import listens_for
@@ -78,6 +79,35 @@ def reset_connection(dbapi_conn, connection_record, reset_state=None):
7879
OPERATORS[json_getitem_op] = operator_lookup["json_getitem_op"]
7980

8081

82+
# PickleType that can be used with Spanner.
83+
# Binary values are automatically encoded/decoded to/from base64.
84+
# Usage:
85+
# class User(Base):
86+
# __tablename__ = 'users'
87+
#
88+
# user_id = Column(Integer, primary_key=True)
89+
# username = Column(String(50), nullable=False)
90+
# preferences = Column(PickleType(impl=SpannerPickleType))
91+
class SpannerPickleType(TypeDecorator):
92+
impl = PickleType
93+
94+
def bind_processor(self, dialect):
95+
def process(value):
96+
if value is None:
97+
return None
98+
return base64.standard_b64encode(value)
99+
100+
return process
101+
102+
def result_processor(self, dialect, coltype):
103+
def process(value):
104+
if value is None:
105+
return None
106+
return base64.standard_b64decode(value)
107+
108+
return process
109+
110+
81111
# Spanner-to-SQLAlchemy types map
82112
_type_map = {
83113
"BOOL": types.Boolean,

‎samples/model.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,10 @@
3232
Sequence,
3333
TextClause,
3434
Index,
35+
PickleType,
3536
)
3637
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
38+
from google.cloud.sqlalchemy_spanner.sqlalchemy_spanner import SpannerPickleType
3739

3840

3941
class Base(DeclarativeBase):
@@ -64,6 +66,9 @@ class Singer(Base):
6466
)
6567
birthdate: Mapped[Optional[datetime.date]] = mapped_column(Date, nullable=True)
6668
picture: Mapped[Optional[bytes]] = mapped_column(LargeBinary, nullable=True)
69+
preferences: Mapped[Optional[object]] = mapped_column(
70+
PickleType(impl=SpannerPickleType), nullable=True
71+
)
6772
albums: Mapped[List["Album"]] = relationship(
6873
back_populates="singer", cascade="all, delete-orphan"
6974
)

‎samples/pickle_type_sample.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,55 @@
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+
17+
from sqlalchemy import create_engine
18+
from sqlalchemy.orm import Session
19+
20+
from sample_helper import run_sample
21+
from model import Singer
22+
23+
# Shows how to use PickleType with Spanner.
24+
def pickle_type():
25+
engine = create_engine(
26+
"spanner:///projects/sample-project/"
27+
"instances/sample-instance/"
28+
"databases/sample-database",
29+
echo=True,
30+
)
31+
with Session(engine) as session:
32+
singer = Singer(
33+
id=str(uuid.uuid4()),
34+
first_name="John",
35+
last_name="Smith",
36+
# Preferences are stored as an opaque BYTES column
37+
# in the database.
38+
preferences={
39+
"wakeup_call": "yes",
40+
"vegetarian": "no",
41+
},
42+
)
43+
session.add(singer)
44+
session.commit()
45+
46+
# Use AUTOCOMMIT for sessions that only read. This is more
47+
# efficient than using a read/write transaction to only read.
48+
session.connection(execution_options={"isolation_level": "AUTOCOMMIT"})
49+
print(
50+
f"Inserted singer {singer.full_name} has these preferences: {singer.preferences}"
51+
)
52+
53+
54+
if __name__ == "__main__":
55+
run_sample(pickle_type)

‎test/mockserver_tests/pickle_type_model.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
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 sqlalchemy import Column, Integer, String, PickleType
16+
from sqlalchemy.orm import DeclarativeBase
17+
18+
from google.cloud.sqlalchemy_spanner.sqlalchemy_spanner import SpannerPickleType
19+
20+
21+
class Base(DeclarativeBase):
22+
pass
23+
24+
25+
class UserPreferences(Base):
26+
__tablename__ = "user_preferences"
27+
28+
user_id = Column(Integer, primary_key=True)
29+
username = Column(String(50), nullable=False)
30+
preferences = Column(PickleType(impl=SpannerPickleType), nullable=True)
31+
created_at = Column(String(30), nullable=False)

‎test/mockserver_tests/test_pickle_type.py

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
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 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+
BeginTransactionRequest,
25+
TypeCode,
26+
)
27+
from test.mockserver_tests.mock_server_test_base import (
28+
MockServerTestBase,
29+
add_result,
30+
add_update_count,
31+
)
32+
from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest
33+
import google.cloud.spanner_v1.types.type as spanner_type
34+
import google.cloud.spanner_v1.types.result_set as result_set
35+
36+
37+
class TestPickleType(MockServerTestBase):
38+
def test_create_table(self):
39+
from test.mockserver_tests.pickle_type_model import Base
40+
41+
add_result(
42+
"""SELECT true
43+
FROM INFORMATION_SCHEMA.TABLES
44+
WHERE TABLE_SCHEMA="" AND TABLE_NAME="user_preferences"
45+
LIMIT 1
46+
""",
47+
ResultSet(),
48+
)
49+
engine = create_engine(
50+
"spanner:///projects/p/instances/i/databases/d",
51+
connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
52+
)
53+
Base.metadata.create_all(engine)
54+
requests = self.database_admin_service.requests
55+
eq_(1, len(requests))
56+
is_instance_of(requests[0], UpdateDatabaseDdlRequest)
57+
eq_(1, len(requests[0].statements))
58+
eq_(
59+
"CREATE TABLE user_preferences (\n"
60+
"\tuser_id INT64 NOT NULL GENERATED BY DEFAULT"
61+
" AS IDENTITY (BIT_REVERSED_POSITIVE), \n"
62+
"\tusername STRING(50) NOT NULL, \n"
63+
"\tpreferences BYTES(MAX), \n"
64+
"\tcreated_at STRING(30) NOT NULL\n"
65+
") PRIMARY KEY (user_id)",
66+
requests[0].statements[0],
67+
)
68+
69+
def test_insert_and_query(self):
70+
from test.mockserver_tests.pickle_type_model import UserPreferences
71+
72+
add_update_count(
73+
"INSERT INTO user_preferences (user_id, username, preferences, created_at) "
74+
"VALUES (@a0, @a1, @a2, @a3)",
75+
1,
76+
)
77+
engine = create_engine(
78+
"spanner:///projects/p/instances/i/databases/d",
79+
connect_args={"client": self.client, "pool": FixedSizePool(size=10)},
80+
)
81+
preferences = {"setting": "true"}
82+
preferences_base64 = "gAWVFQAAAAAAAAB9lIwHc2V0dGluZ5SMBHRydWWUcy4="
83+
with Session(engine) as session:
84+
new_user = UserPreferences(
85+
user_id=1,
86+
username="test_user",
87+
preferences=preferences,
88+
created_at="2025-05-04T00:00:00.000000",
89+
)
90+
91+
session.add(new_user)
92+
session.commit()
93+
94+
# Verify the requests that we got.
95+
requests = self.spanner_service.requests
96+
eq_(4, len(requests))
97+
is_instance_of(requests[0], BatchCreateSessionsRequest)
98+
is_instance_of(requests[1], BeginTransactionRequest)
99+
is_instance_of(requests[2], ExecuteSqlRequest)
100+
is_instance_of(requests[3], CommitRequest)
101+
request: ExecuteSqlRequest = requests[2]
102+
eq_(4, len(request.params))
103+
eq_("1", request.params["a0"])
104+
eq_("test_user", request.params["a1"])
105+
eq_(preferences_base64, request.params["a2"])
106+
eq_(TypeCode.INT64, request.param_types["a0"].code)
107+
eq_(TypeCode.STRING, request.param_types["a1"].code)
108+
eq_(TypeCode.BYTES, request.param_types["a2"].code)
109+
110+
add_user_preferences_result(
111+
"SELECT user_preferences.user_id AS user_preferences_user_id, "
112+
"user_preferences.username AS user_preferences_username, "
113+
"user_preferences.preferences AS user_preferences_preferences, "
114+
"user_preferences.created_at AS user_preferences_created_at\n"
115+
"FROM user_preferences\n"
116+
"WHERE user_preferences.user_id = @a0\n"
117+
" LIMIT @a1",
118+
preferences_base64,
119+
)
120+
user = session.query(UserPreferences).filter_by(user_id=1).first()
121+
eq_(preferences, user.preferences)
122+
123+
124+
def add_user_preferences_result(sql: str, preferences_base64: object):
125+
result = result_set.ResultSet(
126+
dict(
127+
metadata=result_set.ResultSetMetadata(
128+
dict(
129+
row_type=spanner_type.StructType(
130+
dict(
131+
fields=[
132+
spanner_type.StructType.Field(
133+
dict(
134+
name="user_id",
135+
type=spanner_type.Type(
136+
dict(code=spanner_type.TypeCode.INT64)
137+
),
138+
)
139+
),
140+
spanner_type.StructType.Field(
141+
dict(
142+
name="user_name",
143+
type=spanner_type.Type(
144+
dict(code=spanner_type.TypeCode.STRING)
145+
),
146+
)
147+
),
148+
spanner_type.StructType.Field(
149+
dict(
150+
name="preferences",
151+
type=spanner_type.Type(
152+
dict(code=spanner_type.TypeCode.BYTES)
153+
),
154+
)
155+
),
156+
spanner_type.StructType.Field(
157+
dict(
158+
name="created_at",
159+
type=spanner_type.Type(
160+
dict(code=spanner_type.TypeCode.TIMESTAMP)
161+
),
162+
)
163+
),
164+
]
165+
)
166+
)
167+
)
168+
),
169+
)
170+
)
171+
result.rows.extend(
172+
[
173+
(
174+
"1",
175+
"Test",
176+
preferences_base64,
177+
"2025-05-05T00:00:00.000000Z",
178+
),
179+
]
180+
)
181+
add_result(sql, result)

0 commit comments

Comments
 (0)