diff --git a/core/src/index.ts b/core/src/index.ts index 39cae5d89..113fe93e9 100644 --- a/core/src/index.ts +++ b/core/src/index.ts @@ -40,6 +40,7 @@ export { } from './code_executors/unsafe_local_code_executor.js'; export * from './common.js'; export {DatabaseSessionService} from './sessions/database_session_service.js'; +export {upgradeSessionDatabaseSchema} from './sessions/db/schema_version.js'; export {getSessionServiceFromUri} from './sessions/registry.js'; export {VertexAiSessionService} from './sessions/vertex_ai_session_service.js'; export type { diff --git a/core/src/sessions/database_session_service.ts b/core/src/sessions/database_session_service.ts index 033f05b9a..6c8bc3bec 100644 --- a/core/src/sessions/database_session_service.ts +++ b/core/src/sessions/database_session_service.ts @@ -27,7 +27,6 @@ import { import { ensureDatabaseCreated, getConnectionOptionsFromUri, - validateDatabaseSchemaVersion, } from './db/operations.js'; import { ENTITIES, @@ -36,6 +35,7 @@ import { StorageSession, StorageUserState, } from './db/schema.js'; +import {validateDatabaseSchemaVersion} from './db/schema_version.js'; import {createSession, Session} from './session.js'; import {State} from './state.js'; diff --git a/core/src/sessions/db/README.md b/core/src/sessions/db/README.md new file mode 100644 index 000000000..ccdb1ad2b --- /dev/null +++ b/core/src/sessions/db/README.md @@ -0,0 +1,76 @@ +# Adding a session database schema version + +`DatabaseSessionService` stores its schema version in the `adk_internal_metadata` +table, under the key `schema_version`. `schema_version.ts` owns the negotiation: + +- `LATEST_SCHEMA_VERSION` — the version stamped on a database this build creates. +- `SUPPORTED_SCHEMA_VERSIONS` — the versions this build can read and write. +- `upgradeSessionDatabaseSchema()` — the operator entry point that brings an + older database up to the latest version in place. + +The table name and the key string are a cross-language contract. adk-python reads +and writes the same row, so neither may change. + +## Backward compatibility policy + +adk-python states the policy this package follows: "The `DatabaseSessionService` +is designed to be backward-compatible with the previous schema for a few releases +(at least 2)." A version therefore stays in `SUPPORTED_SCHEMA_VERSIONS` for at +least two releases after it stops being the latest. + +## Steps to add version N+1 + +Assume the current version is `1` and you are adding `2`. + +1. **Keep the change additive.** Add a nullable column, a new table, or a new + index. `ensureDatabaseCreated()` runs `updateSchema({safe: true})`, which + applies exactly those and never drops anything. An additive change is readable + by both the old and the new client. A non-additive change is a much larger job + — see "If the change cannot be additive" below. +2. **Edit the entities** in `schema.ts`. +3. **Add the constant** in `schema_version.ts`: + `export const SCHEMA_VERSION_2 = '2';` +4. **Point `LATEST_SCHEMA_VERSION` at it** and add it to + `SUPPORTED_SCHEMA_VERSIONS`. Keep `'1'` in the set for its deprecation window. +5. **Branch the business logic** in `DatabaseSessionService` on the stored + version if a `'1'` database needs different reads or writes. Use + `readSchemaVersion()`; do not add a second entity set. +6. **Warn on the older version.** `validateDatabaseSchemaVersion()` accepts an + older supported version silently, because there is no supported-but-older + version yet. Add + `if (version !== LATEST_SCHEMA_VERSION) logger.warn(...)` after the + `assertCompatibleVersion()` call, pointing operators at + `upgradeSessionDatabaseSchema()`. +7. **Test the pair.** Cover a `'1'` database opened by the new build, and the + upgrade from `'1'` to `'2'`. +8. **Deprecate.** After at least two releases, remove `'1'` from + `SUPPORTED_SCHEMA_VERSIONS`. A `'1'` database then fails to open with the + error that names `upgradeSessionDatabaseSchema()`. + +## Why there is only one entity set + +adk-python keeps a complete SQLAlchemy model set per version (`schemas/v0.py`, +`schemas/v1.py`) and selects one at runtime with its `_SchemaClasses` helper. +Two declarative metadatas can coexist in one process and either can be bound to +an engine at call time. + +MikroORM cannot do that. `MikroORM.init()` discovers and freezes entity metadata, +and `DatabaseSessionService` passes `entities: ENTITIES` in before init. Selecting +an entity set per database would need, in order: + +1. a throwaway pre-init connection purely to read `adk_internal_metadata`; +2. closing it and re-initialising the real ORM with the selected set; +3. an indirection object over every `em.find` / `em.create` call in + `DatabaseSessionService`, mirroring `_SchemaClasses`; +4. duplicate entity classes with duplicate `tableName` declarations, which + MikroORM's discovery rejects unless the lists are disjoint. + +That cost is paid on every startup, so this package keeps one entity set and +makes each bump additive instead. + +## If the change cannot be additive + +A retyped column, a removed column, or a changed payload encoding cannot be read +by both clients. That is the only case that forces the `_SchemaClasses` +equivalent, at the four costs above. Consider a new table that the old client +ignores before you take that on. diff --git a/core/src/sessions/db/operations.ts b/core/src/sessions/db/operations.ts index 92f7bcb25..8ce9529b8 100644 --- a/core/src/sessions/db/operations.ts +++ b/core/src/sessions/db/operations.ts @@ -6,12 +6,7 @@ import {MikroORM, Options as MikroORMOptions} from '@mikro-orm/core'; import {redactUriPassword} from '../../utils/redact_uri.js'; -import { - ENTITIES, - SCHEMA_VERSION_1_JSON, - SCHEMA_VERSION_KEY, - StorageMetadata, -} from './schema.js'; +import {ENTITIES} from './schema.js'; /** * Parses a database connection URI and returns MikroORM Options. @@ -75,32 +70,3 @@ export async function ensureDatabaseCreated(orm: MikroORM): Promise { // creates tables if they don't exist. Safe mode prevents dropping columns or tables. await orm.schema.updateSchema({safe: true}); } - -/** - * Validates the schema version. - * - * @param orm The MikroORM instance. - * @throws Error if the schema version is not compatible. - */ -export async function validateDatabaseSchemaVersion(orm: MikroORM) { - const em = orm.em.fork(); - const existing = await em.findOne(StorageMetadata, { - key: SCHEMA_VERSION_KEY, - }); - - if (existing) { - if (existing.value !== SCHEMA_VERSION_1_JSON) { - throw new Error( - `ADK Database schema version ${existing.value} is not compatible.`, - ); - } - return; - } - - const newVersion = em.create(StorageMetadata, { - key: SCHEMA_VERSION_KEY, - value: SCHEMA_VERSION_1_JSON, - }); - - await em.persist(newVersion).flush(); -} diff --git a/core/src/sessions/db/schema.ts b/core/src/sessions/db/schema.ts index 4649ae710..2687bbd13 100644 --- a/core/src/sessions/db/schema.ts +++ b/core/src/sessions/db/schema.ts @@ -11,8 +11,6 @@ import { transformToSnakeCaseEvent, } from '../../events/event.js'; -export const SCHEMA_VERSION_KEY = 'schema_version'; -export const SCHEMA_VERSION_1_JSON = '1'; export const STORAGE_KEY_COLUMN_LENGTH = 191; /** diff --git a/core/src/sessions/db/schema_version.ts b/core/src/sessions/db/schema_version.ts new file mode 100644 index 000000000..75cc822d5 --- /dev/null +++ b/core/src/sessions/db/schema_version.ts @@ -0,0 +1,161 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {MikroORM, Options as MikroORMOptions} from '@mikro-orm/core'; +import {logger} from '../../utils/logger.js'; +import { + ensureDatabaseCreated, + getConnectionOptionsFromUri, +} from './operations.js'; +import {ENTITIES, StorageMetadata} from './schema.js'; + +/** Key of the `adk_internal_metadata` row that holds the schema version. */ +export const SCHEMA_VERSION_KEY = 'schema_version'; + +/** JSON-serialized event payload schema, shared with adk-python's v1. */ +export const SCHEMA_VERSION_1_JSON = '1'; + +/** Version stamped on a database this build creates. */ +export const LATEST_SCHEMA_VERSION = SCHEMA_VERSION_1_JSON; + +/** + * Versions this build can read and write. + * + * A release keeps the previous version here for its deprecation window rather + * than removing it, so an existing database keeps opening after a bump. See + * `core/src/sessions/db/README.md`. + */ +export const SUPPORTED_SCHEMA_VERSIONS: ReadonlySet = new Set([ + SCHEMA_VERSION_1_JSON, +]); + +/** + * Throws when the database holds a version this build cannot read. + * + * An absent version row is a new database, which is compatible by definition. + * + * @param version The stored version, or undefined when the row is absent. + * @throws Error naming the accepted versions and both remedies. + */ +function assertCompatibleVersion(version: string | undefined): void { + if (version === undefined || SUPPORTED_SCHEMA_VERSIONS.has(version)) { + return; + } + + const supported = [...SUPPORTED_SCHEMA_VERSIONS].join(', '); + throw new Error( + `ADK Database schema version ${version} is not compatible. ` + + `This build of ADK supports schema version(s) ${supported}. ` + + `Upgrade the @google/adk package if the database was written by a ` + + `newer release, or call upgradeSessionDatabaseSchema() to bring an ` + + `older database up to version ${LATEST_SCHEMA_VERSION}.`, + ); +} + +/** + * Reads the schema version stored in the database. + * + * @param orm The MikroORM instance. + * @returns The stored version, or undefined when the row is absent. + */ +export async function readSchemaVersion( + orm: MikroORM, +): Promise { + const em = orm.em.fork(); + const existing = await em.findOne(StorageMetadata, {key: SCHEMA_VERSION_KEY}); + + return existing?.value; +} + +/** + * Writes the schema version, replacing any value already stored. + * + * @param orm The MikroORM instance. + * @param version The version to store. + */ +export async function stampSchemaVersion( + orm: MikroORM, + version: string, +): Promise { + const em = orm.em.fork(); + const existing = await em.findOne(StorageMetadata, {key: SCHEMA_VERSION_KEY}); + + if (existing) { + existing.value = version; + return em.flush(); + } + + const created = em.create(StorageMetadata, { + key: SCHEMA_VERSION_KEY, + value: version, + }); + + return em.persist(created).flush(); +} + +/** + * Validates the schema version of an open database, stamping a new one. + * + * An existing version row is never rewritten; moving a database to a newer + * version is `upgradeSessionDatabaseSchema`'s job. + * + * @param orm The MikroORM instance. + * @throws Error if the stored version is outside `SUPPORTED_SCHEMA_VERSIONS`. + */ +export async function validateDatabaseSchemaVersion( + orm: MikroORM, +): Promise { + const version = await readSchemaVersion(orm); + assertCompatibleVersion(version); + + if (version === undefined) { + await stampSchemaVersion(orm, LATEST_SCHEMA_VERSION); + } +} + +/** + * Brings a session database up to `LATEST_SCHEMA_VERSION` in place. + * + * Applies the additive DDL this build expects, then stamps the version row. + * Idempotent: a database already at the latest version is left untouched. The + * connection is closed on every exit path. + * + * @param connectionStringOrOptions A database URI, or MikroORM options + * carrying a driver. + * @throws Error if the stored version is outside `SUPPORTED_SCHEMA_VERSIONS`. + */ +export async function upgradeSessionDatabaseSchema( + connectionStringOrOptions: string | MikroORMOptions, +): Promise { + const options = + typeof connectionStringOrOptions === 'string' + ? await getConnectionOptionsFromUri(connectionStringOrOptions) + : {...connectionStringOrOptions, entities: ENTITIES}; + + if (!options.driver) { + throw new Error('Driver is required when passing options object.'); + } + + const orm = await MikroORM.init(options); + + try { + await ensureDatabaseCreated(orm); + + const version = await readSchemaVersion(orm); + assertCompatibleVersion(version); + + if (version === LATEST_SCHEMA_VERSION) { + logger.debug( + `Session database is already at schema version ${LATEST_SCHEMA_VERSION}.`, + ); + return; + } + + await stampSchemaVersion(orm, LATEST_SCHEMA_VERSION); + } finally { + await orm.close(); + } +} diff --git a/core/test/sessions/database_session_service_test.ts b/core/test/sessions/database_session_service_test.ts index 76fbe15ce..a2c16720d 100644 --- a/core/test/sessions/database_session_service_test.ts +++ b/core/test/sessions/database_session_service_test.ts @@ -15,7 +15,7 @@ import {MikroORM} from '@mikro-orm/core'; import {SqliteDriver} from '@mikro-orm/sqlite'; import {afterEach, beforeEach, describe, expect, it} from 'vitest'; import {isDatabaseConnectionString} from '../../src/sessions/database_session_service.js'; -import {validateDatabaseSchemaVersion} from '../../src/sessions/db/operations.js'; +import {validateDatabaseSchemaVersion} from '../../src/sessions/db/schema_version.js'; describe('DatabaseSessionService', () => { let service: DatabaseSessionService; diff --git a/core/test/sessions/db/operations_test.ts b/core/test/sessions/db/operations_test.ts index 041f41e80..3b7925c0b 100644 --- a/core/test/sessions/db/operations_test.ts +++ b/core/test/sessions/db/operations_test.ts @@ -6,16 +6,13 @@ import {MikroORM} from '@mikro-orm/core'; import {SqliteDriver} from '@mikro-orm/sqlite'; -import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest'; +import {afterEach, describe, expect, it, vi} from 'vitest'; import { ensureDatabaseCreated, getConnectionOptionsFromUri, - validateDatabaseSchemaVersion, } from '../../../src/sessions/db/operations.js'; import { ENTITIES, - SCHEMA_VERSION_1_JSON, - SCHEMA_VERSION_KEY, STORAGE_KEY_COLUMN_LENGTH, StorageEvent, StorageMetadata, @@ -163,59 +160,4 @@ describe('operations', () => { await expect(ensureDatabaseCreated(orm)).resolves.not.toThrow(); }); }); - - describe('validateDatabaseSchemaVersion', () => { - let orm: MikroORM; - - beforeEach(async () => { - orm = await MikroORM.init({ - dbName: ':memory:', - driver: SqliteDriver, - entities: [StorageMetadata], - }); - // Ensure schema is updated so StorageMetadata table exists - await orm.schema.updateSchema(); - }); - - afterEach(async () => { - await orm.close(); - }); - - it('should initialize schema version if missing', async () => { - const em = orm.em.fork(); - const initial = await em.find(StorageMetadata, {}); - expect(initial.length).toBe(0); - - await validateDatabaseSchemaVersion(orm); - - const after = await em.find(StorageMetadata, {}); - expect(after.length).toBe(1); - expect(after[0].key).toBe(SCHEMA_VERSION_KEY); - expect(after[0].value).toBe(SCHEMA_VERSION_1_JSON); - }); - - it('should do nothing if schema version is correct', async () => { - const em = orm.em.fork(); - const version = em.create(StorageMetadata, { - key: SCHEMA_VERSION_KEY, - value: SCHEMA_VERSION_1_JSON, - }); - await em.persist(version).flush(); - - await expect(validateDatabaseSchemaVersion(orm)).resolves.not.toThrow(); - }); - - it('should throw error if schema version is incompatible', async () => { - const em = orm.em.fork(); - const version = em.create(StorageMetadata, { - key: SCHEMA_VERSION_KEY, - value: '999', - }); - await em.persist(version).flush(); - - await expect(validateDatabaseSchemaVersion(orm)).rejects.toThrow( - 'ADK Database schema version 999 is not compatible', - ); - }); - }); }); diff --git a/core/test/sessions/db/schema_version_test.ts b/core/test/sessions/db/schema_version_test.ts new file mode 100644 index 000000000..e8fa935d4 --- /dev/null +++ b/core/test/sessions/db/schema_version_test.ts @@ -0,0 +1,271 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {upgradeSessionDatabaseSchema} from '@google/adk'; +import {MikroORM} from '@mikro-orm/core'; +import {SqliteDriver} from '@mikro-orm/sqlite'; +import {mkdtemp, rm} from 'node:fs/promises'; +import {tmpdir} from 'node:os'; +import {join} from 'node:path'; +import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest'; +import { + ENTITIES, + StorageMetadata, + StorageSession, +} from '../../../src/sessions/db/schema.js'; +import { + LATEST_SCHEMA_VERSION, + readSchemaVersion, + SCHEMA_VERSION_1_JSON, + SCHEMA_VERSION_KEY, + stampSchemaVersion, + SUPPORTED_SCHEMA_VERSIONS, + validateDatabaseSchemaVersion, +} from '../../../src/sessions/db/schema_version.js'; + +const INCOMPATIBLE_VERSION = '999'; + +describe('schema_version', () => { + describe('constants', () => { + it('accepts the version it stamps', () => { + expect(SUPPORTED_SCHEMA_VERSIONS.has(LATEST_SCHEMA_VERSION)).toBe(true); + }); + + it('pins the metadata row shared with adk-python', () => { + expect(SCHEMA_VERSION_KEY).toBe('schema_version'); + expect(SCHEMA_VERSION_1_JSON).toBe('1'); + }); + }); + + describe('against an open database', () => { + let orm: MikroORM; + + beforeEach(async () => { + orm = await MikroORM.init({ + dbName: ':memory:', + driver: SqliteDriver, + entities: [StorageMetadata], + }); + // Ensure schema is updated so StorageMetadata table exists + await orm.schema.updateSchema(); + }); + + afterEach(async () => { + await orm.close(); + }); + + async function metadataRows(): Promise { + return orm.em.fork().find(StorageMetadata, {}); + } + + describe('readSchemaVersion', () => { + it('resolves undefined when the version row is absent', async () => { + await expect(readSchemaVersion(orm)).resolves.toBeUndefined(); + }); + }); + + describe('stampSchemaVersion', () => { + it('stores a version that readSchemaVersion reads back', async () => { + await stampSchemaVersion(orm, SCHEMA_VERSION_1_JSON); + + await expect(readSchemaVersion(orm)).resolves.toBe( + SCHEMA_VERSION_1_JSON, + ); + expect(await metadataRows()).toHaveLength(1); + }); + + it('replaces the stored value instead of adding a row', async () => { + await stampSchemaVersion(orm, '1'); + await stampSchemaVersion(orm, '2'); + + const rows = await metadataRows(); + expect(rows).toHaveLength(1); + expect(rows[0].value).toBe('2'); + }); + }); + + describe('validateDatabaseSchemaVersion', () => { + it('should initialize schema version if missing', async () => { + const em = orm.em.fork(); + const initial = await em.find(StorageMetadata, {}); + expect(initial.length).toBe(0); + + await validateDatabaseSchemaVersion(orm); + + const after = await em.find(StorageMetadata, {}); + expect(after.length).toBe(1); + expect(after[0].key).toBe(SCHEMA_VERSION_KEY); + expect(after[0].value).toBe(SCHEMA_VERSION_1_JSON); + }); + + it('should do nothing if schema version is correct', async () => { + const em = orm.em.fork(); + const version = em.create(StorageMetadata, { + key: SCHEMA_VERSION_KEY, + value: SCHEMA_VERSION_1_JSON, + }); + await em.persist(version).flush(); + + await expect(validateDatabaseSchemaVersion(orm)).resolves.not.toThrow(); + }); + + it('should throw error if schema version is incompatible', async () => { + const em = orm.em.fork(); + const version = em.create(StorageMetadata, { + key: SCHEMA_VERSION_KEY, + value: '999', + }); + await em.persist(version).flush(); + + await expect(validateDatabaseSchemaVersion(orm)).rejects.toThrow( + 'ADK Database schema version 999 is not compatible', + ); + }); + + it('names the accepted versions and the upgrade step when it throws', async () => { + const em = orm.em.fork(); + await em + .persist( + em.create(StorageMetadata, { + key: SCHEMA_VERSION_KEY, + value: INCOMPATIBLE_VERSION, + }), + ) + .flush(); + + await expect(validateDatabaseSchemaVersion(orm)).rejects.toThrow( + /supports schema version\(s\) 1\..*upgradeSessionDatabaseSchema\(\)/s, + ); + }); + + it('keeps a single version row when called twice', async () => { + await validateDatabaseSchemaVersion(orm); + await validateDatabaseSchemaVersion(orm); + + expect(await metadataRows()).toHaveLength(1); + }); + }); + }); + + describe('upgradeSessionDatabaseSchema', () => { + let directory: string; + let dbPath: string; + + beforeEach(async () => { + directory = await mkdtemp(join(tmpdir(), 'adk-schema-version-')); + dbPath = join(directory, 'sessions.db'); + }); + + afterEach(async () => { + vi.restoreAllMocks(); + await rm(directory, {recursive: true, force: true}); + }); + + function openDatabase(): Promise { + return MikroORM.init({ + dbName: dbPath, + driver: SqliteDriver, + entities: ENTITIES, + }); + } + + async function readDatabase( + read: (orm: MikroORM) => Promise, + ): Promise { + const orm = await openDatabase(); + try { + return await read(orm); + } finally { + await orm.close(); + } + } + + async function seedVersion(version: string): Promise { + const orm = await openDatabase(); + try { + await orm.schema.updateSchema(); + const em = orm.em.fork(); + await em + .persist( + em.create(StorageMetadata, { + key: SCHEMA_VERSION_KEY, + value: version, + }), + ) + .flush(); + } finally { + await orm.close(); + } + } + + function storedVersions(): Promise { + return readDatabase((orm) => orm.em.fork().find(StorageMetadata, {})); + } + + it('creates the tables and stamps a database that does not exist yet', async () => { + await upgradeSessionDatabaseSchema(`sqlite://${dbPath}`); + + const rows = await storedVersions(); + expect(rows).toHaveLength(1); + expect(rows[0].key).toBe(SCHEMA_VERSION_KEY); + expect(rows[0].value).toBe(LATEST_SCHEMA_VERSION); + await expect( + readDatabase((orm) => orm.em.fork().find(StorageSession, {})), + ).resolves.toEqual([]); + }); + + it('leaves a database that is already at the latest version untouched', async () => { + await upgradeSessionDatabaseSchema(`sqlite://${dbPath}`); + await upgradeSessionDatabaseSchema(`sqlite://${dbPath}`); + + const rows = await storedVersions(); + expect(rows).toHaveLength(1); + expect(rows[0].value).toBe(LATEST_SCHEMA_VERSION); + }); + + it('rejects a version outside the accepted set', async () => { + await seedVersion(INCOMPATIBLE_VERSION); + + await expect( + upgradeSessionDatabaseSchema(`sqlite://${dbPath}`), + ).rejects.toThrow( + `ADK Database schema version ${INCOMPATIBLE_VERSION} is not compatible`, + ); + const rows = await storedVersions(); + expect(rows[0].value).toBe(INCOMPATIBLE_VERSION); + }); + + it('accepts an options object carrying a driver', async () => { + await upgradeSessionDatabaseSchema({ + dbName: dbPath, + driver: SqliteDriver, + }); + + const rows = await storedVersions(); + expect(rows).toHaveLength(1); + expect(rows[0].value).toBe(LATEST_SCHEMA_VERSION); + }); + + it('rejects an options object without a driver', async () => { + await expect( + upgradeSessionDatabaseSchema({dbName: dbPath}), + ).rejects.toThrow('Driver is required when passing options object.'); + }); + + it('closes the connection when the stored version is rejected', async () => { + await seedVersion(INCOMPATIBLE_VERSION); + const orm = await openDatabase(); + const close = vi.spyOn(orm, 'close'); + vi.spyOn(MikroORM, 'init').mockResolvedValue(orm); + + await expect( + upgradeSessionDatabaseSchema(`sqlite://${dbPath}`), + ).rejects.toThrow('is not compatible'); + + expect(close).toHaveBeenCalledTimes(1); + }); + }); +});