diff --git a/pyproject.toml b/pyproject.toml index 269e1bc22af..1e9d4ab6518 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -106,6 +106,7 @@ eval = [ test = [ # go/keep-sorted start "a2a-sdk>=0.3.0,<0.4.0;python_version>='3.10'", + "aiosqlite>=0.21.0", "anthropic>=0.43.0", # For anthropic model tests "langchain-community>=0.3.17", "langgraph>=0.2.60, <= 0.4.10", # For LangGraphAgent diff --git a/src/google/adk/sessions/database_session_service.py b/src/google/adk/sessions/database_session_service.py index 959524c689c..fed016a1a3b 100644 --- a/src/google/adk/sessions/database_session_service.py +++ b/src/google/adk/sessions/database_session_service.py @@ -29,20 +29,21 @@ from sqlalchemy import event from sqlalchemy import ForeignKeyConstraint from sqlalchemy import func +from sqlalchemy import select from sqlalchemy import Text from sqlalchemy.dialects import mysql from sqlalchemy.dialects import postgresql -from sqlalchemy.engine import create_engine -from sqlalchemy.engine import Engine from sqlalchemy.exc import ArgumentError +from sqlalchemy.ext.asyncio import async_sessionmaker +from sqlalchemy.ext.asyncio import AsyncEngine +from sqlalchemy.ext.asyncio import AsyncSession +from sqlalchemy.ext.asyncio import create_async_engine from sqlalchemy.ext.mutable import MutableDict from sqlalchemy.inspection import inspect from sqlalchemy.orm import DeclarativeBase from sqlalchemy.orm import Mapped from sqlalchemy.orm import mapped_column from sqlalchemy.orm import relationship -from sqlalchemy.orm import Session as DatabaseSessionFactory -from sqlalchemy.orm import sessionmaker from sqlalchemy.schema import MetaData from sqlalchemy.types import DateTime from sqlalchemy.types import PickleType @@ -385,17 +386,19 @@ class DatabaseSessionService(BaseSessionService): """A session service that uses a database for storage.""" def __init__(self, db_url: str, **kwargs: Any): - """Initializes the database session service with a database URL.""" + """ + Initializes the database session service with a database URL. + + To initialize tables use `initialize_tables` method. + """ # 1. Create DB engine for db connection - # 2. Create all tables based on schema - # 3. Initialize all properties + # 2. Initialize all properties try: - db_engine = create_engine(db_url, **kwargs) + db_engine = create_async_engine(db_url, **kwargs) if db_engine.dialect.name == "sqlite": # Set sqlite pragma to enable foreign keys constraints - event.listen(db_engine, "connect", set_sqlite_pragma) - + event.listen(db_engine.sync_engine, "connect", set_sqlite_pragma) except Exception as e: if isinstance(e, ArgumentError): raise ValueError( @@ -413,18 +416,22 @@ def __init__(self, db_url: str, **kwargs: Any): local_timezone = get_localzone() logger.info("Local timezone: %s", local_timezone) - self.db_engine: Engine = db_engine + self.db_engine: AsyncEngine = db_engine self.metadata: MetaData = MetaData() - self.inspector = inspect(self.db_engine) # DB session factory method - self.database_session_factory: sessionmaker[DatabaseSessionFactory] = ( - sessionmaker(bind=self.db_engine) + self.database_session_factory: async_sessionmaker[AsyncSession] = ( + async_sessionmaker(bind=self.db_engine) ) - # Uncomment to recreate DB every time - # Base.metadata.drop_all(self.db_engine) - Base.metadata.create_all(self.db_engine) + async def initialize_tables(self): + """Initialize database tables. Call this after creating the service.""" + async with self.db_engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + + async def close(self): + """Close the database engine and clean up connections.""" + await self.db_engine.dispose() @override async def create_session( @@ -441,11 +448,11 @@ async def create_session( # 4. Build the session object with generated id # 5. Return the session - with self.database_session_factory() as sql_session: + async with self.database_session_factory() as sql_session: # Fetch app and user states from storage - storage_app_state = sql_session.get(StorageAppState, (app_name)) - storage_user_state = sql_session.get( + storage_app_state = await sql_session.get(StorageAppState, (app_name)) + storage_user_state = await sql_session.get( StorageUserState, (app_name, user_id) ) @@ -485,9 +492,9 @@ async def create_session( state=session_state, ) sql_session.add(storage_session) - sql_session.commit() + await sql_session.commit() - sql_session.refresh(storage_session) + await sql_session.refresh(storage_session) # Merge states for response merged_state = _merge_state(app_state, user_state, session_state) @@ -506,8 +513,8 @@ async def get_session( # 1. Get the storage session entry from session table # 2. Get all the events based on session id and filtering config # 3. Convert and return the session - with self.database_session_factory() as sql_session: - storage_session = sql_session.get( + async with self.database_session_factory() as sql_session: + storage_session = await sql_session.get( StorageSession, (app_name, user_id, session_id) ) if storage_session is None: @@ -519,8 +526,8 @@ async def get_session( else: timestamp_filter = True - storage_events = ( - sql_session.query(StorageEvent) + stmt = ( + select(StorageEvent) .filter(StorageEvent.app_name == app_name) .filter(StorageEvent.session_id == storage_session.id) .filter(StorageEvent.user_id == user_id) @@ -531,12 +538,13 @@ async def get_session( if config and config.num_recent_events else None ) - .all() ) + result = await sql_session.execute(stmt) + storage_events = result.scalars().all() # Fetch states from storage - storage_app_state = sql_session.get(StorageAppState, (app_name)) - storage_user_state = sql_session.get( + storage_app_state = await sql_session.get(StorageAppState, (app_name)) + storage_user_state = await sql_session.get( StorageUserState, (app_name, user_id) ) @@ -557,16 +565,17 @@ async def list_sessions( self, *, app_name: str, user_id: str ) -> ListSessionsResponse: with self.database_session_factory() as sql_session: - results = ( - sql_session.query(StorageSession) + stmt = ( + select(StorageSession) .filter(StorageSession.app_name == app_name) .filter(StorageSession.user_id == user_id) - .all() ) + result = await sql_session.execute(stmt) + results = result.scalars().all() # Fetch states from storage - storage_app_state = sql_session.get(StorageAppState, (app_name)) - storage_user_state = sql_session.get( + storage_app_state = await sql_session.get(StorageAppState, (app_name)) + storage_user_state = await sql_session.get( StorageUserState, (app_name, user_id) ) @@ -585,14 +594,14 @@ async def list_sessions( async def delete_session( self, app_name: str, user_id: str, session_id: str ) -> None: - with self.database_session_factory() as sql_session: + async with self.database_session_factory() as sql_session: stmt = delete(StorageSession).where( StorageSession.app_name == app_name, StorageSession.user_id == user_id, StorageSession.id == session_id, ) - sql_session.execute(stmt) - sql_session.commit() + await sql_session.execute(stmt) + await sql_session.commit() @override async def append_event(self, session: Session, event: Event) -> Event: @@ -603,7 +612,7 @@ async def append_event(self, session: Session, event: Event) -> Event: # 2. Update session attributes based on event config # 3. Store event to table with self.database_session_factory() as sql_session: - storage_session = sql_session.get( + storage_session = await sql_session.get( StorageSession, (session.app_name, session.user_id, session.id) ) @@ -617,8 +626,10 @@ async def append_event(self, session: Session, event: Event) -> Event: ) # Fetch states from storage - storage_app_state = sql_session.get(StorageAppState, (session.app_name)) - storage_user_state = sql_session.get( + storage_app_state = await sql_session.get( + StorageAppState, (session.app_name) + ) + storage_user_state = await sql_session.get( StorageUserState, (session.app_name, session.user_id) ) @@ -649,8 +660,8 @@ async def append_event(self, session: Session, event: Event) -> Event: sql_session.add(StorageEvent.from_event(session, event)) - sql_session.commit() - sql_session.refresh(storage_session) + await sql_session.commit() + await sql_session.refresh(storage_session) # Update timestamp with commit time session.last_update_time = storage_session.update_timestamp_tz diff --git a/tests/unittests/sessions/test_session_service.py b/tests/unittests/sessions/test_session_service.py index c2a3a1d98f5..b0654b7868c 100644 --- a/tests/unittests/sessions/test_session_service.py +++ b/tests/unittests/sessions/test_session_service.py @@ -30,12 +30,13 @@ class SessionServiceType(enum.Enum): DATABASE = 'DATABASE' -def get_session_service( +async def get_session_service( service_type: SessionServiceType = SessionServiceType.IN_MEMORY, ): """Creates a session service for testing.""" if service_type == SessionServiceType.DATABASE: - return DatabaseSessionService('sqlite:///:memory:') + service = DatabaseSessionService('sqlite+aiosqlite:///:memory:') + await service.initialize_tables() return InMemorySessionService() @@ -44,7 +45,7 @@ def get_session_service( 'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE] ) async def test_get_empty_session(service_type): - session_service = get_session_service(service_type) + session_service = await get_session_service(service_type) assert not await session_service.get_session( app_name='my_app', user_id='test_user', session_id='123' ) @@ -55,7 +56,7 @@ async def test_get_empty_session(service_type): 'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE] ) async def test_create_get_session(service_type): - session_service = get_session_service(service_type) + session_service = await get_session_service(service_type) app_name = 'my_app' user_id = 'test_user' state = {'key': 'value'} @@ -99,7 +100,7 @@ async def test_create_get_session(service_type): 'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE] ) async def test_create_and_list_sessions(service_type): - session_service = get_session_service(service_type) + session_service = await get_session_service(service_type) app_name = 'my_app' user_id = 'test_user' @@ -126,7 +127,7 @@ async def test_create_and_list_sessions(service_type): 'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE] ) async def test_session_state(service_type): - session_service = get_session_service(service_type) + session_service = await get_session_service(service_type) app_name = 'my_app' user_id_1 = 'user1' user_id_2 = 'user2' @@ -218,7 +219,7 @@ async def test_session_state(service_type): 'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE] ) async def test_create_new_session_will_merge_states(service_type): - session_service = get_session_service(service_type) + session_service = await get_session_service(service_type) app_name = 'my_app' user_id = 'user' session_id_1 = 'session1' @@ -264,7 +265,7 @@ async def test_create_new_session_will_merge_states(service_type): 'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE] ) async def test_append_event_bytes(service_type): - session_service = get_session_service(service_type) + session_service = await get_session_service(service_type) app_name = 'my_app' user_id = 'user' @@ -305,7 +306,7 @@ async def test_append_event_bytes(service_type): 'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE] ) async def test_append_event_complete(service_type): - session_service = get_session_service(service_type) + session_service = await get_session_service(service_type) app_name = 'my_app' user_id = 'user' @@ -345,7 +346,7 @@ async def test_append_event_complete(service_type): 'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE] ) async def test_get_session_with_config(service_type): - session_service = get_session_service(service_type) + session_service = await get_session_service(service_type) app_name = 'my_app' user_id = 'user' @@ -409,7 +410,7 @@ async def test_get_session_with_config(service_type): 'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE] ) async def test_append_event_with_fields(service_type): - session_service = get_session_service(service_type) + session_service = await get_session_service(service_type) app_name = 'my_app' user_id = 'test_user' session = await session_service.create_session(