feat: Add SqliteSessionService and a migration script to migrate existing DB using DatabaseSessionService to SqliteSessionService

The new Sqlite version has fixed schema and use a single column to store Event data, this should avoid DB migration for future add to the Event object.

- This change introduces `SqliteSessionService`, an asynchronous session service using `aiosqlite` that stores event data as JSON within SQLite.
- A migration script, `migrate_from_sqlalchemy_sqlite.py`, is included to transition data from the older SQLAlchemy-based SQLite schema to this new format.
- The CLI service registry is updated to use SqliteSessionService for sqlite:// URIs.
- Throw error when user trying to access a legacy DB and advice the user to do the migration.

Co-authored-by: Shangjie Chen <deanchen@google.com>
PiperOrigin-RevId: 829971174
This commit is contained in:
Shangjie Chen
2025-11-08 21:19:38 -08:00
committed by Copybara-Service
parent 50ceda00cf
commit e218254495
6 changed files with 850 additions and 48 deletions
+1
View File
@@ -26,6 +26,7 @@ classifiers = [ # List of https://pypi.org/classifiers/
dependencies = [
# go/keep-sorted start
"PyYAML>=6.0.2, <7.0.0", # For APIHubToolset.
"aiosqlite>=0.21.0", # For SQLite database
"anyio>=4.9.0, <5.0.0;python_version>='3.10'", # For MCP Session Manager
"authlib>=1.5.1, <2.0.0", # For RestAPI Tool
"click>=8.1.8, <9.0.0", # For CLI tools
+16 -1
View File
@@ -22,6 +22,7 @@ from urllib.parse import urlparse
from ..artifacts.base_artifact_service import BaseArtifactService
from ..memory.base_memory_service import BaseMemoryService
from ..sessions import InMemorySessionService
from ..sessions.base_session_service import BaseSessionService
@@ -170,8 +171,22 @@ def _register_builtin_services(registry: ServiceRegistry) -> None:
kwargs_copy.pop("agents_dir", None)
return DatabaseSessionService(db_url=uri, **kwargs_copy)
def sqlite_session_factory(uri: str, **kwargs):
from ..sessions.sqlite_session_service import SqliteSessionService
parsed = urlparse(uri)
db_path = parsed.path
if not db_path:
return InMemorySessionService()
elif db_path.startswith("/"):
db_path = db_path[1:]
kwargs_copy = kwargs.copy()
kwargs_copy.pop("agents_dir", None)
return SqliteSessionService(db_path=db_path, **kwargs_copy)
registry.register_session_service("agentengine", agentengine_session_factory)
for scheme in ["sqlite", "postgresql", "mysql"]:
registry.register_session_service("sqlite", sqlite_session_factory)
for scheme in ["postgresql", "mysql"]:
registry.register_session_service(scheme, database_session_factory)
# -- Artifact Services --
@@ -0,0 +1,164 @@
# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Migration script from SQLAlchemy SQLite to the new SQLite JSON schema."""
from __future__ import annotations
import argparse
from datetime import timezone
import json
import logging
import sqlite3
import sys
from google.adk.sessions import database_session_service as dss
from google.adk.sessions import sqlite_session_service as sss
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
logger = logging.getLogger("google_adk." + __name__)
def migrate(source_db_url: str, dest_db_path: str):
"""Migrates data from a SQLAlchemy-based SQLite DB to the new schema."""
logger.info(f"Connecting to source database: {source_db_url}")
try:
engine = create_engine(source_db_url)
dss.Base.metadata.create_all(engine) # Ensure tables exist for inspection
SourceSession = sessionmaker(bind=engine)
source_session = SourceSession()
except Exception as e:
logger.error(f"Failed to connect to source database: {e}")
sys.exit(1)
logger.info(f"Connecting to destination database: {dest_db_path}")
try:
dest_conn = sqlite3.connect(dest_db_path)
dest_cursor = dest_conn.cursor()
dest_cursor.execute(sss.PRAGMA_FOREIGN_KEYS)
dest_cursor.executescript(sss.CREATE_SCHEMA_SQL)
except Exception as e:
logger.error(f"Failed to connect to destination database: {e}")
sys.exit(1)
try:
# Migrate app_states
logger.info("Migrating app_states...")
app_states = source_session.query(dss.StorageAppState).all()
for item in app_states:
dest_cursor.execute(
"INSERT INTO app_states (app_name, state, update_time) VALUES (?,"
" ?, ?)",
(
item.app_name,
json.dumps(item.state),
item.update_time.replace(tzinfo=timezone.utc).timestamp(),
),
)
logger.info(f"Migrated {len(app_states)} app_states.")
# Migrate user_states
logger.info("Migrating user_states...")
user_states = source_session.query(dss.StorageUserState).all()
for item in user_states:
dest_cursor.execute(
"INSERT INTO user_states (app_name, user_id, state, update_time)"
" VALUES (?, ?, ?, ?)",
(
item.app_name,
item.user_id,
json.dumps(item.state),
item.update_time.replace(tzinfo=timezone.utc).timestamp(),
),
)
logger.info(f"Migrated {len(user_states)} user_states.")
# Migrate sessions
logger.info("Migrating sessions...")
sessions = source_session.query(dss.StorageSession).all()
for item in sessions:
dest_cursor.execute(
"INSERT INTO sessions (app_name, user_id, id, state, create_time,"
" update_time) VALUES (?, ?, ?, ?, ?, ?)",
(
item.app_name,
item.user_id,
item.id,
json.dumps(item.state),
item.create_time.replace(tzinfo=timezone.utc).timestamp(),
item.update_time.replace(tzinfo=timezone.utc).timestamp(),
),
)
logger.info(f"Migrated {len(sessions)} sessions.")
# Migrate events
logger.info("Migrating events...")
events = source_session.query(dss.StorageEvent).all()
for item in events:
try:
event_obj = item.to_event()
event_data = event_obj.model_dump_json(exclude_none=True)
dest_cursor.execute(
"INSERT INTO events (id, app_name, user_id, session_id,"
" invocation_id, timestamp, event_data) VALUES (?, ?, ?, ?, ?,"
" ?, ?)",
(
event_obj.id,
item.app_name,
item.user_id,
item.session_id,
event_obj.invocation_id,
event_obj.timestamp,
event_data,
),
)
except Exception as e:
logger.warning(f"Failed to migrate event {item.id}: {e}")
logger.info(f"Migrated {len(events)} events.")
dest_conn.commit()
logger.info("Migration completed successfully.")
except Exception as e:
logger.error(f"An error occurred during migration: {e}", exc_info=True)
dest_conn.rollback()
sys.exit(1)
finally:
source_session.close()
dest_conn.close()
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description=(
"Migrate ADK sessions from an existing SQLAlchemy-based "
"SQLite database to a new SQLite database with JSON events."
)
)
parser.add_argument(
"--source_db_path",
required=True,
help="Path to the source SQLite database file (e.g., /path/to/old.db)",
)
parser.add_argument(
"--dest_db_path",
required=True,
help=(
"Path to the destination SQLite database file (e.g., /path/to/new.db)"
),
)
args = parser.parse_args()
source_url = f"sqlite:///{args.source_db_path}"
migrate(source_url, args.dest_db_path)
File diff suppressed because it is too large Load Diff
+7 -5
View File
@@ -27,6 +27,9 @@ def mock_services():
patch(
"google.adk.sessions.database_session_service.DatabaseSessionService"
) as mock_db_session,
patch(
"google.adk.sessions.sqlite_session_service.SqliteSessionService"
) as mock_sqlite_session,
patch(
"google.adk.artifacts.gcs_artifact_service.GcsArtifactService"
) as mock_gcs_artifact,
@@ -40,6 +43,7 @@ def mock_services():
yield {
"vertex_session": mock_vertex_session,
"db_session": mock_db_session,
"sqlite_session": mock_sqlite_session,
"gcs_artifact": mock_gcs_artifact,
"rag_memory": mock_rag_memory,
"agentengine_memory": mock_agentengine_memory,
@@ -56,17 +60,15 @@ def registry():
# Session Service Tests
def test_create_session_service_sqlite(registry, mock_services):
registry.create_session_service("sqlite:///test.db")
mock_services["db_session"].assert_called_once_with(
db_url="sqlite:///test.db"
)
mock_services["sqlite_session"].assert_called_once_with(db_path="test.db")
def test_create_session_service_sqlite_with_kwargs(registry, mock_services):
registry.create_session_service(
"sqlite:///test.db", pool_size=10, agents_dir="foo"
)
mock_services["db_session"].assert_called_once_with(
db_url="sqlite:///test.db", pool_size=10
mock_services["sqlite_session"].assert_called_once_with(
db_path="test.db", pool_size=10
)
+123 -42
View File
@@ -22,6 +22,7 @@ from google.adk.events.event_actions import EventActions
from google.adk.sessions.base_session_service import GetSessionConfig
from google.adk.sessions.database_session_service import DatabaseSessionService
from google.adk.sessions.in_memory_session_service import InMemorySessionService
from google.adk.sessions.sqlite_session_service import SqliteSessionService
from google.genai import types
import pytest
@@ -29,23 +30,32 @@ import pytest
class SessionServiceType(enum.Enum):
IN_MEMORY = 'IN_MEMORY'
DATABASE = 'DATABASE'
SQLITE = 'SQLITE'
def get_session_service(
service_type: SessionServiceType = SessionServiceType.IN_MEMORY,
tmp_path=None,
):
"""Creates a session service for testing."""
if service_type == SessionServiceType.DATABASE:
return DatabaseSessionService('sqlite:///:memory:')
if service_type == SessionServiceType.SQLITE:
return SqliteSessionService(str(tmp_path / 'sqlite.db'))
return InMemorySessionService()
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_get_empty_session(service_type):
session_service = get_session_service(service_type)
async def test_get_empty_session(service_type, tmp_path):
session_service = get_session_service(service_type, tmp_path)
assert not await session_service.get_session(
app_name='my_app', user_id='test_user', session_id='123'
)
@@ -53,10 +63,15 @@ async def test_get_empty_session(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_create_get_session(service_type):
session_service = get_session_service(service_type)
async def test_create_get_session(service_type, tmp_path):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
user_id = 'test_user'
state = {'key': 'value'}
@@ -97,10 +112,15 @@ async def test_create_get_session(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_create_and_list_sessions(service_type):
session_service = get_session_service(service_type)
async def test_create_and_list_sessions(service_type, tmp_path):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
user_id = 'test_user'
@@ -125,10 +145,15 @@ async def test_create_and_list_sessions(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_list_sessions_all_users(service_type):
session_service = get_session_service(service_type)
async def test_list_sessions_all_users(service_type, tmp_path):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
user_id_1 = 'user1'
user_id_2 = 'user2'
@@ -185,10 +210,15 @@ async def test_list_sessions_all_users(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_app_state_is_shared_by_all_users_of_app(service_type):
session_service = get_session_service(service_type)
async def test_app_state_is_shared_by_all_users_of_app(service_type, tmp_path):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
# User 1 creates a session, establishing app:k1
session1 = await session_service.create_session(
@@ -218,10 +248,17 @@ async def test_app_state_is_shared_by_all_users_of_app(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_user_state_is_shared_only_by_user_sessions(service_type):
session_service = get_session_service(service_type)
async def test_user_state_is_shared_only_by_user_sessions(
service_type, tmp_path
):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
# User 1 creates a session, establishing user:k1 for user 1
session1 = await session_service.create_session(
@@ -250,10 +287,15 @@ async def test_user_state_is_shared_only_by_user_sessions(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_session_state_is_not_shared(service_type):
session_service = get_session_service(service_type)
async def test_session_state_is_not_shared(service_type, tmp_path):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
# User 1 creates a session session1, establishing sk1 only for session1
session1 = await session_service.create_session(
@@ -283,10 +325,17 @@ async def test_session_state_is_not_shared(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_temp_state_is_not_persisted_in_state_or_events(service_type):
session_service = get_session_service(service_type)
async def test_temp_state_is_not_persisted_in_state_or_events(
service_type, tmp_path
):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
user_id = 'u1'
session = await session_service.create_session(
@@ -313,10 +362,15 @@ async def test_temp_state_is_not_persisted_in_state_or_events(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_get_session_respects_user_id(service_type):
session_service = get_session_service(service_type)
async def test_get_session_respects_user_id(service_type, tmp_path):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
# u1 creates session 's1' and adds an event
session1 = await session_service.create_session(
@@ -339,10 +393,17 @@ async def test_get_session_respects_user_id(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_create_session_with_existing_id_raises_error(service_type):
session_service = get_session_service(service_type)
async def test_create_session_with_existing_id_raises_error(
service_type, tmp_path
):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
user_id = 'test_user'
session_id = 'existing_session'
@@ -365,10 +426,15 @@ async def test_create_session_with_existing_id_raises_error(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_append_event_bytes(service_type):
session_service = get_session_service(service_type)
async def test_append_event_bytes(service_type, tmp_path):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
user_id = 'user'
@@ -406,10 +472,15 @@ async def test_append_event_bytes(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_append_event_complete(service_type):
session_service = get_session_service(service_type)
async def test_append_event_complete(service_type, tmp_path):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
user_id = 'user'
@@ -454,10 +525,15 @@ async def test_append_event_complete(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_get_session_with_config(service_type):
session_service = get_session_service(service_type)
async def test_get_session_with_config(service_type, tmp_path):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
user_id = 'user'
@@ -518,10 +594,15 @@ async def test_get_session_with_config(service_type):
@pytest.mark.asyncio
@pytest.mark.parametrize(
'service_type', [SessionServiceType.IN_MEMORY, SessionServiceType.DATABASE]
'service_type',
[
SessionServiceType.IN_MEMORY,
SessionServiceType.DATABASE,
SessionServiceType.SQLITE,
],
)
async def test_partial_events_are_not_persisted(service_type):
session_service = get_session_service(service_type)
async def test_partial_events_are_not_persisted(service_type, tmp_path):
session_service = get_session_service(service_type, tmp_path)
app_name = 'my_app'
user_id = 'user'
session = await session_service.create_session(