8000 Set the max size of strings in database columns. · JCRCS/adk-python@2a9ddec · GitHub
[go: up one dir, main page]
Skip to content

Commit 2a9ddec

Browse files
DeanChensjcopybara-github
DeanChensj
authored and
copybara-github
committed
Set the max size of strings in database columns.
PiperOrigin-RevId: 752918808
1 parent 2ea4315 commit 2a9ddec

File tree

1 file changed

+49
-16
lines changed

1 file changed

+49
-16
lines changed

‎src/google/adk/sessions/database_session_service.py

Lines changed: 49 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,8 @@
5858

5959
logger = logging.getLogger(__name__)
6060

61+
DEFAULT_MAX_VARCHAR_LENGTH = 256
62+
6163

6264
class DynamicJSON(TypeDecorator):
6365
"""A JSON-like type that uses JSONB on PostgreSQL and TEXT with JSON
@@ -92,17 +94,25 @@ def process_result_value(self, value, dialect: Dialect):
9294

9395
class Base(DeclarativeBase):
9496
"""Base class for database tables."""
97+
9598
pass
9699

97100

98101
class StorageSession(Base):
99102
"""Represents a session stored in the database."""
103+
100104
__tablename__ = "sessions"
101105

102-
app_name: Mapped[str] = mapped_column(String, primary_key=True)
103-
user_id: Mapped[str] = mapped_column(String, primary_key=True)
106+
app_name: Mapped[str] = mapped_column(
107+
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
108+
)
109+
user_id: Mapped[str] = mapped_column(
110+
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
111+
)
104112
id: Mapped[str] = mapped_column(
105-
String, primary_key=True, default=lambda: str(uuid.uuid4())
113+
String(DEFAULT_MAX_VARCHAR_LENGTH),
114+
primary_key=True,
115+
default=lambda: str(uuid.uuid4()),
106116
)
107117

108118
state: Mapped[MutableDict[str, Any]] = mapped_column(
@@ -125,16 +135,27 @@ def __repr__(self):
125135

126136
class StorageEvent(Base):
127137
"""Represents an event stored in the database."""
138+
128139
__tablename__ = "events"
129140

130-
id: Mapped[str] = mapped_column(String, primary_key=True)
131-
app_name: Mapped[str] = mapped_column(String, primary_key=True)
132-
user_id: Mapped[str] = mapped_column(String, primary_key=True)
133-
session_id: Mapped[str] = mapped_column(String, primary_key=True)
141+
id: Mapped[str] = mapped_column(
142+
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
143+
)
144+
app_name: Mapped[str] = mapped_column(
145+
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
146+
)
147+
user_id: Mapped[str] = mapped_column(
148+
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
149+
)
150+
session_id: Mapped[str] = mapped_column(
151+
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
152+
)
134153

135-
invocation_id: Mapped[str] = mapped_column(String)
136-
author: Mapped[str] = mapped_column(String)
137-
branch: Mapped[str] = mapped_column(String, nullable=True)
154+
invocation_id: Mapped[str] = mapped_column(String(DEFAULT_MAX_VARCHAR_LENGTH))
155+
author: Mapped[str] = mapped_column(String(DEFAULT_MAX_VARCHAR_LENGTH))
156+
branch: Mapped[str] = mapped_column(
157+
String(DEFAULT_MAX_VARCHAR_LENGTH), nullable=True
158+
)
138159
timestamp: Mapped[DateTime] = mapped_column(DateTime(), default=func.now())
139160
content: Mapped[dict[str, Any]] = mapped_column(DynamicJSON, nullable=True)
140161
actions: Mapped[MutableDict[str, Any]] = mapped_column(PickleType)
@@ -147,8 +168,10 @@ class StorageEvent(Base):
147168
)
148169
partial: Mapped[bool] = mapped_column(Boolean, nullable=True)
149170
turn_complete: Mapped[bool] = mapped_column(Boolean, nullable=True)
150-
error_code: Mapped[str] = mapped_column(String, nullable=True)
151-
error_message: Mapped[str] = mapped_column(String, nullable=True)
171+
error_code: Mapped[str] = mapped_column(
172+
String(DEFAULT_MAX_VARCHAR_LENGTH), nullable=True
173+
)
174+
error_message: Mapped[str] = mapped_column(String(1024), nullable=True)
152175
interrupted: Mapped[bool] = mapped_column(Boolean, nullable=True)
153176

154177
storage_session: Mapped[StorageSession] = relationship(
@@ -182,23 +205,33 @@ def long_running_tool_ids(self, value: set[str]):
182205

183206
class StorageAppState(Base):
184207
"""Represents an app state stored in the database."""
208+
185209
__tablename__ = "app_states"
186210

187-
app_name: Mapped[str] = mapped_column(String, primary_key=True)
211+
app_name: Mapped[str] = mapped_column(
212+
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
213+
)
188214
state: Mapped[MutableDict[str, Any]] = mapped_column(
189215
MutableDict.as_mutable(DynamicJSON), default={}
190216
)
191217
update_time: Mapped[DateTime] = mapped_column(
192218
DateTime(), default=func.now(), onupdate=func.now()
193219
)
194220

195-
196221
class StorageUserState(Base):
197222
"""Represents a user state stored in the database."""
223+
198224
__tablename__ = "user_states"
199225

200-
app_name: Mapped[s 67F4 tr] = mapped_column(String, primary_key=True)
201-
user_id: Mapped[str] = mapped_column(String, primary_key=True)
226+
app_name: Mapped[str] = mapped_column(
227+
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
228+
)
229+
user_id: Mapped[str] = mapped_column(
230+
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
231+
)
232+
state: Mapped[MutableDict[str, Any]] = mapped_column(
233+
MutableDict.as_mutable(DynamicJSON), default={}
234+
)
202235
state: Mapped[MutableDict[str, Any]] = mapped_column(
203236
MutableDict.as_mutable(DynamicJSON), default={}
204237
)

0 commit comments

Comments
 (0)
0