Skip to content

Commit d0aef8a

Browse files
wuliang229copybara-github
authored andcommitted
docs: Update the migration guide for adding a new database session schema
Co-authored-by: Liang Wu <wuliang@google.com> PiperOrigin-RevId: 856882485
1 parent e162bb8 commit d0aef8a

File tree

1 file changed

+46
-26
lines changed

1 file changed

+46
-26
lines changed
Lines changed: 46 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,31 @@
1-
# Process for a New Schema Version
1+
# Process for Adding a New Schema Version
22

33
This document outlines the steps required to introduce a new database schema
44
version for `DatabaseSessionService`. Let's assume you are introducing schema
55
version `2.0`, migrating from `1.0`.
66

77
## 1. Update SQLAlchemy Models
88

9-
Modify the SQLAlchemy model classes (`StorageSession`, `StorageEvent`,
10-
`StorageAppState`, `StorageUserState`, `StorageMetadata`) in
11-
`database_session_service.py` to reflect the new `2.0` schema. This could
12-
involve adding new `mapped_column` definitions, changing types, or adding new
13-
classes for new tables.
9+
Fork from the latest schema version in `google/adk/sessions/schemas/` folder and
10+
modify the SQLAlchemy model classes (`StorageSession`, `StorageEvent`,
11+
`StorageAppState`, `StorageUserState`, `StorageMetadata`) to reflect the new
12+
`2.0` schema, call it `v2.py`. Changes might be adding new `mapped_column`
13+
definitions, changing types, or adding new classes for new tables.
1414

1515
## 2. Create a New Migration Script
1616

1717
You need to create a script that migrates data from schema `1.0` to `2.0`.
1818

1919
* Create a new file, for example:
20-
`google/adk/sessions/migration/migrate_1_0_to_2_0.py`.
20+
`google/adk/sessions/migration/migrate_from_1_0_to_2_0.py`.
2121
* This script must contain a `migrate(source_db_url: str, dest_db_url: str)`
2222
function, similar to `migrate_from_sqlalchemy_pickle.py`.
2323
* Inside this function:
2424
* Connect to the `source_db_url` (which has schema 1.0) and `dest_db_url`
2525
engines using SQLAlchemy.
2626
* **Important**: Create the tables in the destination database using the
2727
new 2.0 schema definition by calling
28-
`dss.Base.metadata.create_all(dest_engine)`.
28+
`v2.Base.metadata.create_all(dest_engine)`.
2929
* Read data from the source tables (schema 1.0). The recommended way to do
3030
this without relying on outdated models is to use `sqlalchemy.text`,
3131
like:
@@ -38,19 +38,19 @@ You need to create a script that migrates data from schema `1.0` to `2.0`.
3838

3939
* For each row read from the source, transform the data as necessary to
4040
fit the `2.0` schema, and create an instance of the corresponding new
41-
SQLAlchemy model (e.g., `dss.StorageSession(...)`).
41+
SQLAlchemy model (e.g., `v2.StorageSession(...)`).
4242
* Add these new `2.0` objects to the destination session, ideally using
4343
`dest_session.merge()` to upsert.
4444
* After migrating data for all tables, ensure the destination database is
45-
marked with the new schema version:
45+
marked with the new schema version using the `adk_internal_metadata`
46+
table:
4647

4748
```python
48-
from google.adk.sessions import database_session_service as dss
49-
from google.adk.sessions.migration import _schema_check
49+
from google.adk.sessions.migration import _schema_check_utils
5050
...
5151
dest_session.merge(
52-
dss.StorageMetadata(
53-
key=_schema_check.SCHEMA_VERSION_KEY,
52+
v2.StorageMetadata(
53+
key=_schema_check_utils.SCHEMA_VERSION_KEY,
5454
value="2.0",
5555
)
5656
)
@@ -59,35 +59,38 @@ You need to create a script that migrates data from schema `1.0` to `2.0`.
5959

6060
## 3. Update Schema Version Constant
6161

62-
You need to update `CURRENT_SCHEMA_VERSION` in
63-
`google/adk/sessions/migration/_schema_check.py` to reflect the new version:
62+
You need to add the new version and update `LATEST_SCHEMA_VERSION` in
63+
`google/adk/sessions/migration/_schema_check_utils.py` to reflect the new version:
6464

6565
```python
66-
CURRENT_SCHEMA_VERSION = "2.0"
66+
SCHEMA_VERSION_2_0 = "2.0"
67+
LATEST_SCHEMA_VERSION = SCHEMA_VERSION_2_0
6768
```
6869

6970
This will also update `LATEST_VERSION` in `migration_runner.py`, as it uses this
7071
constant.
7172

72-
## 4. Register the New Migration in Migration Runner
73+
## 4. Register the New Migration Script in Migration Runner
7374

7475
In `google/adk/sessions/migration/migration_runner.py`, import your new
7576
migration script and add it to the `MIGRATIONS` dictionary. This tells the
7677
runner how to get from version `1.0` to `2.0`. For example:
7778

7879
```python
79-
from google.adk.sessions.migration import _schema_check
80+
from google.adk.sessions.migration import _schema_check_utils
8081
from google.adk.sessions.migration import migrate_from_sqlalchemy_pickle
81-
from google.adk.sessions.migration import migrate_1_0_to_2_0
82+
from google.adk.sessions.migration import migrate_from_1_0_to_2_0
8283
...
8384
MIGRATIONS = {
84-
_schema_check.SCHEMA_VERSION_0_1_PICKLE: (
85-
_schema_check.SCHEMA_VERSION_1_0_JSON,
85+
# Previous migrations
86+
_schema_check_utils.SCHEMA_VERSION_0_PICKLE: (
87+
_schema_check_utils.SCHEMA_VERSION_1_JSON,
8688
migrate_from_sqlalchemy_pickle.migrate,
8789
),
88-
_schema_check.SCHEMA_VERSION_1_0_JSON: (
89-
"2.0",
90-
migrate_1_0_to_2_0.migrate,
90+
# Your new migration
91+
_schema_check_utils.SCHEMA_VERSION_1_JSON: (
92+
_schema_check_utils.SCHEMA_VERSION_2_0,
93+
migrate_from_1_0_to_2_0.migrate,
9194
),
9295
}
9396
```
@@ -100,10 +103,27 @@ creation), update the methods within `DatabaseSessionService` (`create_session`,
100103
`get_session`, `append_event`, etc.) in `database_session_service.py`
101104
accordingly.
102105

106+
The `DatabaseSessionService` is designed to be backward-compatible with the
107+
previous schema for a few releases (at least 2). It detects the current database
108+
schema, and if it's using the previous version of schema, it will still function
109+
correctly. But for new databases, it will create tables using the latest schema.
110+
Therefore, you should modify `_prepare_tables` method and the
111+
DatabaseSessionService's methods (`create_session`, `get_session`,
112+
`append_event`, etc.) to branch based on the `_db_schema_version` variable
113+
accordingly.
114+
103115
## 6. CLI Command Changes
104116

105117
No changes are needed for the Click command definition in `cli_tools_click.py`.
106118
The `adk migrate session` command calls `migration_runner.upgrade()`, which will
107119
now automatically detect the source database version and apply the necessary
108120
migration steps (e.g., `0.1 -> 1.0 -> 2.0`, or `1.0 -> 2.0`) to reach
109-
`LATEST_VERSION`.
121+
`LATEST_VERSION`.
122+
123+
## 7. Deprecate the Previous Schema
124+
125+
After a few releases (at least 2), remove the logic for the previous schema.
126+
Only use the latest schema in the `DatabaseSessionService`, and raise an
127+
Exception if detecting legacy schema versions. Keep the schema files like
128+
`schemas/v1.py` and the migration scripts for documentation and not-yet-migrated
129+
users.

0 commit comments

Comments
 (0)