mptyl/ThothII
1
0
Fork 0
Code Issues 25 Pull Requests Actions Packages Projects Releases 1 Wiki Activity

feat: classify sensitive columns locally

Browse Source
This commit is contained in:
Codex
2026-09-03 02:11:13 +02:00
parent 7b87e95427
commit f114d0065a
57 changed files with 4038 additions and 1149 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/test/catalog-description-generation-routes.test.ts
+68 -256
View File
@@ -14,6 +14,7 @@ import {
type ModelCompletionRequest,
} from "../src/catalog/model-completer.js";
import { CatalogOperationCoordinator } from "../src/catalog/operation-coordinator.js";
import type { SensitivityValueSource } from "../src/catalog/sensitivity-classifier.js";
import type {
CatalogDatabaseClient,
CatalogPostgresAccess,
@@ -69,6 +70,16 @@ async function setup(
sample: vi.fn(async () => []),
},
catalogPostgresAccess?: CatalogPostgresAccess,
sensitivityValueSource: SensitivityValueSource = {
scanTable: vi.fn(async (request, consume) => {
await consume(request.columns.map((column) => ({
columnId: column.id,
value: "ordinary",
characterLength: 8,
})));
return { kind: "complete", observedRows: 1 };
}),
},
) {
const repository = new MemoryCatalogRepository();
const database = await repository.create({
@@ -115,10 +126,11 @@ async function setup(
catalogOperationCoordinator: operations,
metadataGenerationModels: models(),
modelCompleter,
sensitivityValueSource,
...(descriptionSourceSampler ? { descriptionSourceSampler } : {}),
...(catalogPostgresAccess ? { catalogPostgresAccess } : {}),
});
return { app, repository, database, table, column, operations };
return { app, repository, database, table, column, operations, sensitivityValueSource };
}
async function waitForTerminalRun(app: ReturnType<typeof buildApp>, runId: string) {
@@ -136,22 +148,15 @@ async function waitForTerminalRun(app: ReturnType<typeof buildApp>, runId: strin
throw new Error(`Description Generation Run ${runId} did not finish`);
}
test("suggests sensitive flags from structural metadata without persisting them", async () => {
const modelCompleter = {
complete: vi.fn(async () => JSON.stringify({
suggestions: [{ columnId: expect.any(String), sensitive: true }],
})),
};
test("assesses sensitive flags locally without persisting them or calling an LLM", async () => {
const modelCompleter: ModelCompleter = { complete: vi.fn(async () => "unused") };
const { app, repository, database, table, column } = await setup(modelCompleter);
modelCompleter.complete.mockResolvedValueOnce(JSON.stringify({
suggestions: [{ columnId: column.id, sensitive: true }],
}));
try {
const response = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { modelId: configuredModel.id, scope: "all" },
payload: { scope: "all" },
});
expect(response.statusCode).toBe(200);
@@ -160,11 +165,14 @@ test("suggests sensitive flags from structural metadata without persisting them"
run: {
databaseId: database.id,
scope: "all",
modelId: configuredModel.id,
engine: "local",
modelId: null,
policyVersion: "sensitivity-v1",
status: "completed",
total: 1,
suggestedSensitive: 1,
suggestedNonSensitive: 0,
unknown: 0,
errorSummary: null,
},
suggestions: [{
@@ -175,6 +183,8 @@ test("suggests sensitive flags from structural metadata without persisting them"
version: column.version,
currentSensitive: false,
sensitive: true,
assessment: "sensitive",
evidence: [{ kind: "metadata", ruleId: "metadata.direct_identifier" }],
}],
});
expect(await repository.getColumn(database.id, column.tableId, column.id))
@@ -207,49 +217,56 @@ test("suggests sensitive flags from structural metadata without persisting them"
runId: responseBody.run.id,
sequence: 1,
level: "info",
message: "Sensitive-field suggestion generation started.",
message: "Local sensitivity analysis started.",
},
{
runId: responseBody.run.id,
sequence: 2,
level: "info",
message: "Classified 1 of 1 columns.",
message: "Assessed 1 of 1 columns locally.",
},
{
runId: responseBody.run.id,
sequence: 3,
level: "info",
message: "Sensitive-field suggestion generation completed for 1 column.",
message: "Local sensitivity analysis completed for 1 column.",
},
]);
const request = modelCompleter.complete.mock.calls[0]![0] as ModelCompletionRequest;
const prompt = request.messages.map((message) => message.content).join("\n");
expect(prompt).toContain("patients");
expect(prompt).toContain("birth_date");
expect(prompt).toContain("date");
expect(prompt).not.toContain("Patient date of birth");
expect(prompt).not.toContain("test-provider-secret");
expect(modelCompleter.complete).not.toHaveBeenCalled();
} finally {
await app.close();
}
});
test("limits sensitive-data suggestions to the selected tables or columns", async () => {
const modelCompleter: ModelCompleter = {
complete: vi.fn(async (request) => {
const payload = JSON.parse(request.messages.find((message) => message.role === "user")!.content) as {
columns: Array<{ columnId: string; column: string }>;
};
return JSON.stringify({
suggestions: payload.columns.map((column) => ({
columnId: column.columnId,
sensitive: column.column.includes("name") || column.column.includes("note"),
})),
});
}),
};
const { app, repository, database } = await setup(modelCompleter);
test("stops sensitivity analysis at the HTTP deadline without creating a review", async () => {
const controller = new AbortController();
controller.abort();
const timeout = vi.spyOn(AbortSignal, "timeout").mockReturnValue(controller.signal);
const { app, repository, database } = await setup({ complete: vi.fn(async () => "unused") });
try {
const response = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { scope: "all" },
});
expect(response.statusCode).toBe(504);
expect(response.json()).toEqual({
code: "sensitivity_analysis_timeout",
message: "Sensitivity analysis reached its time limit. No assessments were applied.",
});
expect(await repository.listSensitivityAnalysisRuns()).toEqual([]);
} finally {
timeout.mockRestore();
await app.close();
}
});
test("limits sensitivity analysis to the selected tables or columns", async () => {
const modelCompleter: ModelCompleter = { complete: vi.fn(async () => "unused") };
const { app, repository, database, sensitivityValueSource } = await setup(modelCompleter);
await repository.applySchemaSync(database.id, database.version, "all", [], {
schemaVersion: 1,
capabilities: { tables: "available", columns: "available", relationships: "available" },
@@ -281,7 +298,6 @@ test("limits sensitive-data suggestions to the selected tables or columns", asyn
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: {
modelId: configuredModel.id,
scope: "selected_tables",
targetIds: [visits.id, patients.id],
},
@@ -301,7 +317,6 @@ test("limits sensitive-data suggestions to the selected tables or columns", asyn
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: {
modelId: configuredModel.id,
scope: "selected_columns",
targetIds: [clinicalNote.id, status.id],
},
@@ -313,44 +328,41 @@ test("limits sensitive-data suggestions to the selected tables or columns", asyn
expect.objectContaining({ tableId: visits.id, columnId: clinicalNote.id, sensitive: true }),
]));
const prompts = vi.mocked(modelCompleter.complete).mock.calls.map(([request]) => (
JSON.parse(request.messages.find((message) => message.role === "user")!.content) as {
columns: Array<{ columnId: string }>;
}
));
expect(prompts[0]!.columns.map((column) => column.columnId).sort()).toEqual(
[...patientColumns, ...visitColumns].map((column) => column.id).sort(),
);
expect(prompts[0]!.columns.map((column) => column.columnId)).not.toContain(billingColumns[0]!.id);
expect(prompts[1]!.columns.map((column) => column.columnId).sort()).toEqual(
[status.id, clinicalNote.id].sort(),
const scannedColumnIds = vi.mocked(sensitivityValueSource.scanTable).mock.calls.flatMap(
([request]) => request.columns.map((column) => column.id),
);
expect(scannedColumnIds).toEqual([status.id, status.id]);
expect(scannedColumnIds).not.toContain(patientColumns.find(
(column) => column.name === "patient_name",
)!.id);
expect(scannedColumnIds).not.toContain(clinicalNote.id);
expect(scannedColumnIds).not.toContain(billingColumns[0]!.id);
expect(modelCompleter.complete).not.toHaveBeenCalled();
} finally {
await app.close();
}
});
test("explains invalid sensitive-data suggestion selections without calling the model", async () => {
test("explains invalid sensitivity-analysis selections without reading source values", async () => {
const modelCompleter: ModelCompleter = { complete: vi.fn(async () => "unused") };
const { app, database, table } = await setup(modelCompleter);
const { app, database, table, sensitivityValueSource } = await setup(modelCompleter);
try {
const empty = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { modelId: configuredModel.id, scope: "selected_tables", targetIds: [] },
payload: { scope: "selected_tables", targetIds: [] },
});
expect(empty.statusCode).toBe(400);
expect(empty.json()).toEqual({
code: "sensitive_data_suggestion_request_invalid",
message: "Choose a database, one or more tables, or one or more columns to classify.",
message: "Choose a database, one or more tables, or one or more columns to assess.",
});
const duplicate = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: {
modelId: configuredModel.id,
scope: "selected_tables",
targetIds: [table.id, table.id],
},
@@ -365,7 +377,6 @@ test("explains invalid sensitive-data suggestion selections without calling the
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: {
modelId: configuredModel.id,
scope: "selected_tables",
targetIds: ["00000000-0000-4000-8000-000000000001"],
},
@@ -380,7 +391,6 @@ test("explains invalid sensitive-data suggestion selections without calling the
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: {
modelId: configuredModel.id,
scope: "selected_columns",
targetIds: ["00000000-0000-4000-8000-000000000002"],
},
@@ -391,205 +401,7 @@ test("explains invalid sensitive-data suggestion selections without calling the
message: "One or more selected Catalog Columns were not found in this database.",
});
expect(modelCompleter.complete).not.toHaveBeenCalled();
} finally {
await app.close();
}
});
test("batches sensitive-data suggestions for schemas larger than one helper message", async () => {
const maxHelperMessageBytes = 64 * 1024;
const seenColumnIds: string[] = [];
const modelCompleter: ModelCompleter = {
complete: vi.fn(async (request) => {
const userMessage = request.messages.find((message) => message.role === "user")!;
expect(Buffer.byteLength(userMessage.content, "utf8")).toBeLessThanOrEqual(maxHelperMessageBytes);
const payload = JSON.parse(userMessage.content) as {
columns: Array<{ columnId: string; column: string }>;
};
expect(payload.columns.length).toBeLessThanOrEqual(10);
seenColumnIds.push(...payload.columns.map((column) => column.columnId));
return JSON.stringify({
suggestions: payload.columns.map((column) => ({
columnId: column.columnId,
sensitive: column.column.endsWith("_private"),
})),
});
}),
};
const { app, repository, database } = await setup(modelCompleter);
const columnCount = 900;
await repository.applySchemaSync(database.id, database.version, "all", [], {
schemaVersion: 1,
capabilities: { tables: "available", columns: "available", relationships: "available" },
tables: [{ name: "wide_table", sourceComment: null }],
columns: Array.from({ length: columnCount }, (_, index) => ({
tableName: "wide_table",
name: `field_${index.toString().padStart(4, "0")}${index % 10 === 0 ? "_private" : ""}`,
ordinalPosition: index + 1,
dataType: "character varying(255)",
isNullable: true,
defaultExpression: null,
primaryKeyPosition: null,
sourceComment: null,
})),
relationships: [],
});
const wideTable = (await repository.listTables(database.id)).find((table) => table.name === "wide_table")!;
const expectedColumnIds = (await repository.listColumns(database.id, wideTable.id)).map((column) => column.id);
try {
const response = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { modelId: configuredModel.id, scope: "all" },
});
expect(response.statusCode).toBe(200);
const suggestions = response.json().suggestions as Array<{
columnName: string;
currentSensitive: boolean;
sensitive: boolean;
}>;
expect(suggestions).toHaveLength(columnCount);
expect(suggestions).toEqual(expect.arrayContaining([
expect.objectContaining({ columnName: "field_0000_private", currentSensitive: false, sensitive: true }),
expect.objectContaining({ columnName: "field_0001", currentSensitive: false, sensitive: false }),
]));
expect(vi.mocked(modelCompleter.complete).mock.calls.length).toBeGreaterThan(1);
expect(seenColumnIds.slice().sort()).toEqual(expectedColumnIds.slice().sort());
expect(new Set(seenColumnIds).size).toBe(columnCount);
} finally {
await app.close();
}
});
test("retries one invalid sensitive-data classification before returning the review draft", async () => {
const modelCompleter: ModelCompleter = {
complete: vi.fn(async () => "unused"),
};
const { app, database, column } = await setup(modelCompleter);
vi.mocked(modelCompleter.complete)
.mockResolvedValueOnce("not-json")
.mockResolvedValueOnce(JSON.stringify({
suggestions: [{ columnId: column.id, sensitive: true }],
}));
try {
const response = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { modelId: configuredModel.id, scope: "all" },
});
expect(response.statusCode).toBe(200);
expect(response.json().suggestions).toEqual([
expect.objectContaining({ columnId: column.id, sensitive: true }),
]);
expect(modelCompleter.complete).toHaveBeenCalledTimes(2);
} finally {
await app.close();
}
});
test.each(["malformed", "incomplete", "duplicate"] as const)(
"fails safely when sensitive-data suggestions are %s",
async (kind) => {
const modelCompleter: ModelCompleter = {
complete: vi.fn(async () => "unused"),
};
const { app, repository, database, column } = await setup(modelCompleter);
const rawResponse = kind === "malformed"
? "RAW_PROVIDER_RESPONSE_DO_NOT_EXPOSE_{"
: kind === "incomplete"
? JSON.stringify({ suggestions: [] })
: JSON.stringify({
suggestions: [
{ columnId: column.id, sensitive: true },
{ columnId: column.id, sensitive: true },
],
});
vi.mocked(modelCompleter.complete).mockResolvedValueOnce(rawResponse);
try {
const response = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { modelId: configuredModel.id, scope: "all" },
});
expect(response.statusCode).toBe(502);
expect(response.json()).toEqual({
code: "sensitive_data_suggestion_invalid_response",
message: "The LLM returned an incomplete or invalid classification. No suggestions were applied.",
});
expect(response.body).not.toContain(rawResponse);
expect(await repository.getColumn(database.id, column.tableId, column.id))
.toMatchObject({ sensitive: false });
} finally {
await app.close();
}
},
);
test("explains a sensitive-data suggestion provider failure without exposing provider details", async () => {
const modelCompleter: ModelCompleter = {
complete: vi.fn(async () => {
throw new ModelCompletionProviderError();
}),
};
const { app, repository, database, column } = await setup(modelCompleter);
try {
const response = await app.inject({
method: "POST",
url: `/catalog/databases/${database.id}/sensitive-data-suggestions`,
payload: { modelId: configuredModel.id, scope: "all" },
});
expect(response.statusCode).toBe(502);
expect(response.json()).toEqual({
code: "sensitive_data_suggestion_provider_unavailable",
message: "The selected LLM service could not complete the request. No suggestions were applied.",
});
expect(response.body).not.toContain("model completion failed");
expect(await repository.getColumn(database.id, column.tableId, column.id))
.toMatchObject({ sensitive: false });
const history = await app.inject({
method: "GET",
url: "/catalog/sensitive-data-suggestion-runs",
});
expect(history.statusCode).toBe(200);
const [failedRun] = history.json();
expect(failedRun).toMatchObject({
databaseId: database.id,
status: "failed",
total: 1,
suggestedSensitive: 0,
suggestedNonSensitive: 0,
errorSummary: "Sensitive-field suggestion generation failed.",
});
const events = await app.inject({
method: "GET",
url: `/catalog/sensitive-data-suggestion-runs/${failedRun.id}/events-list`,
});
expect(events.statusCode).toBe(200);
expect(events.json()).toMatchObject([
{
runId: failedRun.id,
sequence: 1,
level: "info",
message: "Sensitive-field suggestion generation started.",
},
{
runId: failedRun.id,
sequence: 2,
level: "error",
message: "Sensitive-field suggestion generation failed.",
},
]);
expect(events.body).not.toContain("model completion failed");
expect(sensitivityValueSource.scanTable).not.toHaveBeenCalled();
} finally {
await app.close();
}