From d0aef8a4fadc1c435ec75afecebd84f9ab0ffd7a Mon Sep 17 00:00:00 2001 From: Liang Wu Date: Thu, 15 Jan 2026 17:12:33 -0800 Subject: [PATCH] docs: Update the migration guide for adding a new database session schema Co-authored-by: Liang Wu PiperOrigin-RevId: 856882485 --- src/google/adk/sessions/migration/README.md | 72 +++++++++++++-------- 1 file changed, 46 insertions(+), 26 deletions(-) diff --git a/src/google/adk/sessions/migration/README.md b/src/google/adk/sessions/migration/README.md index 6a907953..77fb5fbe 100644 --- a/src/google/adk/sessions/migration/README.md +++ b/src/google/adk/sessions/migration/README.md @@ -1,4 +1,4 @@ -# Process for a New Schema Version +# Process for Adding a New Schema Version This document outlines the steps required to introduce a new database schema version for `DatabaseSessionService`. Let's assume you are introducing schema @@ -6,18 +6,18 @@ version `2.0`, migrating from `1.0`. ## 1. Update SQLAlchemy Models -Modify the SQLAlchemy model classes (`StorageSession`, `StorageEvent`, -`StorageAppState`, `StorageUserState`, `StorageMetadata`) in -`database_session_service.py` to reflect the new `2.0` schema. This could -involve adding new `mapped_column` definitions, changing types, or adding new -classes for new tables. +Fork from the latest schema version in `google/adk/sessions/schemas/` folder and +modify the SQLAlchemy model classes (`StorageSession`, `StorageEvent`, +`StorageAppState`, `StorageUserState`, `StorageMetadata`) to reflect the new +`2.0` schema, call it `v2.py`. Changes might be adding new `mapped_column` +definitions, changing types, or adding new classes for new tables. ## 2. Create a New Migration Script You need to create a script that migrates data from schema `1.0` to `2.0`. * Create a new file, for example: - `google/adk/sessions/migration/migrate_1_0_to_2_0.py`. + `google/adk/sessions/migration/migrate_from_1_0_to_2_0.py`. * This script must contain a `migrate(source_db_url: str, dest_db_url: str)` function, similar to `migrate_from_sqlalchemy_pickle.py`. * Inside this function: @@ -25,7 +25,7 @@ You need to create a script that migrates data from schema `1.0` to `2.0`. engines using SQLAlchemy. * **Important**: Create the tables in the destination database using the new 2.0 schema definition by calling - `dss.Base.metadata.create_all(dest_engine)`. + `v2.Base.metadata.create_all(dest_engine)`. * Read data from the source tables (schema 1.0). The recommended way to do this without relying on outdated models is to use `sqlalchemy.text`, like: @@ -38,19 +38,19 @@ You need to create a script that migrates data from schema `1.0` to `2.0`. * For each row read from the source, transform the data as necessary to fit the `2.0` schema, and create an instance of the corresponding new - SQLAlchemy model (e.g., `dss.StorageSession(...)`). + SQLAlchemy model (e.g., `v2.StorageSession(...)`). * Add these new `2.0` objects to the destination session, ideally using `dest_session.merge()` to upsert. * After migrating data for all tables, ensure the destination database is - marked with the new schema version: + marked with the new schema version using the `adk_internal_metadata` + table: ```python - from google.adk.sessions import database_session_service as dss - from google.adk.sessions.migration import _schema_check + from google.adk.sessions.migration import _schema_check_utils ... dest_session.merge( - dss.StorageMetadata( - key=_schema_check.SCHEMA_VERSION_KEY, + v2.StorageMetadata( + key=_schema_check_utils.SCHEMA_VERSION_KEY, value="2.0", ) ) @@ -59,35 +59,38 @@ You need to create a script that migrates data from schema `1.0` to `2.0`. ## 3. Update Schema Version Constant -You need to update `CURRENT_SCHEMA_VERSION` in -`google/adk/sessions/migration/_schema_check.py` to reflect the new version: +You need to add the new version and update `LATEST_SCHEMA_VERSION` in +`google/adk/sessions/migration/_schema_check_utils.py` to reflect the new version: ```python -CURRENT_SCHEMA_VERSION = "2.0" +SCHEMA_VERSION_2_0 = "2.0" +LATEST_SCHEMA_VERSION = SCHEMA_VERSION_2_0 ``` This will also update `LATEST_VERSION` in `migration_runner.py`, as it uses this constant. -## 4. Register the New Migration in Migration Runner +## 4. Register the New Migration Script in Migration Runner In `google/adk/sessions/migration/migration_runner.py`, import your new migration script and add it to the `MIGRATIONS` dictionary. This tells the runner how to get from version `1.0` to `2.0`. For example: ```python -from google.adk.sessions.migration import _schema_check +from google.adk.sessions.migration import _schema_check_utils from google.adk.sessions.migration import migrate_from_sqlalchemy_pickle -from google.adk.sessions.migration import migrate_1_0_to_2_0 +from google.adk.sessions.migration import migrate_from_1_0_to_2_0 ... MIGRATIONS = { - _schema_check.SCHEMA_VERSION_0_1_PICKLE: ( - _schema_check.SCHEMA_VERSION_1_0_JSON, + # Previous migrations + _schema_check_utils.SCHEMA_VERSION_0_PICKLE: ( + _schema_check_utils.SCHEMA_VERSION_1_JSON, migrate_from_sqlalchemy_pickle.migrate, ), - _schema_check.SCHEMA_VERSION_1_0_JSON: ( - "2.0", - migrate_1_0_to_2_0.migrate, + # Your new migration + _schema_check_utils.SCHEMA_VERSION_1_JSON: ( + _schema_check_utils.SCHEMA_VERSION_2_0, + migrate_from_1_0_to_2_0.migrate, ), } ``` @@ -100,10 +103,27 @@ creation), update the methods within `DatabaseSessionService` (`create_session`, `get_session`, `append_event`, etc.) in `database_session_service.py` accordingly. +The `DatabaseSessionService` is designed to be backward-compatible with the +previous schema for a few releases (at least 2). It detects the current database +schema, and if it's using the previous version of schema, it will still function +correctly. But for new databases, it will create tables using the latest schema. +Therefore, you should modify `_prepare_tables` method and the +DatabaseSessionService's methods (`create_session`, `get_session`, +`append_event`, etc.) to branch based on the `_db_schema_version` variable +accordingly. + ## 6. CLI Command Changes No changes are needed for the Click command definition in `cli_tools_click.py`. The `adk migrate session` command calls `migration_runner.upgrade()`, which will now automatically detect the source database version and apply the necessary migration steps (e.g., `0.1 -> 1.0 -> 2.0`, or `1.0 -> 2.0`) to reach -`LATEST_VERSION`. \ No newline at end of file +`LATEST_VERSION`. + +## 7. Deprecate the Previous Schema + +After a few releases (at least 2), remove the logic for the previous schema. +Only use the latest schema in the `DatabaseSessionService`, and raise an +Exception if detecting legacy schema versions. Keep the schema files like +`schemas/v1.py` and the migration scripts for documentation and not-yet-migrated +users. \ No newline at end of file