From 219d162e1a9b43b8e74af5c87654174bd1c3bbdf Mon Sep 17 00:00:00 2001 From: LiuVaayne <10231735+vaayne@users.noreply.github.com> Date: Wed, 17 Sep 2025 12:08:59 +0800 Subject: [PATCH] feat: Add automatic database migration system for agents service (#10215) * feat: Add automatic database migration system for agents service - Add migrations tracking schema with version, tag, and timestamp - Implement MigrationService to automatically run pending migrations - Integrate migration check into BaseService initialization - Read migration files from drizzle/ directory and journal.json - Track applied migrations to prevent re-execution - Ensure database is always at latest version on service startup Generated with [Claude Code](https://claude.ai/code) via [Happy](https://happy.engineering) Co-Authored-By: Claude Co-Authored-By: Happy * refactor: Improve migration logging and enhance database path configuration * chore: harden migration bootstrap flow --------- Co-authored-by: Claude Co-authored-by: Happy --- src/main/services/agents/BaseService.ts | 5 + .../agents/database/MigrationService.ts | 169 ++++++++++++++++++ .../services/agents/database/schema/index.ts | 1 + .../database/schema/migrations.schema.ts | 14 ++ src/main/services/agents/drizzle.config.ts | 11 +- 5 files changed, 198 insertions(+), 2 deletions(-) create mode 100644 src/main/services/agents/database/MigrationService.ts create mode 100644 src/main/services/agents/database/schema/migrations.schema.ts diff --git a/src/main/services/agents/BaseService.ts b/src/main/services/agents/BaseService.ts index 53830df243..de33e0899d 100644 --- a/src/main/services/agents/BaseService.ts +++ b/src/main/services/agents/BaseService.ts @@ -5,6 +5,7 @@ import fs from 'fs' import path from 'path' import * as schema from './database/schema' +import { MigrationService } from './database/MigrationService' import { dbPath } from './drizzle.config' const logger = loggerService.withContext('BaseService') @@ -66,6 +67,10 @@ export abstract class BaseService { BaseService.db = drizzle(BaseService.client, { schema }) + // Run database migrations + const migrationService = new MigrationService(BaseService.db, BaseService.client) + await migrationService.runMigrations() + BaseService.isInitialized = true logger.info('Agent database initialized successfully') return diff --git a/src/main/services/agents/database/MigrationService.ts b/src/main/services/agents/database/MigrationService.ts new file mode 100644 index 0000000000..16e13f56ca --- /dev/null +++ b/src/main/services/agents/database/MigrationService.ts @@ -0,0 +1,169 @@ +import { type Client } from '@libsql/client' +import { loggerService } from '@logger' +import { type LibSQLDatabase } from 'drizzle-orm/libsql' +import fs from 'fs' +import path from 'path' + +import * as schema from './schema' +import { migrations, type NewMigration } from './schema/migrations.schema' + +const logger = loggerService.withContext('MigrationService') + +interface MigrationJournal { + version: string + dialect: string + entries: Array<{ + idx: number + version: string + when: number + tag: string + breakpoints: boolean + }> +} + +export class MigrationService { + private db: LibSQLDatabase + private client: Client + private migrationDir: string + + constructor(db: LibSQLDatabase, client: Client) { + this.db = db + this.client = client + this.migrationDir = path.join(__dirname, 'drizzle') + } + + async runMigrations(): Promise { + try { + logger.info('Starting migration check...') + + // Ensure migrations table exists + await this.ensureMigrationsTable() + + // Read migration journal + const journal = await this.readMigrationJournal() + if (!journal.entries.length) { + logger.info('No migrations found in journal') + return + } + + // Get applied migrations + const appliedMigrations = await this.getAppliedMigrations() + const appliedVersions = new Set(appliedMigrations.map((m) => Number(m.version))) + + const latestAppliedVersion = appliedMigrations.reduce( + (max, migration) => Math.max(max, Number(migration.version)), + 0 + ) + const latestJournalVersion = journal.entries.reduce((max, entry) => Math.max(max, entry.idx), 0) + + logger.info(`Latest applied migration: v${latestAppliedVersion}, latest available: v${latestJournalVersion}`) + + // Find pending migrations (compare journal idx with stored version, which is the same value) + const pendingMigrations = journal.entries + .filter((entry) => !appliedVersions.has(entry.idx)) + .sort((a, b) => a.idx - b.idx) + + if (pendingMigrations.length === 0) { + logger.info('Database is up to date') + return + } + + logger.info(`Found ${pendingMigrations.length} pending migrations`) + + // Execute pending migrations + for (const migration of pendingMigrations) { + await this.executeMigration(migration) + } + + logger.info('All migrations completed successfully') + } catch (error) { + logger.error('Migration failed:', { error }) + throw error + } + } + + private async ensureMigrationsTable(): Promise { + try { + const tableExists = await this.client.execute(`SELECT name FROM sqlite_master WHERE type='table' AND name='migrations'`) + + if (tableExists.rows.length === 0) { + logger.info('Migrations table missing, creating...') + + await this.client.execute(` + CREATE TABLE IF NOT EXISTS migrations ( + version INTEGER PRIMARY KEY, + tag TEXT NOT NULL, + executed_at INTEGER NOT NULL + ) + `) + + logger.info('Migrations table created successfully') + } else { + logger.debug('Migrations table already exists') + } + } catch (error) { + logger.error('Failed to ensure migrations table exists:', { error }) + throw error + } + } + + private async readMigrationJournal(): Promise { + const journalPath = path.join(this.migrationDir, 'meta', '_journal.json') + + if (!fs.existsSync(journalPath)) { + logger.warn('Migration journal not found:', { journalPath }) + return { version: '7', dialect: 'sqlite', entries: [] } + } + + try { + const journalContent = fs.readFileSync(journalPath, 'utf-8') + return JSON.parse(journalContent) + } catch (error) { + logger.error('Failed to read migration journal:', { error }) + throw error + } + } + + private async getAppliedMigrations(): Promise { + try { + return await this.db.select().from(migrations) + } catch (error) { + // This should not happen since we ensure the table exists in runMigrations() + logger.error('Failed to query applied migrations:', { error }) + throw error + } + } + + private async executeMigration(migration: MigrationJournal['entries'][0]): Promise { + const sqlFilePath = path.join(this.migrationDir, `${migration.tag}.sql`) + + if (!fs.existsSync(sqlFilePath)) { + throw new Error(`Migration SQL file not found: ${sqlFilePath}`) + } + + try { + logger.info(`Executing migration ${migration.tag}...`) + const startTime = Date.now() + + // Read and execute SQL + const sqlContent = fs.readFileSync(sqlFilePath, 'utf-8') + await this.client.execute(sqlContent) + + // Record migration as applied (store journal idx as version for tracking) + const newMigration: NewMigration = { + version: migration.idx, + tag: migration.tag, + executedAt: Date.now() + } + + await this.db.insert(migrations).values(newMigration) + + const executionTime = Date.now() - startTime + logger.info(`Migration ${migration.tag} completed in ${executionTime}ms`) + } catch (error) { + logger.error(`Migration ${migration.tag} failed:`, { error }) + throw error + } + } + +} diff --git a/src/main/services/agents/database/schema/index.ts b/src/main/services/agents/database/schema/index.ts index f15f0c7014..c8d3a38012 100644 --- a/src/main/services/agents/database/schema/index.ts +++ b/src/main/services/agents/database/schema/index.ts @@ -5,3 +5,4 @@ export * from './agents.schema' export * from './sessions.schema' export * from './messages.schema' +export * from './migrations.schema' diff --git a/src/main/services/agents/database/schema/migrations.schema.ts b/src/main/services/agents/database/schema/migrations.schema.ts new file mode 100644 index 0000000000..b3485ca911 --- /dev/null +++ b/src/main/services/agents/database/schema/migrations.schema.ts @@ -0,0 +1,14 @@ +/** + * Migration tracking schema + */ + +import { integer, sqliteTable, text } from 'drizzle-orm/sqlite-core' + +export const migrations = sqliteTable('migrations', { + version: integer('version').primaryKey(), + tag: text('tag').notNull(), + executedAt: integer('executed_at').notNull() +}) + +export type Migration = typeof migrations.$inferSelect +export type NewMigration = typeof migrations.$inferInsert \ No newline at end of file diff --git a/src/main/services/agents/drizzle.config.ts b/src/main/services/agents/drizzle.config.ts index adeb0f0191..af31934f56 100644 --- a/src/main/services/agents/drizzle.config.ts +++ b/src/main/services/agents/drizzle.config.ts @@ -5,16 +5,23 @@ import os from 'node:os' import path from 'node:path' +import { isDev } from '@main/constant' import { defineConfig } from 'drizzle-kit' +import { app } from 'electron' -export const dbPath = path.join(os.homedir(), '.cherrystudio', 'data', 'agents.db') +function getDbPath() { + if (isDev) { + return path.join(os.homedir(), '.cherrystudio', 'data', 'agents.db') + } + return path.join(app.getPath('userData'), 'agents.db') +} export default defineConfig({ dialect: 'sqlite', schema: './src/main/services/agents/database/schema/index.ts', out: './src/main/services/agents/database/drizzle', dbCredentials: { - url: `file:${dbPath}` + url: `file:${getDbPath()}` }, verbose: true, strict: true