feat: classify sensitive columns locally

2026-09-03 02:11:13 +02:00
parent 7b87e95427
commit f114d0065a
57 changed files with 4038 additions and 1149 deletions
@@ -0,0 +1,99 @@
+import { expect, test, vi } from "vitest";
+import {
+  SensitivityAnalysisInterruptedError,
+  SensitivityAnalysisService,
+} from "../src/catalog/sensitivity-analysis-service.js";
+import { SensitivityAnalysisRunner } from "../src/catalog/sensitivity-analysis-runner.js";
+import type { SensitivityClassifier } from "../src/catalog/sensitivity-classifier.js";
+import type {
+  CatalogRepository,
+  SensitivityAnalysisRun,
+  WorkspaceDatabase,
+} from "../src/catalog/types.js";
+
+const database = {
+  id: "11111111-1111-4111-8111-111111111111",
+  workspaceId: "psd-clinical",
+  engine: "postgres",
+  databaseName: "warehouse",
+  schema: "public",
+  version: 1,
+  createdAt: "2026-09-02T08:00:00Z",
+  updatedAt: "2026-09-02T08:00:00Z",
+  connectionStatus: "reachable",
+  binding: { transport: "postgres_direct", host: "db.internal", port: 5432, username: "reader" },
+} satisfies WorkspaceDatabase;
+
+const running: SensitivityAnalysisRun = {
+  id: "22222222-2222-4222-8222-222222222222",
+  databaseId: database.id,
+  scope: "all",
+  engine: "local",
+  modelId: null,
+  policyVersion: "sensitivity-v1",
+  status: "running",
+  total: 0,
+  suggestedSensitive: 0,
+  suggestedNonSensitive: 0,
+  unknown: 0,
+  inputTokens: 0,
+  cacheReadTokens: 0,
+  outputTokens: 0,
+  createdAt: "2026-09-02T08:00:00Z",
+  startedAt: "2026-09-02T08:00:00Z",
+  updatedAt: "2026-09-02T08:00:00Z",
+  finishedAt: null,
+  errorSummary: null,
+};
+
+test("stops catalog selection when the request expires during a catalog read", async () => {
+  const controller = new AbortController();
+  const listTables = vi.fn();
+  const repository = {
+    get: vi.fn(async () => {
+      controller.abort();
+      return database;
+    }),
+    listTables,
+  } as unknown as CatalogRepository;
+  const classifier = { assessTable: vi.fn() } as unknown as SensitivityClassifier;
+  const analysis = new SensitivityAnalysisService(repository, classifier);
+
+  await expect(analysis.analyze(
+    database.id,
+    "all",
+    [],
+    controller.signal,
+  )).rejects.toBeInstanceOf(SensitivityAnalysisInterruptedError);
+  expect(listTables).not.toHaveBeenCalled();
+  expect(classifier.assessTable).not.toHaveBeenCalled();
+});
+
+test("marks a created run interrupted if the request deadline expires during persistence", async () => {
+  const controller = new AbortController();
+  const update = vi.fn(async (_runId: string, changes: Partial<SensitivityAnalysisRun>) => ({
+    ...running,
+    ...changes,
+  }));
+  const repository = {
+    get: vi.fn(async () => database),
+    createSensitivityAnalysisRun: vi.fn(async () => {
+      controller.abort();
+      return running;
+    }),
+    updateSensitivityAnalysisRun: update,
+    appendSensitivityAnalysisEvent: vi.fn(async () => undefined),
+  } as unknown as CatalogRepository;
+  const analysis = { analyze: vi.fn() } as unknown as SensitivityAnalysisService;
+  const runner = new SensitivityAnalysisRunner(repository, analysis);
+
+  await expect(runner.run(database.id, "all", [], controller.signal))
+    .rejects.toBeInstanceOf(SensitivityAnalysisInterruptedError);
+  expect(analysis.analyze).not.toHaveBeenCalled();
+  expect(update).toHaveBeenCalledWith(running.id, expect.objectContaining({
+    status: "interrupted",
+    total: 0,
+    unknown: 0,
+    errorSummary: "Local sensitivity analysis reached its time limit.",
+  }));
+});