feat(storage): pgvector adapter support gated on tier=federated (FED-M1-03) (#472)

packages/storage/src/adapters/postgres.spec.ts

@@ -0,0 +1,107 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+import type { DbHandle } from '@mosaicstack/db';
+
+// Mock @mosaicstack/db before importing the adapter
+vi.mock('@mosaicstack/db', async (importOriginal) => {
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  const actual = await importOriginal<Record<string, any>>();
+  return {
+    ...actual,
+    createDb: vi.fn(),
+    runMigrations: vi.fn().mockResolvedValue(undefined),
+  };
+});
+
+import { createDb, runMigrations } from '@mosaicstack/db';
+import { PostgresAdapter } from './postgres.js';
+
+describe('PostgresAdapter — vector extension gating', () => {
+  let mockExecute: ReturnType<typeof vi.fn>;
+  let mockDb: { execute: ReturnType<typeof vi.fn> };
+  let mockHandle: Pick<DbHandle, 'close'> & { db: typeof mockDb };
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    mockExecute = vi.fn().mockResolvedValue(undefined);
+    mockDb = { execute: mockExecute };
+    mockHandle = { db: mockDb, close: vi.fn().mockResolvedValue(undefined) };
+    vi.mocked(createDb).mockReturnValue(mockHandle as unknown as DbHandle);
+  });
+
+  it('calls db.execute with CREATE EXTENSION IF NOT EXISTS vector when enableVector=true', async () => {
+    const adapter = new PostgresAdapter({
+      type: 'postgres',
+      url: 'postgresql://test:test@localhost:5432/test',
+      enableVector: true,
+    });
+
+    await adapter.migrate();
+
+    // Should have called execute
+    expect(mockExecute).toHaveBeenCalledTimes(1);
+
+    // Verify the SQL contains the extension creation statement.
+    // Prefer Drizzle's public toSQL() API; fall back to queryChunks if unavailable.
+    // NOTE: queryChunks is an undocumented Drizzle internal (drizzle-orm ^0.45.x).
+    // toSQL() was not present on the raw sql`` result in this version — if a future
+    // Drizzle upgrade adds it, remove the fallback path and delete this comment.
+    const sqlObj = mockExecute.mock.calls[0]![0] as {
+      toSQL?: () => { sql: string; params: unknown[] };
+      queryChunks?: Array<{ value: string[] }>;
+    };
+    const sqlText = sqlObj.toSQL
+      ? sqlObj.toSQL().sql.toLowerCase()
+      : (sqlObj.queryChunks ?? [])
+          .flatMap((chunk) => chunk.value)
+          .join('')
+          .toLowerCase();
+    expect(sqlText).toContain('create extension if not exists vector');
+  });
+
+  it('does NOT call db.execute for extension when enableVector is false', async () => {
+    const adapter = new PostgresAdapter({
+      type: 'postgres',
+      url: 'postgresql://test:test@localhost:5432/test',
+      enableVector: false,
+    });
+
+    await adapter.migrate();
+
+    expect(mockExecute).not.toHaveBeenCalled();
+    expect(vi.mocked(runMigrations)).toHaveBeenCalledOnce();
+  });
+
+  it('does NOT call db.execute for extension when enableVector is unset', async () => {
+    const adapter = new PostgresAdapter({
+      type: 'postgres',
+      url: 'postgresql://test:test@localhost:5432/test',
+    });
+
+    await adapter.migrate();
+
+    expect(mockExecute).not.toHaveBeenCalled();
+    expect(vi.mocked(runMigrations)).toHaveBeenCalledOnce();
+  });
+
+  it('calls runMigrations after the extension is created', async () => {
+    const callOrder: string[] = [];
+    mockExecute.mockImplementation(() => {
+      callOrder.push('execute');
+      return Promise.resolve(undefined);
+    });
+    vi.mocked(runMigrations).mockImplementation(() => {
+      callOrder.push('runMigrations');
+      return Promise.resolve();
+    });
+
+    const adapter = new PostgresAdapter({
+      type: 'postgres',
+      url: 'postgresql://test:test@localhost:5432/test',
+      enableVector: true,
+    });
+
+    await adapter.migrate();
+
+    expect(callOrder).toEqual(['execute', 'runMigrations']);
+  });
+});

packages/storage/src/adapters/postgres.ts

@@ -66,13 +66,19 @@ export class PostgresAdapter implements StorageAdapter {
   private handle: DbHandle;
   private db: Db;
   private url: string;
+  private enableVector: boolean;
 
   constructor(config: Extract<StorageConfig, { type: 'postgres' }>) {
     this.url = config.url;
+    this.enableVector = config.enableVector ?? false;
     this.handle = createDb(config.url);
     this.db = this.handle.db;
   }
 
+  private async ensureVectorExtension(): Promise<void> {
+    await this.db.execute(sql`CREATE EXTENSION IF NOT EXISTS vector`);
+  }
+
   async create<T extends Record<string, unknown>>(
     collection: string,
     data: T,
@@ -149,6 +155,9 @@ export class PostgresAdapter implements StorageAdapter {
   }
 
   async migrate(): Promise<void> {
+    if (this.enableVector) {
+      await this.ensureVectorExtension();
+    }
     await runMigrations(this.url);
   }
 

packages/storage/src/types.ts

@@ -38,6 +38,6 @@ export interface StorageAdapter {
 }
 
 export type StorageConfig =
-  | { type: 'postgres'; url: string }
+  | { type: 'postgres'; url: string; enableVector?: boolean }
   | { type: 'pglite'; dataDir?: string }
   | { type: 'files'; dataDir: string; format?: 'json' | 'md' };