Skip to content

Commit 54ed079

Browse files
google-genai-botcopybara-github
authored andcommitted
fix: Fix Spanner DatabaseSessionService support
Introduce `DynamicPickleType` to handle session actions, using sqlalchemy-spanner `SpannerPickleType` when the database dialect is Spanner. Connects to a Spanner database to store session data persistently in tables. # Example using Spanner database: `session_service = DatabaseSessionService(db_url="spanner+spanner:///projects/project-id/instances/instance-id/databases/database-id")` # Example adk web command: `adk web --session_service_uri="spanner+spanner:///projects/project-id/instances/instance-id/databases/database-id"` PiperOrigin-RevId: 797416610
1 parent c144b53 commit 54ed079

File tree

2 files changed

+30
-1
lines changed

2 files changed

+30
-1
lines changed

pyproject.toml

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

src/google/adk/sessions/database_session_service.py

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
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
@@ -104,6 +105,33 @@ def load_dialect_impl(self, dialect):
104105
return self.impl
105106

106107

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

@@ -209,7 +237,7 @@ class StorageEvent(Base):
209237
PreciseTimestamp, default=func.now()
210238
)
211239
content: Mapped[dict[str, Any]] = mapped_column(DynamicJSON, nullable=True)
212-
actions: Mapped[MutableDict[str, Any]] = mapped_column(PickleType)
240+
actions: Mapped[MutableDict[str, Any]] = mapped_column(DynamicPickleType)
213241

214242
long_running_tool_ids_json: Mapped[Optional[str]] = mapped_column(
215243
Text, nullable=True

0 commit comments

Comments
 (0)