feat: add AI catalog description generation

2026-08-29 16:42:56 +02:00
parent b0afba81ca
commit 376dd5a09d
76 changed files with 14860 additions and 102 deletions
@@ -0,0 +1,111 @@
+import { expect, test, vi } from "vitest";
+import {
+  PostgresDescriptionSourceSampler,
+  type DescriptionSourceSamplingTarget,
+} from "../src/catalog/description-source-sampler.js";
+import type {
+  CatalogDatabaseClient,
+  CatalogPostgresAccess,
+} from "../src/catalog/postgres-access.js";
+import type { WorkspaceDatabase } from "../src/catalog/types.js";
+
+const database: WorkspaceDatabase = {
+  id: "11111111-1111-4111-8111-111111111111",
+  workspaceId: "psd-clinical",
+  engine: "postgres",
+  databaseName: "warehouse",
+  schema: 'clinical"data',
+  version: 1,
+  createdAt: "2026-08-28T08:00:00Z",
+  updatedAt: "2026-08-28T08:00:00Z",
+  connectionStatus: "reachable",
+  binding: {
+    transport: "postgres_direct",
+    host: "db.internal",
+    port: 5432,
+    username: "reader",
+  },
+};
+
+const target: DescriptionSourceSamplingTarget = {
+  targetId: "22222222-2222-4222-8222-222222222222",
+  tableName: 'patient"facts',
+  columnNames: ['status"code', "ward"],
+};
+
+test("samples at most five source rows and five distinct non-null examples in a read-only transaction", async () => {
+  const query = vi.fn(async (sql: string) => {
+    if (!sql.startsWith("SELECT")) return { rows: [] };
+    return {
+      rows: [
+        { 'status"code': "active", ward: null },
+        { 'status"code': "pending", ward: "A" },
+        { 'status"code': "closed", ward: "A" },
+        { 'status"code': "transferred", ward: "B" },
+        { 'status"code': "unknown", ward: "C" },
+        { 'status"code': "must-not-be-sampled", ward: "D" },
+      ],
+    };
+  });
+  const end = vi.fn(async () => undefined);
+  const access: CatalogPostgresAccess = {
+    connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
+  };
+  const sampler = new PostgresDescriptionSourceSampler(access);
+  const controller = new AbortController();
+
+  const samples = await sampler.sample(database, [target], controller.signal);
+
+  expect(samples).toEqual([{
+    targetId: target.targetId,
+    tableName: target.tableName,
+    rows: [
+      { fields: [{ name: 'status"code', value: "active" }, { name: "ward", value: null }] },
+      { fields: [{ name: 'status"code', value: "pending" }, { name: "ward", value: "A" }] },
+      { fields: [{ name: 'status"code', value: "closed" }, { name: "ward", value: "A" }] },
+      { fields: [{ name: 'status"code', value: "transferred" }, { name: "ward", value: "B" }] },
+      { fields: [{ name: 'status"code', value: "unknown" }, { name: "ward", value: "C" }] },
+    ],
+    representativeValues: [
+      {
+        column: 'status"code',
+        values: ["active", "pending", "closed", "transferred"],
+      },
+      { column: "ward", values: ["A"] },
+    ],
+  }]);
+  expect(access.connect).toHaveBeenCalledWith(database, controller.signal);
+  expect(samples[0]!.representativeValues.flatMap((entry) => entry.values)).toHaveLength(5);
+  expect(query.mock.calls).toEqual([
+    ["BEGIN TRANSACTION READ ONLY", []],
+    [
+      'SELECT LEFT(("status""code")::text, $1) AS "status""code", LEFT(("ward")::text, $1) AS "ward" FROM "clinical""data"."patient""facts" LIMIT $2',
+      [256, 5],
+    ],
+    ["ROLLBACK", []],
+  ]);
+  expect(query.mock.calls.map(([sql]) => String(sql).split(" ")[0])).toEqual([
+    "BEGIN",
+    "SELECT",
+    "ROLLBACK",
+  ]);
+  expect(end).toHaveBeenCalledOnce();
+});
+
+test("rolls back and closes the source connection when sampling fails", async () => {
+  const query = vi.fn(async (sql: string) => {
+    if (sql.startsWith("SELECT")) throw new Error("distinctive-source-secret");
+    return { rows: [] };
+  });
+  const end = vi.fn(async () => undefined);
+  const access: CatalogPostgresAccess = {
+    connect: vi.fn(async () => ({ query, end }) as CatalogDatabaseClient),
+  };
+  const sampler = new PostgresDescriptionSourceSampler(access);
+  const controller = new AbortController();
+
+  await expect(sampler.sample(database, [target], controller.signal)).rejects.toThrow();
+
+  expect(query).toHaveBeenCalledWith("ROLLBACK", []);
+  expect(end).toHaveBeenCalledOnce();
+});