Skip to content

Commit 7740113

Browse files
DeanChensjcopybara-github
authored andcommitted
chore: Add migration guide for DatabaseSessionService
Co-authored-by: Shangjie Chen <deanchen@google.com> PiperOrigin-RevId: 839376632
1 parent 0907119 commit 7740113

1 file changed

Lines changed: 109 additions & 0 deletions

File tree

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
# Process for a New Schema Version
2+
3+
This document outlines the steps required to introduce a new database schema
4+
version for `DatabaseSessionService`. Let's assume you are introducing schema
5+
version `2.0`, migrating from `1.0`.
6+
7+
## 1. Update SQLAlchemy Models
8+
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.
14+
15+
## 2. Create a New Migration Script
16+
17+
You need to create a script that migrates data from schema `1.0` to `2.0`.
18+
19+
* Create a new file, for example:
20+
`google/adk/sessions/migration/migrate_1_0_to_2_0.py`.
21+
* This script must contain a `migrate(source_db_url: str, dest_db_url: str)`
22+
function, similar to `migrate_from_sqlalchemy_pickle.py`.
23+
* Inside this function:
24+
* Connect to the `source_db_url` (which has schema 1.0) and `dest_db_url`
25+
engines using SQLAlchemy.
26+
* **Important**: Create the tables in the destination database using the
27+
new 2.0 schema definition by calling
28+
`dss.Base.metadata.create_all(dest_engine)`.
29+
* Read data from the source tables (schema 1.0). The recommended way to do
30+
this without relying on outdated models is to use `sqlalchemy.text`,
31+
like:
32+
33+
```python
34+
from sqlalchemy import text
35+
...
36+
rows = source_session.execute(text("SELECT * FROM sessions")).mappings().all()
37+
```
38+
39+
* For each row read from the source, transform the data as necessary to
40+
fit the `2.0` schema, and create an instance of the corresponding new
41+
SQLAlchemy model (e.g., `dss.StorageSession(...)`).
42+
* Add these new `2.0` objects to the destination session, ideally using
43+
`dest_session.merge()` to upsert.
44+
* After migrating data for all tables, ensure the destination database is
45+
marked with the new schema version:
46+
47+
```python
48+
from google.adk.sessions import database_session_service as dss
49+
from google.adk.sessions.migration import _schema_check
50+
...
51+
dest_session.merge(
52+
dss.StorageMetadata(
53+
key=_schema_check.SCHEMA_VERSION_KEY,
54+
value="2.0",
55+
)
56+
)
57+
dest_session.commit()
58+
```
59+
60+
## 3. Update Schema Version Constant
61+
62+
You need to update `CURRENT_SCHEMA_VERSION` in
63+
`google/adk/sessions/migration/_schema_check.py` to reflect the new version:
64+
65+
```python
66+
CURRENT_SCHEMA_VERSION = "2.0"
67+
```
68+
69+
This will also update `LATEST_VERSION` in `migration_runner.py`, as it uses this
70+
constant.
71+
72+
## 4. Register the New Migration in Migration Runner
73+
74+
In `google/adk/sessions/migration/migration_runner.py`, import your new
75+
migration script and add it to the `MIGRATIONS` dictionary. This tells the
76+
runner how to get from version `1.0` to `2.0`. For example:
77+
78+
```python
79+
from google.adk.sessions.migration import _schema_check
80+
from google.adk.sessions.migration import migrate_from_sqlalchemy_pickle
81+
from google.adk.sessions.migration import migrate_1_0_to_2_0
82+
...
83+
MIGRATIONS = {
84+
_schema_check.SCHEMA_VERSION_0_1_PICKLE: (
85+
_schema_check.SCHEMA_VERSION_1_0_JSON,
86+
migrate_from_sqlalchemy_pickle.migrate,
87+
),
88+
_schema_check.SCHEMA_VERSION_1_0_JSON: (
89+
"2.0",
90+
migrate_1_0_to_2_0.migrate,
91+
),
92+
}
93+
```
94+
95+
## 5. Update `DatabaseSessionService` Business Logic
96+
97+
If your schema change affects how data should be read or written during normal
98+
operation (e.g., you added a new column that needs to be populated on session
99+
creation), update the methods within `DatabaseSessionService` (`create_session`,
100+
`get_session`, `append_event`, etc.) in `database_session_service.py`
101+
accordingly.
102+
103+
## 6. CLI Command Changes
104+
105+
No changes are needed for the Click command definition in `cli_tools_click.py`.
106+
The `adk migrate session` command calls `migration_runner.upgrade()`, which will
107+
now automatically detect the source database version and apply the necessary
108+
migration steps (e.g., `0.1 -> 1.0 -> 2.0`, or `1.0 -> 2.0`) to reach
109+
`LATEST_VERSION`.

0 commit comments

Comments
 (0)