|
| 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