feat: sample sensitive columns progressively

backend/test/catalog-sensitivity-analysis.test.ts

@@ -6,7 +6,9 @@ import {
 import { SensitivityAnalysisRunner } from "../src/catalog/sensitivity-analysis-runner.js";
 import type { SensitivityClassifier } from "../src/catalog/sensitivity-classifier.js";
 import type {
+  CatalogColumn,
   CatalogRepository,
+  CatalogTable,
   SensitivityAnalysisRun,
   WorkspaceDatabase,
 } from "../src/catalog/types.js";
@@ -24,13 +26,54 @@ const database = {
   binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
 } satisfies WorkspaceDatabase;
 
+function catalogTable(id: string, name: string): CatalogTable {
+  return {
+    id,
+    databaseId: database.id,
+    name,
+    sourceComment: null,
+    description: null,
+    generatedDescription: null,
+    lastSyncedDatabaseVersion: 1,
+    lastSyncedAt: "2026-09-02T08:00:00Z",
+    version: 1,
+    createdAt: "2026-09-02T08:00:00Z",
+    updatedAt: "2026-09-02T08:00:00Z",
+  };
+}
+
+function catalogColumn(id: string, tableId: string, name: string): CatalogColumn {
+  return {
+    id,
+    tableId,
+    name,
+    ordinalPosition: 1,
+    dataType: "text",
+    isNullable: true,
+    defaultExpression: null,
+    primaryKeyPosition: null,
+    isPrimaryKey: false,
+    isForeignKey: false,
+    foreignKeyCount: 0,
+    sourceComment: null,
+    description: null,
+    generatedDescription: null,
+    sensitive: false,
+    lastSyncedDatabaseVersion: 1,
+    lastSyncedAt: "2026-09-02T08:00:00Z",
+    version: 1,
+    createdAt: "2026-09-02T08:00:00Z",
+    updatedAt: "2026-09-02T08:00:00Z",
+  };
+}
+
 const running: SensitivityAnalysisRun = {
   id: "22222222-2222-4222-8222-222222222222",
   databaseId: database.id,
   scope: "all",
   engine: "local",
   modelId: null,
-  policyVersion: "sensitivity-v1",
+  policyVersion: "sensitivity-v2",
   status: "running",
   total: 0,
   suggestedSensitive: 0,
@@ -56,7 +99,7 @@ test("stops catalog selection when the request expires during a catalog read", a
     }),
     listTables,
   } as unknown as CatalogRepository;
-  const classifier = { assessTable: vi.fn() } as unknown as SensitivityClassifier;
+  const classifier = { assess: vi.fn() } as unknown as SensitivityClassifier;
   const analysis = new SensitivityAnalysisService(repository, classifier);
 
   await expect(analysis.analyze(
@@ -66,7 +109,71 @@ test("stops catalog selection when the request expires during a catalog read", a
     controller.signal,
   )).rejects.toBeInstanceOf(SensitivityAnalysisInterruptedError);
   expect(listTables).not.toHaveBeenCalled();
-  expect(classifier.assessTable).not.toHaveBeenCalled();
+  expect(classifier.assess).not.toHaveBeenCalled();
+});
+
+test("classifies all selected tables in one breadth-first run and reports coverage", async () => {
+  const firstTable = catalogTable("33333333-3333-4333-8333-333333333333", "patients");
+  const secondTable = catalogTable("44444444-4444-4444-8444-444444444444", "encounters");
+  const firstColumn = catalogColumn(
+    "55555555-5555-4555-8555-555555555555",
+    firstTable.id,
+    "status",
+  );
+  const secondColumn = catalogColumn(
+    "66666666-6666-4666-8666-666666666666",
+    secondTable.id,
+    "note",
+  );
+  const repository = {
+    get: vi.fn(async () => database),
+    listTables: vi.fn(async () => [firstTable, secondTable]),
+    listColumns: vi.fn(async (_databaseId: string, tableId: string) => (
+      tableId === firstTable.id ? [firstColumn] : [secondColumn]
+    )),
+  } as unknown as CatalogRepository;
+  const assess = vi.fn(async () => [
+    {
+      columnId: firstColumn.id,
+      assessment: "non_sensitive" as const,
+      proposedSensitive: false,
+      evidence: [{ kind: "coverage" as const, ruleId: "coverage.sampled_1000" }],
+      observedValues: 1_000,
+      coverage: "sampled" as const,
+    },
+    {
+      columnId: secondColumn.id,
+      assessment: "sensitive" as const,
+      proposedSensitive: true,
+      evidence: [{ kind: "content" as const, ruleId: "pii.email" }],
+      observedValues: 12,
+      coverage: "sampled" as const,
+    },
+  ]);
+  const classifier = { assess } as unknown as SensitivityClassifier;
+  const onPrepared = vi.fn();
+  const onProgress = vi.fn();
+
+  const suggestions = await new SensitivityAnalysisService(repository, classifier).analyze(
+    database.id,
+    "all",
+    [],
+    new AbortController().signal,
+    onPrepared,
+    onProgress,
+  );
+
+  expect(assess).toHaveBeenCalledOnce();
+  expect(assess.mock.calls[0]![0]).toEqual([
+    { database, table: firstTable, columns: [firstColumn] },
+    { database, table: secondTable, columns: [secondColumn] },
+  ]);
+  expect(onPrepared).toHaveBeenCalledWith(2);
+  expect(onProgress.mock.calls.map(([processed]) => processed)).toEqual([1, 2]);
+  expect(suggestions).toEqual([
+    expect.objectContaining({ columnId: firstColumn.id, sensitive: false, coverage: "sampled" }),
+    expect.objectContaining({ columnId: secondColumn.id, sensitive: true, coverage: "sampled" }),
+  ]);
 });
 
 test("marks a created run interrupted if the request deadline expires during persistence", async () => {
@@ -94,6 +201,6 @@ test("marks a created run interrupted if the request deadline expires during per
     status: "interrupted",
     total: 0,
     unknown: 0,
-    errorSummary: "Local sensitivity analysis reached its time limit.",
+    errorSummary: "Local sensitivity analysis was interrupted before completion.",
   }));
 });