diff --git a/.changeset/quiet-migrations-share.md b/.changeset/quiet-migrations-share.md new file mode 100644 index 0000000000..67b969181b --- /dev/null +++ b/.changeset/quiet-migrations-share.md @@ -0,0 +1,6 @@ +--- +"@electric-ax/agents-server": patch +--- + +Allow `runMigrations` to use a caller-owned PostgreSQL client while always +closing clients created from PostgreSQL URLs. diff --git a/packages/agents-server/src/db/index.ts b/packages/agents-server/src/db/index.ts index 9f970702ce..9d50f397de 100644 --- a/packages/agents-server/src/db/index.ts +++ b/packages/agents-server/src/db/index.ts @@ -40,14 +40,39 @@ export function resolveMigrationsFolder(fromUrl = import.meta.url): string { return folder } -export async function runMigrations(postgresUrl: string): Promise { - const migrationClient = postgres(postgresUrl, { - max: 1, - onnotice: () => {}, - }) - const db = drizzle(migrationClient) +async function migrateClient(client: PgClient): Promise { + const db = drizzle(client) await migrate(db, { migrationsFolder: resolveMigrationsFolder(), }) - await migrationClient.end() +} + +/** + * Runs migrations with a library-owned client that is always closed. + */ +export function runMigrations(postgresUrl: string): Promise +/** + * Runs migrations with a caller-owned client that is never closed. + * + * The caller remains responsible for closing the client after success or + * failure. + */ +export function runMigrations(client: PgClient): Promise +export async function runMigrations( + postgresUrlOrClient: string | PgClient +): Promise { + if (typeof postgresUrlOrClient !== `string`) { + await migrateClient(postgresUrlOrClient) + return + } + + const migrationClient = postgres(postgresUrlOrClient, { + max: 1, + onnotice: () => {}, + }) + try { + await migrateClient(migrationClient) + } finally { + await migrationClient.end() + } } diff --git a/packages/agents-server/test/db-migrations.test.ts b/packages/agents-server/test/db-migrations.test.ts new file mode 100644 index 0000000000..8479303712 --- /dev/null +++ b/packages/agents-server/test/db-migrations.test.ts @@ -0,0 +1,110 @@ +import postgres from 'postgres' +import { beforeEach, describe, expect, expectTypeOf, it, vi } from 'vitest' + +import { runMigrations } from '../src/index' +import type { PgClient } from '../src/index' + +const { drizzleMock, migrateMock, postgresMock } = vi.hoisted(() => ({ + drizzleMock: vi.fn(), + migrateMock: vi.fn(async () => {}), + postgresMock: vi.fn((_postgresUrl: string) => ({ + end: vi.fn(async () => {}), + })), +})) + +vi.mock(`postgres`, () => ({ + default: postgresMock, +})) + +vi.mock(`drizzle-orm/postgres-js`, () => ({ + drizzle: drizzleMock, +})) + +vi.mock(`drizzle-orm/postgres-js/migrator`, () => ({ + migrate: migrateMock, +})) + +const syntheticPostgresUrl = `postgresql://db.invalid/agents` +const migrationDb = { kind: `migration-db` } + +function compilePublicMigrationCalls(client: PgClient): void { + void runMigrations(syntheticPostgresUrl) + void runMigrations(client) +} + +function createCallerOwnedClient() { + const endMock = vi.fn(async () => {}) + postgresMock.mockReturnValueOnce({ end: endMock }) + const client = postgres(syntheticPostgresUrl) + postgresMock.mockClear() + return { client, endMock } +} + +describe(`runMigrations`, () => { + beforeEach(() => { + vi.clearAllMocks() + drizzleMock.mockReturnValue(migrationDb) + migrateMock.mockResolvedValue(undefined) + }) + + it(`accepts URL and PgClient inputs through the package root`, () => { + expectTypeOf(compilePublicMigrationCalls).returns.toEqualTypeOf() + }) + + it(`uses a caller-owned client without closing it after success`, async () => { + const { client, endMock } = createCallerOwnedClient() + + await runMigrations(client) + + expect(postgresMock).not.toHaveBeenCalled() + expect(drizzleMock).toHaveBeenCalledOnce() + expect(drizzleMock.mock.calls.at(0)?.at(0)).toBe(client) + expect(migrateMock).toHaveBeenCalledWith(migrationDb, { + migrationsFolder: expect.any(String), + }) + expect(endMock).not.toHaveBeenCalled() + }) + + it(`does not close a caller-owned client after migration failure`, async () => { + const migrationError = new Error(`migration failed`) + const { client, endMock } = createCallerOwnedClient() + migrateMock.mockRejectedValueOnce(migrationError) + + await expect(runMigrations(client)).rejects.toBe(migrationError) + + expect(postgresMock).not.toHaveBeenCalled() + expect(drizzleMock).toHaveBeenCalledOnce() + expect(drizzleMock.mock.calls.at(0)?.at(0)).toBe(client) + expect(endMock).not.toHaveBeenCalled() + }) + + it(`closes a URL-created client after success`, async () => { + const endMock = vi.fn(async () => {}) + const libraryOwnedClient = { end: endMock } + postgresMock.mockReturnValueOnce(libraryOwnedClient) + + await runMigrations(syntheticPostgresUrl) + + expect(postgresMock).toHaveBeenCalledWith(syntheticPostgresUrl, { + max: 1, + onnotice: expect.any(Function), + }) + expect(drizzleMock.mock.calls.at(0)?.at(0)).toBe(libraryOwnedClient) + expect(endMock).toHaveBeenCalledOnce() + }) + + it(`closes a URL-created client after migration failure`, async () => { + const migrationError = new Error(`migration failed`) + const endMock = vi.fn(async () => {}) + const libraryOwnedClient = { end: endMock } + postgresMock.mockReturnValueOnce(libraryOwnedClient) + migrateMock.mockRejectedValueOnce(migrationError) + + await expect(runMigrations(syntheticPostgresUrl)).rejects.toBe( + migrationError + ) + + expect(drizzleMock.mock.calls.at(0)?.at(0)).toBe(libraryOwnedClient) + expect(endMock).toHaveBeenCalledOnce() + }) +})