diff --git a/core/src/sessions/database_session_service.ts b/core/src/sessions/database_session_service.ts index 033f05b9a..690c30525 100644 --- a/core/src/sessions/database_session_service.ts +++ b/core/src/sessions/database_session_service.ts @@ -481,9 +481,7 @@ export class DatabaseSessionService extends BaseSessionService { } else { const newStorageEvent = txEm.create(StorageEvent, { id: trimmedEvent.id, - appName: session.appName, - userId: session.userId, - sessionId: session.id, + session: storageSession, invocationId: trimmedEvent.invocationId, timestamp: new Date(trimmedEvent.timestamp), eventData: trimmedEvent, diff --git a/core/src/sessions/db/operations.ts b/core/src/sessions/db/operations.ts index ef5897292..4ad0e1b71 100644 --- a/core/src/sessions/db/operations.ts +++ b/core/src/sessions/db/operations.ts @@ -4,7 +4,12 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {MikroORM, Options as MikroORMOptions} from '@mikro-orm/core'; +import { + MikroORM, + Options as MikroORMOptions, + QueryResult, +} from '@mikro-orm/core'; +import {logger} from '../../utils/logger.js'; import { ENTITIES, SCHEMA_VERSION_1_JSON, @@ -12,6 +17,14 @@ import { StorageMetadata, } from './schema.js'; +const DELETE_ORPHANED_EVENTS_SQL = `DELETE FROM events +WHERE NOT EXISTS ( + SELECT 1 FROM sessions + WHERE sessions.app_name = events.app_name + AND sessions.user_id = events.user_id + AND sessions.id = events.session_id +)`; + /** * Parses a database connection URI and returns MikroORM Options. * @@ -61,6 +74,21 @@ export async function getConnectionOptionsFromUri( } as MikroORMOptions; } +/** + * Deletes event rows whose session no longer exists. + * + * @param orm The MikroORM instance. + * @returns Promise + */ +export async function deleteOrphanedEvents(orm: MikroORM): Promise { + const result: QueryResult = await orm.em + .getConnection() + .execute(DELETE_ORPHANED_EVENTS_SQL, [], 'run'); + logger.warn( + `Deleted ${result.affectedRows} event rows whose session no longer exists.`, + ); +} + /** * Creates a database and tables if they don't exist. * @@ -71,8 +99,21 @@ export async function ensureDatabaseCreated(orm: MikroORM): Promise { // creates database if it doesn't exist await orm.schema.ensureDatabase(); - // creates tables if they don't exist. Safe mode prevents dropping columns or tables. - await orm.schema.updateSchema({safe: true}); + try { + // creates tables if they don't exist. Safe mode prevents dropping columns or tables. + await orm.schema.updateSchema({safe: true}); + } catch (error) { + // A database created before the events -> sessions foreign key can hold + // event rows whose session is already gone, and the constraint cannot be + // added while they exist. Those rows are unreachable through the service, + // so drop them and let the schema update finish. + logger.warn( + 'Schema update failed; retrying after deleting orphaned events.', + error, + ); + await deleteOrphanedEvents(orm); + await orm.schema.updateSchema({safe: true}); + } } /** diff --git a/core/src/sessions/db/schema.ts b/core/src/sessions/db/schema.ts index 4649ae710..abbe5942c 100644 --- a/core/src/sessions/db/schema.ts +++ b/core/src/sessions/db/schema.ts @@ -4,7 +4,13 @@ * SPDX-License-Identifier: Apache-2.0 */ -import {Entity, JsonType, PrimaryKey, Property} from '@mikro-orm/core'; +import { + Entity, + JsonType, + ManyToOne, + PrimaryKey, + Property, +} from '@mikro-orm/core'; import { Event, transformToCamelCaseEvent, @@ -96,11 +102,14 @@ export class StorageUserState { [PrimaryKey.name]?: [string, string]; } +/** + * The primary key is declared in the order `app_name, user_id, id` to match + * adk-python's v1 schema. MikroORM binds a composite foreign key in the + * target's declaration order, so this order is what keeps the `events` columns + * laid out as `id, app_name, user_id, session_id`. + */ @Entity({tableName: 'sessions'}) export class StorageSession { - @PrimaryKey({type: 'string', length: STORAGE_KEY_COLUMN_LENGTH}) - id!: string; - @PrimaryKey({ type: 'string', fieldName: 'app_name', @@ -115,6 +124,9 @@ export class StorageSession { }) userId!: string; + @PrimaryKey({type: 'string', length: STORAGE_KEY_COLUMN_LENGTH}) + id!: string; + @Property({type: 'json'}) state!: Record; @@ -135,30 +147,33 @@ export class StorageSession { [PrimaryKey.name]?: [string, string, string]; } +/** + * The `session` relation owns `app_name`, `user_id` and `session_id`, so the + * database enforces that every event belongs to a live session and removes the + * events when that session row is deleted. `appName`, `userId` and `sessionId` + * stay declared as non-persistent mirrors of the same columns so that queries + * can filter on them without loading the relation. + */ @Entity({tableName: 'events'}) export class StorageEvent { @PrimaryKey({type: 'string', length: STORAGE_KEY_COLUMN_LENGTH}) id!: string; - @PrimaryKey({ - type: 'string', - fieldName: 'app_name', - length: STORAGE_KEY_COLUMN_LENGTH, + @ManyToOne(() => StorageSession, { + primary: true, + joinColumns: ['app_name', 'user_id', 'session_id'], + deleteRule: 'cascade', + updateRule: 'cascade', }) + session!: StorageSession; + + @Property({type: 'string', fieldName: 'app_name', persist: false}) appName!: string; - @PrimaryKey({ - type: 'string', - fieldName: 'user_id', - length: STORAGE_KEY_COLUMN_LENGTH, - }) + @Property({type: 'string', fieldName: 'user_id', persist: false}) userId!: string; - @PrimaryKey({ - type: 'string', - fieldName: 'session_id', - length: STORAGE_KEY_COLUMN_LENGTH, - }) + @Property({type: 'string', fieldName: 'session_id', persist: false}) sessionId!: string; @Property({type: 'string', fieldName: 'invocation_id'}) diff --git a/core/test/sessions/db/operations_test.ts b/core/test/sessions/db/operations_test.ts index 041f41e80..d5047279f 100644 --- a/core/test/sessions/db/operations_test.ts +++ b/core/test/sessions/db/operations_test.ts @@ -8,6 +8,7 @@ import {MikroORM} from '@mikro-orm/core'; import {SqliteDriver} from '@mikro-orm/sqlite'; import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest'; import { + deleteOrphanedEvents, ensureDatabaseCreated, getConnectionOptionsFromUri, validateDatabaseSchemaVersion, @@ -16,10 +17,10 @@ import { ENTITIES, SCHEMA_VERSION_1_JSON, SCHEMA_VERSION_KEY, - STORAGE_KEY_COLUMN_LENGTH, StorageEvent, StorageMetadata, } from '../../../src/sessions/db/schema.js'; +import {resetLogger, setLogger} from '../../../src/utils/logger.js'; // Mock dynamic imports for drivers that might not be installed in dev vi.mock('@mikro-orm/postgresql', () => ({ @@ -35,40 +36,65 @@ vi.mock('@mikro-orm/mssql', () => ({ MsSqlDriver: class MockMsSqlDriver {}, })); -describe('operations', () => { - describe('storage schema', () => { - let orm: MikroORM; +const APP_NAME = 'test-app'; +const USER_ID = 'test-user'; +const SESSION_ID = 'test-session'; + +async function insertEventRow( + orm: MikroORM, + id: string, + sessionId: string, +): Promise { + await orm.em + .getConnection() + .execute( + 'INSERT INTO events (id, app_name, user_id, session_id, invocation_id, timestamp, event_data) VALUES (?, ?, ?, ?, ?, ?, ?)', + [id, APP_NAME, USER_ID, sessionId, `invocation-${id}`, Date.now(), '{}'], + ); +} + +async function insertLiveSessionWithEvent(orm: MikroORM): Promise { + const now = Date.now(); + await orm.em + .getConnection() + .execute( + 'INSERT INTO sessions (app_name, user_id, id, state, create_time, update_time) VALUES (?, ?, ?, ?, ?, ?)', + [APP_NAME, USER_ID, SESSION_ID, '{}', now, now], + ); + await insertEventRow(orm, 'live-event', SESSION_ID); +} - afterEach(async () => { - if (orm) { - await orm.close(); - } - }); - - it('keeps events composite key columns within the MySQL index limit', async () => { - orm = await MikroORM.init({ - dbName: ':memory:', - driver: SqliteDriver, - entities: ENTITIES, - }); - - const eventProperties = orm.getMetadata().get(StorageEvent.name) - .properties as Record; - const keyProperties = ['id', 'appName', 'userId', 'sessionId']; - - for (const keyProperty of keyProperties) { - expect(eventProperties[keyProperty].length).toBe( - STORAGE_KEY_COLUMN_LENGTH, - ); - } - - const utf8mb4KeyBytes = keyProperties.reduce((total, keyProperty) => { - return total + eventProperties[keyProperty].length! * 4; - }, 0); - expect(utf8mb4KeyBytes).toBeLessThanOrEqual(3072); - }); - }); +/** + * Writes the row shape that a database created before the foreign key can + * hold: an event whose session was deleted. + */ +async function insertOrphanedEvent(orm: MikroORM): Promise { + const connection = orm.em.getConnection(); + await connection.execute('pragma foreign_keys = off'); + await insertEventRow(orm, 'orphan-event', 'missing-session'); + await connection.execute('pragma foreign_keys = on'); +} + +async function countEvents(orm: MikroORM): Promise { + return orm.em.fork().count(StorageEvent, {}); +} + +function makeWarnCapturingLogger() { + const warnCalls: string[] = []; + const mockLogger = { + setLogLevel: () => {}, + log: () => {}, + debug: () => {}, + info: () => {}, + warn: (...args: unknown[]) => { + warnCalls.push(args.map(String).join(' ')); + }, + error: () => {}, + }; + return {mockLogger, warnCalls}; +} +describe('operations', () => { describe('getConnectionOptionsFromUri', () => { it('should parse postgresql URI', async () => { const options = await getConnectionOptionsFromUri( @@ -162,6 +188,120 @@ describe('operations', () => { // Verify it runs without error await expect(ensureDatabaseCreated(orm)).resolves.not.toThrow(); }); + + it('deletes orphaned events and retries when the schema update fails', async () => { + orm = await MikroORM.init({ + dbName: ':memory:', + driver: SqliteDriver, + entities: ENTITIES, + }); + await orm.schema.createSchema(); + await insertOrphanedEvent(orm); + + const updateSchema = vi.spyOn(orm.schema, 'updateSchema'); + updateSchema.mockRejectedValueOnce( + new Error('constraint "events_app_name_user_id_session_id_foreign"'), + ); + + await expect(ensureDatabaseCreated(orm)).resolves.toBeUndefined(); + + expect(updateSchema).toHaveBeenCalledTimes(2); + expect(await countEvents(orm)).toBe(0); + }); + + it('propagates the error when the schema update fails again', async () => { + orm = await MikroORM.init({ + dbName: ':memory:', + driver: SqliteDriver, + entities: ENTITIES, + }); + await orm.schema.createSchema(); + + const updateSchema = vi.spyOn(orm.schema, 'updateSchema'); + updateSchema.mockRejectedValueOnce(new Error('first failure')); + updateSchema.mockRejectedValueOnce(new Error('second failure')); + + await expect(ensureDatabaseCreated(orm)).rejects.toThrow( + 'second failure', + ); + }); + + it('logs the schema update failure that triggered the retry', async () => { + orm = await MikroORM.init({ + dbName: ':memory:', + driver: SqliteDriver, + entities: ENTITIES, + }); + await orm.schema.createSchema(); + + const updateSchema = vi.spyOn(orm.schema, 'updateSchema'); + updateSchema.mockRejectedValueOnce(new Error('lock timeout')); + const {mockLogger, warnCalls} = makeWarnCapturingLogger(); + setLogger(mockLogger); + + try { + await ensureDatabaseCreated(orm); + } finally { + resetLogger(); + } + + expect(warnCalls[0]).toBe( + 'Schema update failed; retrying after deleting orphaned events. ' + + 'Error: lock timeout', + ); + }); + }); + + describe('deleteOrphanedEvents', () => { + let orm: MikroORM; + + afterEach(async () => { + if (orm) { + await orm.close(); + } + }); + + it('removes only the events whose session is gone', async () => { + orm = await MikroORM.init({ + dbName: ':memory:', + driver: SqliteDriver, + entities: ENTITIES, + }); + await orm.schema.createSchema(); + await insertLiveSessionWithEvent(orm); + await insertOrphanedEvent(orm); + expect(await countEvents(orm)).toBe(2); + + await deleteOrphanedEvents(orm); + + const remaining = await orm.em + .fork() + .find(StorageEvent, {}, {fields: ['id']}); + expect(remaining.map((event) => event.id)).toEqual(['live-event']); + }); + + it('logs how many rows it deleted', async () => { + orm = await MikroORM.init({ + dbName: ':memory:', + driver: SqliteDriver, + entities: ENTITIES, + }); + await orm.schema.createSchema(); + await insertLiveSessionWithEvent(orm); + await insertOrphanedEvent(orm); + const {mockLogger, warnCalls} = makeWarnCapturingLogger(); + setLogger(mockLogger); + + try { + await deleteOrphanedEvents(orm); + } finally { + resetLogger(); + } + + expect(warnCalls).toEqual([ + 'Deleted 1 event rows whose session no longer exists.', + ]); + }); }); describe('validateDatabaseSchemaVersion', () => { diff --git a/core/test/sessions/db/schema_test.ts b/core/test/sessions/db/schema_test.ts new file mode 100644 index 000000000..0a7cbd66d --- /dev/null +++ b/core/test/sessions/db/schema_test.ts @@ -0,0 +1,166 @@ +/** + * @license + * Copyright 2026 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import {createEvent} from '@google/adk'; +import {MikroORM} from '@mikro-orm/core'; +import {SqliteDriver} from '@mikro-orm/sqlite'; +import {afterEach, beforeEach, describe, expect, it} from 'vitest'; +import { + ENTITIES, + STORAGE_KEY_COLUMN_LENGTH, + StorageEvent, + StorageSession, +} from '../../../src/sessions/db/schema.js'; + +const APP_NAME = 'test-app'; +const USER_ID = 'test-user'; +const SESSION_ID = 'test-session'; + +describe('storage schema', () => { + let orm: MikroORM; + + beforeEach(async () => { + orm = await MikroORM.init({ + dbName: ':memory:', + driver: SqliteDriver, + entities: ENTITIES, + }); + await orm.schema.createSchema(); + }); + + afterEach(async () => { + await orm.close(); + }); + + async function createStorageSession(id = SESSION_ID): Promise { + const em = orm.em.fork(); + em.create(StorageSession, { + id, + appName: APP_NAME, + userId: USER_ID, + state: {}, + }); + await em.flush(); + } + + async function createStorageEvent( + id: string, + sessionId = SESSION_ID, + ): Promise { + const em = orm.em.fork(); + const session = em.getReference(StorageSession, [ + APP_NAME, + USER_ID, + sessionId, + ]); + em.create(StorageEvent, { + id, + session, + invocationId: `invocation-${id}`, + timestamp: new Date(), + eventData: createEvent({id}), + }); + await em.flush(); + } + + async function countEvents(): Promise { + return orm.em.fork().count(StorageEvent, {}); + } + + it('keeps events composite key columns within the MySQL index limit', async () => { + // `session` owns app_name, user_id and session_id, so these two + // properties emit every key column of the events table. + const eventProperties = orm.getMetadata().get(StorageEvent.name) + .properties as Record; + const keyColumnLengths = ['id', 'session'].flatMap((keyProperty) => { + const property = eventProperties[keyProperty]; + return property.fieldNames.map(() => property.length); + }); + + expect(keyColumnLengths).toEqual(Array(4).fill(STORAGE_KEY_COLUMN_LENGTH)); + + const utf8mb4KeyBytes = keyColumnLengths.reduce( + (total, length) => total + (length ?? 0) * 4, + 0, + ); + expect(utf8mb4KeyBytes).toBeLessThanOrEqual(3072); + }); + + it('declares an events -> sessions foreign key that cascades on delete', async () => { + const sql = await orm.schema.getCreateSchemaSQL(); + + expect(sql).toContain( + 'foreign key(`app_name`, `user_id`, `session_id`) ' + + 'references `sessions`(`app_name`, `user_id`, `id`) ' + + 'on delete cascade on update cascade', + ); + }); + + it('keeps the events primary key columns in their original order', async () => { + const sql = await orm.schema.getCreateSchemaSQL(); + + expect(sql).toContain( + 'primary key (`id`, `app_name`, `user_id`, `session_id`)', + ); + }); + + it('deletes the events when the session row alone is deleted', async () => { + await createStorageSession(); + await createStorageEvent('event-1'); + await createStorageEvent('event-2'); + expect(await countEvents()).toBe(2); + + await orm.em + .getConnection() + .execute( + 'DELETE FROM sessions WHERE app_name = ? AND user_id = ? AND id = ?', + [APP_NAME, USER_ID, SESSION_ID], + ); + + expect(await countEvents()).toBe(0); + }); + + it('rejects an event that names a session which does not exist', async () => { + await expect( + createStorageEvent('orphan-event', 'missing-session'), + ).rejects.toThrow(/FOREIGN KEY/i); + }); + + it('enforces foreign keys on a SQLite connection without an app-level pragma', async () => { + const rows = await orm.em + .getConnection() + .execute>('pragma foreign_keys'); + + expect(rows[0].foreign_keys).toBe(1); + }); + + it('reads and deletes events through the non-persistent key mirrors', async () => { + await createStorageSession(); + await createStorageEvent('event-1'); + + const em = orm.em.fork(); + const found = await em.find(StorageEvent, { + appName: APP_NAME, + userId: USER_ID, + sessionId: SESSION_ID, + }); + + expect(found).toHaveLength(1); + expect(found[0].appName).toBe(APP_NAME); + expect(found[0].userId).toBe(USER_ID); + expect(found[0].sessionId).toBe(SESSION_ID); + expect(found[0].eventData.id).toBe('event-1'); + + const deleted = await em.nativeDelete(StorageEvent, { + appName: APP_NAME, + userId: USER_ID, + sessionId: SESSION_ID, + }); + + expect(deleted).toBe(1); + expect(await countEvents()).toBe(0); + }); +});