Skip to content

Commit 454b9f0

Browse files
google-genai-botcopybara-github
authored andcommitted
fix: Fix Spanner database session service support
Introduce `DynamicPickleType` to handle session actions, using sqlalchemy-spanner `SpannerPickleType` when the database dialect is Spanner. PiperOrigin-RevId: 794673730
1 parent ddf2e21 commit 454b9f0

File tree

2 files changed

+29
-1
lines changed

2 files changed

+29
-1
lines changed

pyproject.toml

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ dependencies = [
4747
"python-dateutil>=2.9.0.post0", # For Vertext AI Session Service
4848
"python-dotenv>=1.0.0", # To manage environment variables
4949
"requests>=2.32.4",
50+
"sqlalchemy-spanner>=1.14.0", # Spanner database session service
5051
"sqlalchemy>=2.0", # SQL database ORM
5152
"starlette>=0.46.2", # For FastAPI CLI
5253
"tenacity>=8.0.0", # For Retry management

src/google/adk/sessions/database_session_service.py

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,12 @@
1818
from datetime import timezone
1919
import json
2020
import logging
21+
import pickle
2122
from typing import Any
2223
from typing import Optional
2324
import uuid
2425

26+
from google.cloud.sqlalchemy_spanner.sqlalchemy_spanner import SpannerPickleType
2527
from sqlalchemy import Boolean
2628
from sqlalchemy import delete
2729
from sqlalchemy import Dialect
@@ -104,6 +106,31 @@ def load_dialect_impl(self, dialect):
104106
return self.impl
105107

106108

109+
class DynamicPickleType(TypeDecorator):
110+
"""Represents a type that can be pickled."""
111+
112+
impl = PickleType
113+
114+
def load_dialect_impl(self, dialect):
115+
if dialect.name == "spanner+spanner":
116+
return dialect.type_descriptor(SpannerPickleType)
117+
return self.impl
118+
119+
def process_bind_param(self, value, dialect):
120+
"""Ensures the pickled value is a bytes object before passing it to the database dialect."""
121+
if value is not None:
122+
if dialect.name == "spanner+spanner":
123+
return pickle.dumps(value)
124+
return value
125+
126+
def process_result_value(self, value, dialect):
127+
"""Ensures the raw bytes from the database are unpickled back into a Python object."""
128+
if value is not None:
129+
if dialect.name == "spanner+spanner":
130+
return pickle.loads(value)
131+
return value
132+
133+
107134
class Base(DeclarativeBase):
108135
"""Base class for database tables."""
109136

@@ -209,7 +236,7 @@ class StorageEvent(Base):
209236
PreciseTimestamp, default=func.now()
210237
)
211238
content: Mapped[dict[str, Any]] = mapped_column(DynamicJSON, nullable=True)
212-
actions: Mapped[MutableDict[str, Any]] = mapped_column(PickleType)
239+
actions: Mapped[MutableDict[str, Any]] = mapped_column(DynamicPickleType)
213240

214241
long_running_tool_ids_json: Mapped[Optional[str]] = mapped_column(
215242
Text, nullable=True

0 commit comments

Comments
 (0)