@axiom-lattice/pg-stores 5.0.0 → 5.0.1

package/src/__tests__/PostgreSQLModelProviderStore.integration.test.ts

@@ -0,0 +1,313 @@
+import { randomUUID } from "node:crypto";
+import { Pool } from "pg";
+
+// The @axiom-lattice/core barrel cannot load under jest in this package (ESM-only
+// transitive deps), so mock it with the real encryption module — node:crypto only.
+jest.mock("@axiom-lattice/core", () => ({
+  ...jest.requireActual("../../../core/src/util/encryption"),
+}));
+
+import { decrypt } from "@axiom-lattice/core";
+import { PostgreSQLModelProviderStore } from "../stores/PostgreSQLModelProviderStore";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const describeIfDatabase = TEST_DATABASE_URL ? describe : describe.skip;
+
+describeIfDatabase("PostgreSQLModelProviderStore integration", () => {
+  let store: PostgreSQLModelProviderStore;
+  let pool: Pool;
+  const tenants = new Set<string>();
+
+  beforeAll(async () => {
+    store = new PostgreSQLModelProviderStore({
+      poolConfig: TEST_DATABASE_URL!,
+      autoMigrate: true,
+    });
+    await store.initialize();
+    pool = new Pool({ connectionString: TEST_DATABASE_URL });
+  });
+
+  afterEach(async () => {
+    if (tenants.size === 0) {
+      return;
+    }
+    const ids = [...tenants];
+    await pool.query(
+      "DELETE FROM lattice_model_provider_models WHERE tenant_id = ANY($1::text[])",
+      [ids],
+    );
+    await pool.query(
+      "DELETE FROM lattice_model_providers WHERE tenant_id = ANY($1::text[])",
+      [ids],
+    );
+    tenants.clear();
+  });
+
+  afterAll(async () => {
+    await pool.end();
+    await store.dispose();
+  });
+
+  function tenantId(): string {
+    const id = `model-provider-${randomUUID()}`;
+    tenants.add(id);
+    return id;
+  }
+
+  function providerData(index = 0) {
+    return {
+      name: `openai-${index}-${randomUUID()}`,
+      displayName: `OpenAI ${index}`,
+      protocol: "openai-compatible" as const,
+      apiStyle: "chat-completions" as const,
+      baseURL: "https://api.example.com/v1",
+      apiKey: "sk-secret-abcd",
+      enabled: true,
+    };
+  }
+
+  async function createProvider(tenant: string) {
+    const id = randomUUID();
+    const data = providerData();
+    const provider = await store.createProvider(tenant, id, data);
+    return { id, data, provider };
+  }
+
+  it("creates and reads a provider by id and by name", async () => {
+    const tenant = tenantId();
+    const { id, data, provider } = await createProvider(tenant);
+
+    expect(provider).toMatchObject({
+      id,
+      tenantId: tenant,
+      name: data.name,
+      displayName: data.displayName,
+      protocol: "openai-compatible",
+      apiStyle: "chat-completions",
+      llmProvider: "openai",
+      baseURL: data.baseURL,
+      enabled: true,
+      status: "unknown",
+      apiKey: data.apiKey,
+      apiKeyHint: "abcd",
+    });
+
+    const byId = await store.getProviderById(tenant, id);
+    expect(byId?.id).toBe(id);
+    expect(byId?.apiKey).toBe(data.apiKey);
+
+    const byName = await store.getProviderByName(tenant, data.name);
+    expect(byName?.id).toBe(id);
+
+    const listed = await store.listProviders(tenant);
+    expect(listed.map((entry) => entry.id)).toEqual([id]);
+  });
+
+  it("encrypts the API key at rest and decrypts it on read", async () => {
+    const tenant = tenantId();
+    const { id, data } = await createProvider(tenant);
+
+    const raw = await pool.query<{ api_key_enc: string; api_key_hint: string }>(
+      "SELECT api_key_enc, api_key_hint FROM lattice_model_providers WHERE tenant_id = $1 AND id = $2",
+      [tenant, id],
+    );
+
+    expect(raw.rows[0].api_key_enc).not.toBe(data.apiKey);
+    expect(decrypt(raw.rows[0].api_key_enc)).toBe(data.apiKey);
+    expect(raw.rows[0].api_key_hint).toBe("abcd");
+  });
+
+  it("updates provider fields, keeps an omitted key, and can clear it", async () => {
+    const tenant = tenantId();
+    const { id, data } = await createProvider(tenant);
+
+    const renamed = await store.updateProvider(tenant, id, {
+      name: `${data.name}-renamed`,
+      displayName: "Renamed",
+      baseURL: "https://new.example.com/v1",
+      enabled: false,
+    });
+
+    expect(renamed).toMatchObject({
+      name: `${data.name}-renamed`,
+      displayName: "Renamed",
+      baseURL: "https://new.example.com/v1",
+      enabled: false,
+      apiKey: data.apiKey,
+    });
+
+    const rekeyed = await store.updateProvider(tenant, id, { apiKey: "sk-new-wxyz" });
+    expect(rekeyed?.apiKey).toBe("sk-new-wxyz");
+    expect(rekeyed?.apiKeyHint).toBe("wxyz");
+
+    const cleared = await store.updateProvider(tenant, id, { clearApiKey: true });
+    expect(cleared?.apiKey).toBeUndefined();
+    expect(cleared?.apiKeyHint).toBeUndefined();
+  });
+
+  it("records discovery status and clears the error on success", async () => {
+    const tenant = tenantId();
+    const { id } = await createProvider(tenant);
+
+    await store.markProviderStatus(tenant, id, "error", "UPSTREAM_UNAVAILABLE");
+    const failed = await store.getProviderById(tenant, id);
+    expect(failed?.status).toBe("error");
+    expect(failed?.lastError).toBe("UPSTREAM_UNAVAILABLE");
+
+    const discoveredAt = new Date("2024-01-02T03:04:05.000Z");
+    await store.markProviderStatus(tenant, id, "ok", undefined, discoveredAt);
+    const ok = await store.getProviderById(tenant, id);
+    expect(ok?.status).toBe("ok");
+    expect(ok?.lastError).toBeUndefined();
+    // `last_discovered_at` is a zone-naive TIMESTAMP, so the driver interprets
+    // the stored wall-clock value in the local zone. Assert the instant was
+    // persisted rather than an exact ISO string, which is timezone-dependent.
+    expect(ok?.lastDiscoveredAt).toBeInstanceOf(Date);
+    expect(
+      Math.abs(ok!.lastDiscoveredAt!.getTime() - discoveredAt.getTime()),
+    ).toBeLessThan(24 * 60 * 60 * 1000);
+  });
+
+  it("replaces, reads, toggles, and deletes a provider's models", async () => {
+    const tenant = tenantId();
+    const { id: providerId } = await createProvider(tenant);
+
+    const returned = await store.replaceProviderModels(tenant, providerId, [
+      { id: "gpt-4o", displayName: "GPT-4o", ownedBy: "openai" },
+      { id: "gpt-4o-mini", extra: { temperature: 0.2 }, kind: "chat" },
+    ]);
+    expect(returned).toEqual(["gpt-4o", "gpt-4o-mini"]);
+
+    const all = await store.listModels(tenant);
+    expect(all.map((model) => model.modelId).sort()).toEqual(["gpt-4o", "gpt-4o-mini"]);
+
+    const scoped = await store.listModels(tenant, providerId);
+    expect(scoped).toHaveLength(2);
+
+    const model = await store.getModel(tenant, providerId, "gpt-4o-mini");
+    expect(model).toMatchObject({
+      tenantId: tenant,
+      providerId,
+      modelId: "gpt-4o-mini",
+      enabled: true,
+      kind: "chat",
+      extra: { temperature: 0.2 },
+    });
+
+    const disabled = await store.setModelEnabled(tenant, providerId, "gpt-4o-mini", false);
+    expect(disabled?.enabled).toBe(false);
+
+    // A later discovery pass without an explicit `enabled` keeps the tenant choice.
+    await store.replaceProviderModels(tenant, providerId, [
+      { id: "gpt-4o" },
+      { id: "gpt-4o-mini" },
+    ]);
+    expect(
+      (await store.getModel(tenant, providerId, "gpt-4o-mini"))?.enabled,
+    ).toBe(false);
+
+    expect(await store.deleteModelsByProvider(tenant, providerId)).toBe(2);
+    expect(await store.listModels(tenant)).toEqual([]);
+  });
+
+  it("reconciles a snapshot: drops removed non-defaults, keeps the default stale, and clears stale on rediscovery", async () => {
+    const tenant = tenantId();
+    const { id: providerId } = await createProvider(tenant);
+
+    await store.replaceProviderModels(tenant, providerId, [
+      { id: "default-model", isDefault: true },
+      { id: "removed-model" },
+    ]);
+    expect(
+      (await store.getModel(tenant, providerId, "default-model"))?.isDefault,
+    ).toBe(true);
+
+    // The upstream drops both existing models and offers a new one.
+    await store.replaceProviderModels(tenant, providerId, [
+      { id: "replacement-model" },
+    ]);
+
+    expect(await store.getModel(tenant, providerId, "removed-model")).toBeNull();
+
+    const retainedDefault = await store.getModel(
+      tenant,
+      providerId,
+      "default-model",
+    );
+    expect(retainedDefault).toMatchObject({
+      modelId: "default-model",
+      isDefault: true,
+      stale: true,
+    });
+
+    const replacement = await store.getModel(
+      tenant,
+      providerId,
+      "replacement-model",
+    );
+    expect(replacement).toMatchObject({ isDefault: false, stale: false });
+    expect(
+      (await store.listModels(tenant, providerId)).map((model) => model.modelId).sort(),
+    ).toEqual(["default-model", "replacement-model"]);
+
+    // Rediscovery offers the retained default again: stale clears while the
+    // default flag is preserved, because discovery carries no isDefault opinion.
+    await store.replaceProviderModels(tenant, providerId, [
+      { id: "default-model" },
+      { id: "replacement-model" },
+    ]);
+
+    expect(
+      await store.getModel(tenant, providerId, "default-model"),
+    ).toMatchObject({ isDefault: true, stale: false });
+    expect(
+      await store.getModel(tenant, providerId, "replacement-model"),
+    ).toMatchObject({ isDefault: false, stale: false });
+  });
+
+  it("deletes a provider together with its model snapshot", async () => {
+    const tenant = tenantId();
+    const { id: providerId } = await createProvider(tenant);
+    await store.replaceProviderModels(tenant, providerId, [{ id: "m1" }]);
+
+    expect(await store.deleteProvider(tenant, providerId)).toBe(true);
+    expect(await store.getProviderById(tenant, providerId)).toBeNull();
+    expect(await store.listModels(tenant)).toEqual([]);
+    expect(await store.deleteProvider(tenant, providerId)).toBe(false);
+  });
+
+  it("isolates providers and models by tenant", async () => {
+    const tenantA = tenantId();
+    const tenantB = tenantId();
+    const idA = randomUUID();
+    const idB = randomUUID();
+    const name = `shared-name-${randomUUID()}`;
+
+    await store.createProvider(tenantA, idA, { ...providerData(), name });
+    await store.createProvider(tenantB, idB, { ...providerData(), name });
+
+    expect((await store.getProviderById(tenantA, idB))).toBeNull();
+    expect((await store.getProviderByName(tenantA, name))?.id).toBe(idA);
+    expect((await store.listProviders(tenantA)).map((entry) => entry.id)).toEqual([idA]);
+
+    await store.replaceProviderModels(tenantA, idA, [{ id: "same-model" }]);
+    await store.replaceProviderModels(tenantB, idB, [{ id: "same-model" }]);
+
+    expect(await store.listModels(tenantA)).toHaveLength(1);
+    expect(await store.listModels(tenantB)).toHaveLength(1);
+    expect(await store.getModel(tenantA, idB, "same-model")).toBeNull();
+  });
+
+  it("returns absent results for unknown providers and models", async () => {
+    const tenant = tenantId();
+    const missing = randomUUID();
+
+    expect(await store.getProviderById(tenant, missing)).toBeNull();
+    expect(await store.getProviderByName(tenant, "does-not-exist")).toBeNull();
+    expect(await store.updateProvider(tenant, missing, { name: "x" })).toBeNull();
+    expect(await store.deleteProvider(tenant, missing)).toBe(false);
+    expect(await store.getModel(tenant, missing, "m1")).toBeNull();
+    expect(await store.setModelEnabled(tenant, missing, "m1", false)).toBeNull();
+    expect(await store.deleteModelsByProvider(tenant, missing)).toBe(0);
+  });
+});

package/src/__tests__/PostgreSQLSkillStore.integration.test.ts

@@ -0,0 +1,203 @@
+import { randomUUID } from "node:crypto";
+
+import { PostgreSQLSkillStore } from "../stores/PostgreSQLSkillStore";
+
+const TEST_DATABASE_URL = process.env.TEST_DATABASE_URL;
+const describeIfDatabase = TEST_DATABASE_URL ? describe : describe.skip;
+
+describeIfDatabase("PostgreSQLSkillStore integration", () => {
+  let store: PostgreSQLSkillStore;
+  const created: Array<{ tenantId: string; id: string }> = [];
+
+  beforeAll(async () => {
+    store = new PostgreSQLSkillStore({
+      poolConfig: TEST_DATABASE_URL!,
+      autoMigrate: true,
+    });
+    await store.initialize();
+  });
+
+  afterEach(async () => {
+    while (created.length > 0) {
+      const entry = created.pop()!;
+      await store.deleteSkill(entry.tenantId, entry.id);
+    }
+  });
+
+  afterAll(async () => {
+    await store.dispose();
+  });
+
+  function tenantId(): string {
+    return `skill-${randomUUID()}`;
+  }
+
+  async function seedSkill(
+    tenant: string,
+    data: Partial<Parameters<PostgreSQLSkillStore["createSkill"]>[2]> = {},
+  ): Promise<{ id: string; tenantId: string }> {
+    const id = randomUUID();
+    created.push({ tenantId: tenant, id });
+    await store.createSkill(tenant, id, {
+      name: `skill-${id}`,
+      description: "A test skill",
+      ...data,
+    });
+    return { id, tenantId: tenant };
+  }
+
+  it("round-trips a created skill through get by id", async () => {
+    const tenant = tenantId();
+    const id = randomUUID();
+    created.push({ tenantId: tenant, id });
+
+    const createdSkill = await store.createSkill(tenant, id, {
+      name: "pdf-reader",
+      description: "Reads PDFs",
+      license: "MIT",
+      compatibility: ">=1.0.0",
+      metadata: { category: "documents", level: "1" },
+      content: "# Body",
+      subSkills: ["pdf", "ocr"],
+    });
+
+    expect(createdSkill.tenantId).toBe(tenant);
+    expect(createdSkill.name).toBe("pdf-reader");
+
+    const read = await store.getSkillById(tenant, id);
+    expect(read).not.toBeNull();
+    expect(read).toMatchObject({
+      id,
+      tenantId: tenant,
+      name: "pdf-reader",
+      description: "Reads PDFs",
+      license: "MIT",
+      compatibility: ">=1.0.0",
+      metadata: { category: "documents", level: "1" },
+      content: "# Body",
+      subSkills: ["pdf", "ocr"],
+    });
+    expect(read!.createdAt).toBeInstanceOf(Date);
+    expect(read!.updatedAt).toBeInstanceOf(Date);
+  });
+
+  it("returns null and false for an absent skill id", async () => {
+    const tenant = tenantId();
+    const missing = randomUUID();
+
+    expect(await store.getSkillById(tenant, missing)).toBeNull();
+    expect(await store.hasSkill(tenant, missing)).toBe(false);
+    expect(await store.updateSkill(tenant, missing, { name: "nope" })).toBeNull();
+    expect(await store.deleteSkill(tenant, missing)).toBe(false);
+  });
+
+  it("lists only the requested tenant's skills", async () => {
+    const tenantA = tenantId();
+    const tenantB = tenantId();
+    const a1 = await seedSkill(tenantA);
+    const a2 = await seedSkill(tenantA);
+    const b1 = await seedSkill(tenantB);
+
+    const listedA = await store.getAllSkills(tenantA);
+    const idsA = listedA.map((skill) => skill.id).sort();
+    expect(idsA).toEqual([a1.id, a2.id].sort());
+    expect(listedA.every((skill) => skill.tenantId === tenantA)).toBe(true);
+
+    const listedB = await store.getAllSkills(tenantB);
+    expect(listedB.map((skill) => skill.id)).toEqual([b1.id]);
+  });
+
+  it("updates only the provided fields and preserves the rest", async () => {
+    const tenant = tenantId();
+    const { id } = await seedSkill(tenant, {
+      name: "before",
+      description: "keep me",
+      license: "MIT",
+      metadata: { a: "1" },
+    });
+
+    const updated = await store.updateSkill(tenant, id, {
+      name: "after",
+      metadata: { a: "2" },
+    });
+
+    expect(updated).not.toBeNull();
+    expect(updated!.name).toBe("after");
+    expect(updated!.metadata).toEqual({ a: "2" });
+    expect(updated!.description).toBe("keep me");
+    expect(updated!.license).toBe("MIT");
+  });
+
+  it("reports has and deletes a skill", async () => {
+    const tenant = tenantId();
+    const { id } = await seedSkill(tenant);
+
+    expect(await store.hasSkill(tenant, id)).toBe(true);
+    expect(await store.deleteSkill(tenant, id)).toBe(true);
+    expect(await store.hasSkill(tenant, id)).toBe(false);
+    expect(await store.getSkillById(tenant, id)).toBeNull();
+    expect(await store.deleteSkill(tenant, id)).toBe(false);
+  });
+
+  it("searches by metadata within a tenant", async () => {
+    const tenant = tenantId();
+    const otherTenant = tenantId();
+    const match = await seedSkill(tenant, { metadata: { category: "finance" } });
+    await seedSkill(tenant, { metadata: { category: "other" } });
+    await seedSkill(otherTenant, { metadata: { category: "finance" } });
+
+    const results = await store.searchByMetadata(tenant, "category", "finance");
+    expect(results.map((skill) => skill.id)).toEqual([match.id]);
+  });
+
+  it("filters by compatibility within a tenant", async () => {
+    const tenant = tenantId();
+    const otherTenant = tenantId();
+    const match = await seedSkill(tenant, { compatibility: ">=2.0.0" });
+    await seedSkill(tenant, { compatibility: ">=3.0.0" });
+    await seedSkill(otherTenant, { compatibility: ">=2.0.0" });
+
+    const results = await store.filterByCompatibility(tenant, ">=2.0.0");
+    expect(results.map((skill) => skill.id)).toEqual([match.id]);
+  });
+
+  it("filters by license within a tenant", async () => {
+    const tenant = tenantId();
+    const otherTenant = tenantId();
+    const match = await seedSkill(tenant, { license: "Apache-2.0" });
+    await seedSkill(tenant, { license: "MIT" });
+    await seedSkill(otherTenant, { license: "Apache-2.0" });
+
+    const results = await store.filterByLicense(tenant, "Apache-2.0");
+    expect(results.map((skill) => skill.id)).toEqual([match.id]);
+  });
+
+  it("finds skills that list the parent in subSkills within a tenant", async () => {
+    const tenant = tenantId();
+    const otherTenant = tenantId();
+    const parent = "parent-skill";
+
+    await seedSkill(tenant, { name: parent });
+    const child = await seedSkill(tenant, {
+      name: "child-skill",
+      subSkills: [parent],
+    });
+    // Membership, not exact-array equality: the parent is one element among others.
+    const membershipChild = await seedSkill(tenant, {
+      name: "membership-child",
+      subSkills: ["other", parent],
+    });
+    const otherChild = await seedSkill(otherTenant, {
+      name: "other-child",
+      subSkills: [parent],
+    });
+
+    const results = await store.getSubSkills(tenant, parent);
+
+    expect(results.map((skill) => skill.id).sort()).toEqual(
+      [child.id, membershipChild.id].sort(),
+    );
+    expect(results.map((skill) => skill.id)).not.toContain(otherChild.id);
+    expect(results.every((skill) => skill.tenantId === tenant)).toBe(true);
+  });
+});