[GitMarco27]

Implement Full async DatabaseSessionService

Target Issue: #1005

Overview

This PR introduces an asynchronous implementation of the DatabaseSessionService with minimal breaking changes. The primary goal is to enable effective use of ADK in fully async environments and API endpoints while avoiding event loop blocking during database I/O operations.

Changes

• Converted DatabaseSessionService to use async/await patterns throughout

Testing Plan

The implementation has been tested following the project's contribution guidelines:

Unit Tests

• All existing unit tests pass successfully
• Minor update to test requirements added to support aiosqlite

Manual End-to-End Testing

• E2E tests performed using:
  • LLM Provider: LiteLLM
  • Database: PostgreSQL with asyncpg driver

from google.adk.sessions.database_session_service import DatabaseSessionService

connection_string: str = (
   "postgresql+asyncpg://PG_USER:PG_PSWD@PG_HOST:5432/PG_DB"
)
session_service: DatabaseSessionService = DatabaseSessionService(
   db_url=connection_string
)

session = await session_service.create_session(
   app_name="test_app", session_id="test_session", user_id="test_user"
)
assert session is not None

sessions = await session_service.list_sessions(app_name="test_app", user_id="test_user")
assert len(sessions.sessions) > 0


session = await session_service.get_session(
   app_name="test_app", session_id="test_session", user_id="test_user"
)
assert session is not None

await session_service.delete_session(
   app_name="test_app", session_id="test_session", user_id="test_user"
)
assert (
   await session_service.get_session(
       app_name="test_app", session_id="test_session", user_id="test_user"
   )
   is None
)

The implementation have been also tested using the following configurations for llm provider and Runner:

def get_azure_openai_model(deployment_id: str | None = None) -> LiteLlm:
    ...

    if not deployment_id:
        deployment_id = os.getenv("AZURE_OPENAI_DEPLOYMENT_ID")

    logger.info(f"Using Azure OpenAI deployment ID: {deployment_id}")

    return LiteLlm(
        model=f"azure/{os.getenv('AZURE_OPENAI_DEPLOYMENT_ID')}",
        stream=True,
    )

...

    @staticmethod
    def _get_runner(agent: Agent) -> Runner:
        storage=DatabaseSessionService(db_url=get_pg_connection_string())
        return Runner(
            agent=agent,
            app_name=APP_NAME,
            session_service=storage,
        )

...

async for event in self.runner.run_async(
    user_id=user_id,
    session_id=session_id,
    new_message=content,
    run_config=(
        RunConfig(
            streaming_mode=StreamingMode.SSE, response_modalities=["TEXT"]
        )
        if stream
        else RunConfig()
    ),
):
    last_event = event
    if stream:
        yield event

...

Breaking Changes

• Database connection string format may need updates for async drivers

[gemini-code-assist]

Summary of Changes

Hello @GitMarco27, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a comprehensive asynchronous implementation for the DatabaseSessionService. The primary motivation is to facilitate the effective use of the ADK in fully asynchronous environments and API endpoints, specifically by preventing event loop blocking during database I/O operations. This change significantly enhances the service's compatibility and performance in modern async Python applications.

Highlights

• Asynchronous DatabaseSessionService: The DatabaseSessionService has been fully converted to use asynchronous patterns, enabling non-blocking database I/O operations in async environments.
• Lazy Table Creation: Database tables are now created lazily upon the first database interaction, rather than during service initialization, improving startup performance.
• Dependency Update: The aiosqlite library has been added as a test dependency to support asynchronous SQLite connections.
• Minimal Breaking Changes: The implementation aims for minimal breaking changes, primarily affecting database connection string formats for async drivers.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request does a great job of converting the DatabaseSessionService to a fully asynchronous implementation. The changes correctly adopt async/await patterns and modern SQLAlchemy 2.0 style for async operations, which will significantly benefit applications running in an async environment. The code is clean and the approach of lazy table creation is sound. I've found one potential race condition in the table creation logic and provided a detailed suggestion to resolve it using an asyncio.Lock. With that fix, this will be a solid enhancement.

[GitMarco27]

@gemini-code-assist review

[gemini-code-assist]

Code Review

This pull request does an excellent job of converting the DatabaseSessionService to a fully asynchronous implementation. The changes correctly adopt sqlalchemy.ext.asyncio patterns, including the use of an async engine, async sessions, and the modern select statement syntax. The introduction of lazy, thread-safe table creation via _ensure_tables_created is a robust addition. My review focuses on performance optimizations, suggesting the use of asyncio.gather to run independent database queries concurrently, which should reduce latency in several key methods. Overall, this is a high-quality and well-executed feature implementation.

[iborko]

@DeanChensj can you prioritize this? Database session service is very important component for production of the systems based on ADK, it's a real shame this part is still sync.

[matthiasak]

+1 here -- this is a mission critical feature

[fabiofumarola]

really waiting to see it integrated!

[mikeedjones]

Thank you @GitMarco27 very much needed.

I took the liberty of pulling the changes out into a AsyncDatabaseSessionService in this gist so it can be used. I think it should "just work". The tests pass using it anyway.

Cheers!

[prasadskarmarkar]

I think this bug should be addressed in the async implementation as well :
#3351

let me know your thoughts.

[hangfei]

Could you resolve the conflicts?

[DeanChensj]

Thanks @GitMarco27 for the PR!

This change seems will break users who does not use async drives, collecting ideas how we make this change less invasive.
In theory developers don't need to migrate the DB schema to use async drive, just that they can not use adk web --session_service_uri=sqlite:///test.db syntax, instead they need to do adk web --session_service_uri=sqlite+aiosqlite:///test.db.

Could you also verify again this works with the latest service_registry refactoring?

[BertrandDechoux]

Breaking changes could be reduced with two implementations (one async and one non-async) or a single double-implementation (supporting both async and non-async).

BUT there is actually not point in having a non-async implementation of DatabaseSessionService : adk is async by design and for good reasons. Finding out that current DatabaseSessionService is not async is basically a bad discovery for every users targeting a non-trivial production environment.

[BertrandDechoux]

Could a collaborator (@hangfei @DeanChensj) add the conflict resolution ?

[Rainelz]

BUT there is actually not point in having a non-async implementation of DatabaseSessionService : adk is async by design and for good reasons. Finding out that current DatabaseSessionService is not async is basically a bad discovery for every users targeting a non-trivial production environment.

I fully agree that moving this to async is the right direction, and it still feels early enough to make this improvement.

For reference, a similar breaking change was handled quickly in #2677 with a simple note in the 1.14.0 release:

Add custom_metadata to DatabaseSessionService fb009d8

NOTE: This requires DB migration, run ALTER TABLE events ADD COLUMN custom_metadata JSON; to migrate existing database tables

This PR has been open for two months now, and I worry that long delays on straightforward fixes can discourage community contributions.

Is there anything blocking this change that we, as contributors, can help with to move it forward?

[DeanChensj]

Thanks all for the comments, we have added a new SqliteSessionService with a fixed schema, and I'll merge this change with conflict resolved.

[adk-bot]

Thank you @GitMarco27 for your contribution! 🎉

Your changes have been successfully imported and merged via Copybara in commit 7495941.

Closing this PR as the changes are now in the main branch.