NeoAnd/adk-python
1
0
Fork 0
mirror of https://github.com/EvolutionAPI/adk-python.git synced 2025-07-14 09:51:25 -06:00
Code Issues Actions Packages Projects Releases Wiki Activity

Set the max size of strings in database columns.

Browse Source 
PiperOrigin-RevId: 752918808
This commit is contained in:
Shangjie Chen 2025-04-29 16:07:21 -07:00 committed by Copybara-Service
parent 2ea4315e9f
commit 2a9ddec7e3
1 changed files with 49 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

65
src/google/adk/sessions/database_session_service.py
View File

@ -58,6 +58,8 @@ from .state import State
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
DEFAULT_MAX_VARCHAR_LENGTH = 256
class DynamicJSON(TypeDecorator): class DynamicJSON(TypeDecorator):
"""A JSON-like type that uses JSONB on PostgreSQL and TEXT with JSON """A JSON-like type that uses JSONB on PostgreSQL and TEXT with JSON
@ -92,17 +94,25 @@ class DynamicJSON(TypeDecorator):
class Base(DeclarativeBase): class Base(DeclarativeBase):
"""Base class for database tables.""" """Base class for database tables."""
pass pass
class StorageSession(Base): class StorageSession(Base):
"""Represents a session stored in the database.""" """Represents a session stored in the database."""
__tablename__ = "sessions" __tablename__ = "sessions"
app_name: Mapped[str] = mapped_column(String, primary_key=True) app_name: Mapped[str] = mapped_column(
user_id: Mapped[str] = mapped_column(String, primary_key=True) String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
)
user_id: Mapped[str] = mapped_column(
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
)
id: Mapped[str] = mapped_column( id: Mapped[str] = mapped_column(
String, primary_key=True, default=lambda: str(uuid.uuid4()) String(DEFAULT_MAX_VARCHAR_LENGTH),
primary_key=True,
default=lambda: str(uuid.uuid4()),
) )
state: Mapped[MutableDict[str, Any]] = mapped_column( state: Mapped[MutableDict[str, Any]] = mapped_column(
@ -125,16 +135,27 @@ class StorageSession(Base):
class StorageEvent(Base): class StorageEvent(Base):
"""Represents an event stored in the database.""" """Represents an event stored in the database."""
__tablename__ = "events" __tablename__ = "events"
id: Mapped[str] = mapped_column(String, primary_key=True) id: Mapped[str] = mapped_column(
app_name: Mapped[str] = mapped_column(String, primary_key=True) String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
user_id: Mapped[str] = mapped_column(String, primary_key=True) )
session_id: Mapped[str] = mapped_column(String, primary_key=True) app_name: Mapped[str] = mapped_column(
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
)
user_id: Mapped[str] = mapped_column(
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
)
session_id: Mapped[str] = mapped_column(
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
)
invocation_id: Mapped[str] = mapped_column(String) invocation_id: Mapped[str] = mapped_column(String(DEFAULT_MAX_VARCHAR_LENGTH))
author: Mapped[str] = mapped_column(String) author: Mapped[str] = mapped_column(String(DEFAULT_MAX_VARCHAR_LENGTH))
branch: Mapped[str] = mapped_column(String, nullable=True) branch: Mapped[str] = mapped_column(
String(DEFAULT_MAX_VARCHAR_LENGTH), nullable=True
)
timestamp: Mapped[DateTime] = mapped_column(DateTime(), default=func.now()) timestamp: Mapped[DateTime] = mapped_column(DateTime(), default=func.now())
content: Mapped[dict[str, Any]] = mapped_column(DynamicJSON, nullable=True) content: Mapped[dict[str, Any]] = mapped_column(DynamicJSON, nullable=True)
actions: Mapped[MutableDict[str, Any]] = mapped_column(PickleType) actions: Mapped[MutableDict[str, Any]] = mapped_column(PickleType)
@ -147,8 +168,10 @@ class StorageEvent(Base):
) )
partial: Mapped[bool] = mapped_column(Boolean, nullable=True) partial: Mapped[bool] = mapped_column(Boolean, nullable=True)
turn_complete: Mapped[bool] = mapped_column(Boolean, nullable=True) turn_complete: Mapped[bool] = mapped_column(Boolean, nullable=True)
error_code: Mapped[str] = mapped_column(String, nullable=True) error_code: Mapped[str] = mapped_column(
error_message: Mapped[str] = mapped_column(String, nullable=True) String(DEFAULT_MAX_VARCHAR_LENGTH), nullable=True
)
error_message: Mapped[str] = mapped_column(String(1024), nullable=True)
interrupted: Mapped[bool] = mapped_column(Boolean, nullable=True) interrupted: Mapped[bool] = mapped_column(Boolean, nullable=True)
storage_session: Mapped[StorageSession] = relationship( storage_session: Mapped[StorageSession] = relationship(
@ -182,9 +205,12 @@ class StorageEvent(Base):
class StorageAppState(Base): class StorageAppState(Base):
"""Represents an app state stored in the database.""" """Represents an app state stored in the database."""
__tablename__ = "app_states" __tablename__ = "app_states"
app_name: Mapped[str] = mapped_column(String, primary_key=True) app_name: Mapped[str] = mapped_column(
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
)
state: Mapped[MutableDict[str, Any]] = mapped_column( state: Mapped[MutableDict[str, Any]] = mapped_column(
MutableDict.as_mutable(DynamicJSON), default={} MutableDict.as_mutable(DynamicJSON), default={}
) )
@ -192,13 +218,20 @@ class StorageAppState(Base):
DateTime(), default=func.now(), onupdate=func.now() DateTime(), default=func.now(), onupdate=func.now()
) )
class StorageUserState(Base): class StorageUserState(Base):
"""Represents a user state stored in the database.""" """Represents a user state stored in the database."""
__tablename__ = "user_states" __tablename__ = "user_states"
app_name: Mapped[str] = mapped_column(String, primary_key=True) app_name: Mapped[str] = mapped_column(
user_id: Mapped[str] = mapped_column(String, primary_key=True) String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
)
user_id: Mapped[str] = mapped_column(
String(DEFAULT_MAX_VARCHAR_LENGTH), primary_key=True
)
state: Mapped[MutableDict[str, Any]] = mapped_column(
MutableDict.as_mutable(DynamicJSON), default={}
)
state: Mapped[MutableDict[str, Any]] = mapped_column( state: Mapped[MutableDict[str, Any]] = mapped_column(
MutableDict.as_mutable(DynamicJSON), default={} MutableDict.as_mutable(DynamicJSON), default={}
) )