ThothII/backend/test/catalog-sensitivity-classifier.test.ts

import { expect, test, vi } from "vitest";
import {
  SensitivityClassifier,
  type LocalNerDetector,
  type SensitivityNerBudget,
  type SensitivityTableScan,
  type SensitivityValueSource,
} from "../src/catalog/sensitivity-classifier.js";
import type { CatalogColumn, CatalogTable, WorkspaceDatabase } from "../src/catalog/types.js";
import { CatalogConnectorError } 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 table = {
  id: "22222222-2222-4222-8222-222222222222",
  databaseId: database.id,
  name: "observations",
  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",
} satisfies CatalogTable;

function column(overrides: Partial<CatalogColumn> = {}): CatalogColumn {
  return {
    id: "33333333-3333-4333-8333-333333333333",
    tableId: table.id,
    name: "note",
    ordinalPosition: 1,
    dataType: "character varying",
    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",
    ...overrides,
  };
}

function source(scan: SensitivityTableScan): SensitivityValueSource {
  return { scanTable: vi.fn(async (_request, consume) => {
    for (const batch of scan.batches) await consume(batch);
    return scan.coverage;
  }) };
}

test("one email hidden in a generically named column makes the whole column sensitive", async () => {
  const target = column();
  const values = source({
    batches: [[
      { columnId: target.id, value: "nessun contatto", characterLength: 16 },
      { columnId: target.id, value: "mario.rossi@example.it", characterLength: 23 },
    ]],
    coverage: { kind: "complete", observedRows: 2 },
  });
  const classifier = new SensitivityClassifier(values);

  const [assessment] = await classifier.assessTable(
    { database, table, columns: [target] },
    new AbortController().signal,
  );

  expect(assessment).toMatchObject({
    columnId: target.id,
    assessment: "sensitive",
    proposedSensitive: true,
    evidence: [{ kind: "content", ruleId: "pii.email" }],
  });
});

test("one text value longer than 500 characters makes the whole column sensitive", async () => {
  const target = column({ name: "comment" });
  const values = source({
    batches: [[{ columnId: target.id, value: "x".repeat(501), characterLength: 743 }]],
    coverage: { kind: "sampled", observedRows: 1 },
  });

  const [assessment] = await new SensitivityClassifier(values).assessTable(
    { database, table, columns: [target] },
    new AbortController().signal,
  );

  expect(assessment).toMatchObject({
    assessment: "sensitive",
    proposedSensitive: true,
    evidence: [{ kind: "length", ruleId: "text.over_500_characters" }],
  });
});

test("complete coverage permits non-sensitive while empty columns remain unknown", async () => {
  const benign = column({ id: "44444444-4444-4444-8444-444444444444", name: "status" });
  const empty = column({ id: "55555555-5555-4555-8555-555555555555", name: "optional_note" });
  const humanProtected = column({
    id: "66666666-6666-4666-8666-666666666666",
    name: "category",
    sensitive: true,
  });
  const values = source({
    batches: [[
      { columnId: benign.id, value: "active", characterLength: 6 },
      { columnId: empty.id, value: null, characterLength: null },
      { columnId: humanProtected.id, value: "administrative", characterLength: 14 },
    ]],
    coverage: { kind: "complete", observedRows: 1 },
  });

  const assessments = await new SensitivityClassifier(values).assessTable(
    { database, table, columns: [benign, empty, humanProtected] },
    new AbortController().signal,
  );

  expect(assessments).toEqual([
    expect.objectContaining({ columnId: benign.id, assessment: "non_sensitive", proposedSensitive: false }),
    expect.objectContaining({
      columnId: empty.id,
      assessment: "unknown",
      proposedSensitive: false,
      evidence: [{ kind: "coverage", ruleId: "coverage.no_values" }],
    }),
    expect.objectContaining({
      columnId: humanProtected.id,
      assessment: "non_sensitive",
      proposedSensitive: false,
    }),
  ]);
});

test("sampled coverage without a match is unknown and preserves the current human flag", async () => {
  const target = column({ sensitive: true });
  const values = source({
    batches: [[{ columnId: target.id, value: "ordinary", characterLength: 8 }]],
    coverage: { kind: "sampled", observedRows: 1 },
  });

  const [assessment] = await new SensitivityClassifier(values).assessTable(
    { database, table, columns: [target] },
    new AbortController().signal,
  );

  expect(assessment).toMatchObject({
    assessment: "unknown",
    proposedSensitive: true,
    evidence: [{ kind: "coverage", ruleId: "coverage.incomplete" }],
  });
});

test("an unavailable source produces sanitized unknown evidence without losing metadata findings", async () => {
  const unresolved = column();
  const metadataMatch = column({
    id: "44444444-4444-4444-8444-444444444444",
    name: "codice_fiscale",
  });
  const values: SensitivityValueSource = {
    scanTable: vi.fn(async () => {
      throw new CatalogConnectorError("upstream detail must not escape");
    }),
  };

  const assessments = await new SensitivityClassifier(values).assessTable(
    { database, table, columns: [unresolved, metadataMatch] },
    new AbortController().signal,
  );

  expect(assessments).toEqual([
    expect.objectContaining({
      columnId: unresolved.id,
      assessment: "unknown",
      evidence: [{ kind: "coverage", ruleId: "coverage.unavailable" }],
    }),
    expect.objectContaining({
      columnId: metadataMatch.id,
      assessment: "sensitive",
      evidence: [{ kind: "metadata", ruleId: "metadata.direct_identifier" }],
    }),
  ]);
});

test("strong Italian PII metadata is sensitive even when the source column is empty", async () => {
  const target = column({ name: "codice_fiscale" });
  const values = source({
    batches: [],
    coverage: { kind: "complete", observedRows: 0 },
  });

  const [assessment] = await new SensitivityClassifier(values).assessTable(
    { database, table, columns: [target] },
    new AbortController().signal,
  );

  expect(assessment).toMatchObject({
    assessment: "sensitive",
    proposedSensitive: true,
    evidence: [{ kind: "metadata", ruleId: "metadata.direct_identifier" }],
  });
  expect(values.scanTable).not.toHaveBeenCalled();
});

test.each([
  ["RSSMRA85T10A562S", "pii.italian_fiscal_code"],
  ["IT60 X054 2811 1010 0000 0123 456", "financial.iban"],
  ["4111 1111 1111 1111", "financial.payment_card"],
  ["SWIFT DEUTDEFF500", "financial.bic"],
  ["Partita IVA 00743110157", "pii.italian_vat"],
  ["Passaporto YA1234567", "pii.passport_number"],
  ["Carta d'identità CA12345AA", "pii.identity_card"],
  ["Patente di guida U11234567A", "pii.drivers_license_number"],
  ["Chiamare +39 347 123 4567", "pii.phone_number"],
  ["Client 192.168.1.5", "network.ip_address"],
  ["Device 00:1B:44:11:3A:B7", "network.mac_address"],
  ["https://example.org/profiles/mario", "network.url"],
  ["550e8400-e29b-41d4-a716-446655440000", "pii.uuid"],
  ["AWS key AKIAIOSFODNN7EXAMPLE", "credential.access_key"],
  ["Diagnosi: carcinoma mammario con metastasi ossee", "health.clinical_term"],
  ["-----BEGIN PRIVATE KEY----- secret -----END PRIVATE KEY-----", "credential.private_key"],
  ['{"profile":{"email":"not yet supplied"}}', "pii.json_sensitive_key"],
] as const)("recognizes validated sensitive content without relying on the column name: %s", async (
  value,
  ruleId,
) => {
  const target = column();
  const values = source({
    batches: [[{ columnId: target.id, value, characterLength: value.length }]],
    coverage: { kind: "complete", observedRows: 1 },
  });

  const [assessment] = await new SensitivityClassifier(values).assessTable(
    { database, table, columns: [target] },
    new AbortController().signal,
  );

  expect(assessment).toMatchObject({
    assessment: "sensitive",
    evidence: [{ kind: "content", ruleId }],
  });
});

test("does not make a malformed email decisive", async () => {
  const target = column();
  const [assessment] = await new SensitivityClassifier(source({
    batches: [[{
      columnId: target.id,
      value: "contatto a@b..com non valido",
      characterLength: 28,
    }]],
    coverage: { kind: "complete", observedRows: 1 },
  })).assessTable(
    { database, table, columns: [target] },
    new AbortController().signal,
  );

  expect(assessment).toMatchObject({ assessment: "non_sensitive", evidence: [] });
});

test("finds a valid email after a malformed candidate in the same value", async () => {
  const target = column();
  const [assessment] = await new SensitivityClassifier(source({
    batches: [[{
      columnId: target.id,
      value: "contatto a@b..com; indirizzo valido mario.rossi@example.it",
      characterLength: 58,
    }]],
    coverage: { kind: "complete", observedRows: 1 },
  })).assessTable(
    { database, table, columns: [target] },
    new AbortController().signal,
  );

  expect(assessment).toMatchObject({
    assessment: "sensitive",
    evidence: [{ kind: "content", ruleId: "pii.email" }],
  });
});

test("optional local NER evidence can make otherwise ambiguous Italian text sensitive", async () => {
  const target = column();
  const values = source({
    batches: [[{ columnId: target.id, value: "Dimesso Mario Rossi", characterLength: 19 }]],
    coverage: { kind: "sampled", observedRows: 1 },
  });
  const detector: LocalNerDetector = {
    detect: vi.fn(async () => [{ columnId: target.id, label: "person_name", confidence: 0.91 }]),
  };

  const [assessment] = await new SensitivityClassifier(values, detector).assessTable(
    { database, table, columns: [target] },
    new AbortController().signal,
  );

  expect(detector.detect).toHaveBeenCalledWith(
    [{ columnId: target.id, text: "Dimesso Mario Rossi" }],
    expect.any(AbortSignal),
    expect.any(Number),
  );
  expect(assessment).toMatchObject({
    assessment: "sensitive",
    evidence: [{ kind: "ner", ruleId: "ner.entity", label: "person_name", confidence: 0.91 }],
  });
});

test("does not wait for an optional NER worker that is still warming", async () => {
  const target = column();
  const detector: LocalNerDetector = {
    isReady: () => false,
    detect: vi.fn(async () => [{ columnId: target.id, label: "person", confidence: 0.99 }]),
  };

  const [assessment] = await new SensitivityClassifier(source({
    batches: [[{ columnId: target.id, value: "Dimesso Mario Rossi", characterLength: 19 }]],
    coverage: { kind: "sampled", observedRows: 1 },
  }), detector).assessTable(
    { database, table, columns: [target] },
    new AbortController().signal,
  );

  expect(detector.detect).not.toHaveBeenCalled();
  expect(assessment).toMatchObject({ assessment: "unknown" });
});

test("bounds each optional NER request when an installation raises the per-table work limit", async () => {
  const columns = Array.from({ length: 17 }, (_, index) => column({
    id: `00000000-0000-4000-8000-${(index + 1).toString(16).padStart(12, "0")}`,
    name: `attribute_${index + 1}`,
    ordinalPosition: index + 1,
  }));
  const observations = columns.flatMap((item, columnIndex) => Array.from(
    { length: 8 },
    (_, valueIndex) => ({
      columnId: item.id,
      value: `ordinary-${columnIndex}-${valueIndex}`,
      characterLength: 13,
    }),
  ));
  const detector: LocalNerDetector = { detect: vi.fn(async () => []) };

  await new SensitivityClassifier(source({
    batches: [observations],
    coverage: { kind: "complete", observedRows: 8 },
  }), detector, { maxNerCandidatesPerTable: 136 }).assessTable(
    { database, table, columns },
    new AbortController().signal,
  );

  expect(detector.detect).toHaveBeenCalledTimes(2);
  expect(vi.mocked(detector.detect).mock.calls.map(([candidates]) => candidates.length)).toEqual([
    128,
    8,
  ]);
});

test("limits default NER work to two candidates spread across a wide table", async () => {
  const columns = Array.from({ length: 10 }, (_, index) => column({
    id: `10000000-0000-4000-8000-${(index + 1).toString(16).padStart(12, "0")}`,
    name: `attribute_${index + 1}`,
    ordinalPosition: index + 1,
  }));
  const detector: LocalNerDetector = { detect: vi.fn(async () => []) };

  await new SensitivityClassifier(source({
    batches: [columns.flatMap((item, columnIndex) => [0, 1].map((valueIndex) => ({
      columnId: item.id,
      value: `ordinary-${columnIndex}-${valueIndex}`,
      characterLength: 13,
    })))],
    coverage: { kind: "complete", observedRows: 2 },
  }), detector).assessTable(
    { database, table, columns },
    new AbortController().signal,
  );

  expect(detector.detect).toHaveBeenCalledOnce();
  const submitted = vi.mocked(detector.detect).mock.calls[0]![0];
  expect(submitted).toHaveLength(2);
  expect(new Set(submitted.map((candidate) => candidate.columnId)).size).toBe(2);
});

test("shares a bounded NER time allowance across tables in one analysis run", async () => {
  const target = column();
  const values = source({
    batches: [[{ columnId: target.id, value: "Dimesso Mario Rossi", characterLength: 19 }]],
    coverage: { kind: "sampled", observedRows: 1 },
  });
  const detector: LocalNerDetector = {
    detect: vi.fn(async () => {
      await new Promise((resolve) => setTimeout(resolve, 20));
      return [];
    }),
  };
  const classifier = new SensitivityClassifier(values, detector);
  const nerBudget: SensitivityNerBudget = { remainingMs: 1 };

  await classifier.assessTable(
    { database, table, columns: [target] },
    new AbortController().signal,
    Date.now() + 1_000,
    nerBudget,
  );
  await classifier.assessTable(
    { database, table, columns: [target] },
    new AbortController().signal,
    Date.now() + 1_000,
    nerBudget,
  );

  expect(detector.detect).toHaveBeenCalledOnce();
  expect(nerBudget.remainingMs).toBe(0);
});

test("uninterpretable binary content remains unknown after complete coverage", async () => {
  const target = column({ dataType: "bytea" });
  const values = source({
    batches: [[{ columnId: target.id, value: "\\xdeadbeef", characterLength: 10 }]],
    coverage: { kind: "complete", observedRows: 1 },
  });

  const [assessment] = await new SensitivityClassifier(values).assessTable(
    { database, table, columns: [target] },
    new AbortController().signal,
  );

  expect(assessment).toMatchObject({
    assessment: "unknown",
    proposedSensitive: false,
    evidence: [{ kind: "coverage", ruleId: "coverage.unsupported_type" }],
  });
});