feat(agents-server): accept caller-owned migration clients

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

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<void> {
-  const migrationClient = postgres(postgresUrl, {
-    max: 1,
-    onnotice: () => {},
-  })
-  const db = drizzle(migrationClient)
+async function migrateClient(client: PgClient): Promise<void> {
+  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<void>
+/**
+ * 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<void>
+export async function runMigrations(
+  postgresUrlOrClient: string | PgClient
+): Promise<void> {
+  if (typeof postgresUrlOrClient !== `string`) {
+    await migrateClient(postgresUrlOrClient)
+    return
+  }
+
+  const migrationClient = postgres(postgresUrlOrClient, {
+    max: 1,
+    onnotice: () => {},
+  })
+  try {
+    await migrateClient(migrationClient)
+  } finally {
+    await migrationClient.end()
+  }
 }

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<void>()
+  })
+
+  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()
+  })
+})