diff --git a/CONTEXT.md b/CONTEXT.md index f5a7d008..70908c4a 100644 --- a/CONTEXT.md +++ b/CONTEXT.md @@ -373,7 +373,9 @@ per le nuove colonne. **Sensitive Data Policy** — La regola che applica il Sensitive Data Flag ai Catalog Sample: valori sintetici per una colonna protetta, valori reali per una colonna non protetta. L'AI può -suggerire il flag dai soli metadati tecnici, ma soltanto l'utente lo imposta. +suggerire il flag dai soli metadati tecnici di un database, delle tabelle o delle colonne +esplicitamente selezionate; le richieste ampie vengono divise in batch bounded, ma soltanto +l'utente imposta i flag dopo aver rivisto la proposta completa. _Avoid_: PII filter, sample filter **Introspection Capability** — Una categoria di struttura fisica che una Database Binding diff --git a/PROJECT_STATE.md b/PROJECT_STATE.md index 506cbfb2..9dbdfbde 100644 --- a/PROJECT_STATE.md +++ b/PROJECT_STATE.md @@ -111,7 +111,10 @@ deferred to their dedicated slices. AI Description Generation uses the catalog's human-owned Sensitive Data Flag. The flag defaults to `false`, including for newly synchronized columns. An administrator may request an AI proposal based -only on structural metadata, but it remains an unsaved draft until the human reviews and saves it. +only on structural metadata for one selected database, selected tables, or selected columns. The +backend divides large scopes into deterministic model requests of at most ten columns, also bounded +by helper message size, and combines their results, but the proposal remains an unsaved draft until +the human reviews and saves it. For unprotected columns, up to five source rows and five representative non-null values may be sent transiently to the configured model provider. Protected columns are omitted from source reads and replaced in the prompt by deterministic plausible values derived only from their metadata. Existing diff --git a/backend/src/catalog/sensitive-data-suggester.ts b/backend/src/catalog/sensitive-data-suggester.ts index c99c6afb..d6eb17a9 100644 --- a/backend/src/catalog/sensitive-data-suggester.ts +++ b/backend/src/catalog/sensitive-data-suggester.ts @@ -1,28 +1,74 @@ import { z } from "zod"; import type { MetadataGenerationModels } from "./metadata-generation-models.js"; -import type { ModelCompleter } from "./model-completer.js"; -import type { CatalogRepository } from "./types.js"; +import type { ModelCompleter, ModelCompletionMessage } from "./model-completer.js"; +import type { CatalogColumn, CatalogRepository, CatalogTable } from "./types.js"; -const MAX_COLUMNS = 10_000; +// The helper accepts at most 64 KiB per message. Keep the same safety margin used by +// Description Generation so UTF-8 structural metadata never reaches that hard limit. +const MAX_USER_MESSAGE_BYTES = 60 * 1024; +// Preserve ThothAI's proven completion granularity: small batches keep generation time and +// structured-output accuracy predictable even when the helper byte limit would allow much more. +const MAX_COLUMNS_PER_BATCH = 10; const responseSchema = z.object({ suggestions: z.array(z.object({ columnId: z.uuid(), sensitive: z.boolean(), - }).strict()).max(MAX_COLUMNS), + }).strict()), }).strict(); +export type SensitiveDataSuggestionScope = "all" | "selected_tables" | "selected_columns"; + +interface StructuralColumn { + columnId: string; + tableId: string; + table: string; + column: string; + dataType: string; + nullable: boolean; + primaryKey: boolean; + foreignKey: boolean; + version: number; + currentSensitive: boolean; +} + export interface SensitiveDataSuggestion { columnId: string; + tableId: string; + tableName: string; + columnName: string; + version: number; + currentSensitive: boolean; sensitive: boolean; } export class SensitiveDataSuggestionTargetNotFoundError extends Error { - constructor() { - super("database not found"); + constructor(readonly target: "database" | "table" | "column") { + super(`${target} not found`); this.name = "SensitiveDataSuggestionTargetNotFoundError"; } } +export class SensitiveDataSuggestionDuplicateTargetIdsError extends Error { + constructor() { + super("sensitive-data suggestion target IDs must be unique"); + this.name = "SensitiveDataSuggestionDuplicateTargetIdsError"; + } +} + +export class SensitiveDataSuggestionNoEligibleColumnsError extends Error { + constructor(readonly scope: SensitiveDataSuggestionScope) { + super("selected scope has no catalog columns"); + this.name = "SensitiveDataSuggestionNoEligibleColumnsError"; + } +} + +export class SensitiveDataSuggestionPayloadTooLargeError extends Error { + constructor() { + super("sensitive-data suggestion structural metadata is too large"); + this.name = "SensitiveDataSuggestionPayloadTooLargeError"; + } +} + export class SensitiveDataSuggestionInvalidResponseError extends Error { constructor() { super("sensitive-data suggestion response is invalid"); @@ -30,6 +76,77 @@ export class SensitiveDataSuggestionInvalidResponseError extends Error { } } +function userContent( + database: { databaseName: string; schema: string }, + columns: readonly StructuralColumn[], +): string { + return JSON.stringify({ + database: database.databaseName, + schema: database.schema, + columns: columns.map((column) => ({ + columnId: column.columnId, + table: column.table, + column: column.column, + dataType: column.dataType, + nullable: column.nullable, + primaryKey: column.primaryKey, + foreignKey: column.foreignKey, + })), + }); +} + +function batchesFor( + database: { databaseName: string; schema: string }, + columns: readonly StructuralColumn[], +): StructuralColumn[][] { + const batches: StructuralColumn[][] = []; + let current: StructuralColumn[] = []; + for (const column of columns) { + if (current.length === MAX_COLUMNS_PER_BATCH) { + batches.push(current); + current = []; + } + const candidate = [...current, column]; + if (Buffer.byteLength(userContent(database, candidate), "utf8") <= MAX_USER_MESSAGE_BYTES) { + current = candidate; + continue; + } + if (current.length === 0) throw new SensitiveDataSuggestionPayloadTooLargeError(); + batches.push(current); + current = [column]; + if (Buffer.byteLength(userContent(database, current), "utf8") > MAX_USER_MESSAGE_BYTES) { + throw new SensitiveDataSuggestionPayloadTooLargeError(); + } + } + if (current.length > 0) batches.push(current); + return batches; +} + +function structuralColumn(table: CatalogTable, column: CatalogColumn): StructuralColumn { + return { + columnId: column.id, + tableId: table.id, + table: table.name, + column: column.name, + dataType: column.dataType, + nullable: column.isNullable, + primaryKey: column.isPrimaryKey, + foreignKey: column.isForeignKey, + version: column.version, + currentSensitive: column.sensitive, + }; +} + +const systemMessage: ModelCompletionMessage = { + role: "system", + content: [ + "Classify whether each database column is likely to contain sensitive source values.", + "Use only the supplied structural metadata. Return strict JSON with this exact shape:", + '{"suggestions":[{"columnId":"uuid","sensitive":true}]}', + "Return every supplied column exactly once. Do not add explanations or markdown.", + ].join("\n"), +}; + export class SensitiveDataSuggester { constructor( private readonly repository: CatalogRepository, @@ -37,67 +154,86 @@ export class SensitiveDataSuggester { private readonly completer: ModelCompleter, ) {} + private async selectColumns( + databaseId: string, + scope: SensitiveDataSuggestionScope, + targetIds: readonly string[], + ): Promise { + if (new Set(targetIds).size !== targetIds.length) { + throw new SensitiveDataSuggestionDuplicateTargetIdsError(); + } + const tables = await this.repository.listTables(databaseId); + const tableIds = new Set(targetIds); + const selectedTables = scope === "selected_tables" + ? tables.filter((table) => tableIds.has(table.id)) + : tables; + if (scope === "selected_tables" && selectedTables.length !== targetIds.length) { + throw new SensitiveDataSuggestionTargetNotFoundError("table"); + } + + const columns = (await Promise.all(selectedTables.map(async (table) => ( + (await this.repository.listColumns(databaseId, table.id)).map((column) => ( + structuralColumn(table, column) + )) + )))).flat(); + const columnIds = new Set(targetIds); + const selectedColumns = scope === "selected_columns" + ? columns.filter((column) => columnIds.has(column.columnId)) + : columns; + if (scope === "selected_columns" && selectedColumns.length !== targetIds.length) { + throw new SensitiveDataSuggestionTargetNotFoundError("column"); + } + if (selectedColumns.length === 0) { + throw new SensitiveDataSuggestionNoEligibleColumnsError(scope); + } + return selectedColumns; + } + async suggest( databaseId: string, modelId: string, + scope: SensitiveDataSuggestionScope, + targetIds: readonly string[], signal: AbortSignal, ): Promise { const database = await this.repository.get(databaseId); - if (!database) throw new SensitiveDataSuggestionTargetNotFoundError(); + if (!database) throw new SensitiveDataSuggestionTargetNotFoundError("database"); + const columns = await this.selectColumns(databaseId, scope, targetIds); + const model = this.models.resolve(modelId); + const suggestions: SensitiveDataSuggestion[] = []; - const tables = await this.repository.listTables(databaseId); - const columns = (await Promise.all(tables.map(async (table) => ({ - table, - columns: await this.repository.listColumns(databaseId, table.id), - })))).flatMap(({ table, columns: tableColumns }) => tableColumns.map((column) => ({ - columnId: column.id, - table: table.name, - column: column.name, - dataType: column.dataType, - nullable: column.isNullable, - primaryKey: column.isPrimaryKey, - foreignKey: column.isForeignKey, - }))); - if (columns.length === 0) return []; - if (columns.length > MAX_COLUMNS) throw new SensitiveDataSuggestionInvalidResponseError(); - - const content = await this.completer.complete({ - model: this.models.resolve(modelId), - signal, - messages: [ - { - role: "system", - content: [ - "Classify whether each database column is likely to contain sensitive source values.", - "Use only the supplied structural metadata. Return strict JSON with this exact shape:", - '{"suggestions":[{"columnId":"uuid","sensitive":true}]}', - "Return every supplied column exactly once. Do not add explanations or markdown.", - ].join("\n"), - }, - { - role: "user", - content: JSON.stringify({ - database: database.databaseName, - schema: database.schema, - columns, - }), - }, - ], - }); - - try { - const parsed = responseSchema.parse(JSON.parse(content)); - const expected = new Set(columns.map((column) => column.columnId)); - const received = new Set(parsed.suggestions.map((suggestion) => suggestion.columnId)); - if (received.size !== parsed.suggestions.length - || received.size !== expected.size - || [...received].some((columnId) => !expected.has(columnId))) { - throw new SensitiveDataSuggestionInvalidResponseError(); + for (const batch of batchesFor(database, columns)) { + let received: Map | undefined; + for (let attempt = 0; attempt < 2 && !received; attempt += 1) { + const content = await this.completer.complete({ + model, + signal, + messages: [systemMessage, { role: "user", content: userContent(database, batch) }], + }); + try { + const parsed = responseSchema.parse(JSON.parse(content)); + const expected = new Set(batch.map((column) => column.columnId)); + const candidate = new Map(parsed.suggestions.map((suggestion) => [suggestion.columnId, suggestion])); + if (candidate.size !== parsed.suggestions.length + || candidate.size !== expected.size + || [...candidate.keys()].some((columnId) => !expected.has(columnId))) { + throw new SensitiveDataSuggestionInvalidResponseError(); + } + received = candidate; + } catch { + if (attempt === 1) throw new SensitiveDataSuggestionInvalidResponseError(); + } } - return parsed.suggestions; - } catch (error) { - if (error instanceof SensitiveDataSuggestionInvalidResponseError) throw error; - throw new SensitiveDataSuggestionInvalidResponseError(); + suggestions.push(...batch.map((column) => ({ + columnId: column.columnId, + tableId: column.tableId, + tableName: column.table, + columnName: column.column, + version: column.version, + currentSensitive: column.currentSensitive, + sensitive: received!.get(column.columnId)!.sensitive, + }))); } + return suggestions; } } diff --git a/backend/src/routes/catalog-description-generation.ts b/backend/src/routes/catalog-description-generation.ts index bfbaa0f3..4d3732a2 100644 --- a/backend/src/routes/catalog-description-generation.ts +++ b/backend/src/routes/catalog-description-generation.ts @@ -14,7 +14,10 @@ import { MetadataGenerationModelUnavailableError } from "../catalog/metadata-gen import { ModelCompletionProviderError } from "../catalog/model-completer.js"; import { SensitiveDataSuggester, + SensitiveDataSuggestionDuplicateTargetIdsError, SensitiveDataSuggestionInvalidResponseError, + SensitiveDataSuggestionNoEligibleColumnsError, + SensitiveDataSuggestionPayloadTooLargeError, SensitiveDataSuggestionTargetNotFoundError, } from "../catalog/sensitive-data-suggester.js"; import { @@ -28,8 +31,20 @@ import { const idSchema = z.uuid(); const modelIdSchema = z.string().regex(/^[a-z][a-z0-9._-]{0,63}$/); -const suggestionSchema = z.object({ modelId: modelIdSchema }).strict(); const selectedTargetIdsSchema = z.array(idSchema).min(1); +const suggestionSchema = z.discriminatedUnion("scope", [ + z.object({ modelId: modelIdSchema, scope: z.literal("all") }).strict(), + z.object({ + modelId: modelIdSchema, + scope: z.literal("selected_tables"), + targetIds: selectedTargetIdsSchema, + }).strict(), + z.object({ + modelId: modelIdSchema, + scope: z.literal("selected_columns"), + targetIds: selectedTargetIdsSchema, + }).strict(), +]); const startSchema = z.discriminatedUnion("scope", [ z.object({ modelId: modelIdSchema, @@ -123,19 +138,6 @@ function safeError(reply: FastifyReply, error: unknown) { message: "The selected metadata-generation model is unavailable.", }); } - if (error instanceof SensitiveDataSuggestionTargetNotFoundError) { - return reply.code(404).send({ - code: "database_not_found", - message: "Database configuration was not found.", - }); - } - if (error instanceof SensitiveDataSuggestionInvalidResponseError - || error instanceof ModelCompletionProviderError) { - return reply.code(502).send({ - code: "sensitive_data_suggestion_failed", - message: "Sensitive-data suggestions could not be prepared.", - }); - } if (error instanceof DescriptionGenerationDuplicateTargetIdsError) { return reply.code(400).send({ code: "description_generation_target_ids_duplicate", @@ -190,6 +192,74 @@ function safeError(reply: FastifyReply, error: unknown) { }); } +function safeSuggestionError(reply: FastifyReply, error: unknown) { + if (error instanceof CatalogUnavailableError) { + return reply.code(503).send({ + code: "catalog_unavailable", + message: "The database catalog is unavailable, so no sensitive-field suggestions were prepared.", + }); + } + if (error instanceof MetadataGenerationModelUnavailableError) { + return reply.code(409).send({ + code: "metadata_generation_model_unavailable", + message: "The selected metadata-generation model is unavailable.", + }); + } + if (error instanceof SensitiveDataSuggestionTargetNotFoundError) { + const code = error.target === "database" + ? "database_not_found" + : error.target === "table" + ? "catalog_table_not_found" + : "catalog_column_not_found"; + const message = error.target === "database" + ? "The database configuration was not found." + : error.target === "table" + ? "One or more selected Catalog Tables were not found in this database." + : "One or more selected Catalog Columns were not found in this database."; + return reply.code(404).send({ code, message }); + } + if (error instanceof SensitiveDataSuggestionDuplicateTargetIdsError) { + return reply.code(400).send({ + code: "sensitive_data_suggestion_target_ids_duplicate", + message: "Each selected table or column must appear only once.", + }); + } + if (error instanceof SensitiveDataSuggestionNoEligibleColumnsError) { + return reply.code(409).send({ + code: "sensitive_data_suggestion_no_columns", + message: "The selected scope contains no Catalog Columns to classify.", + }); + } + if (error instanceof SensitiveDataSuggestionPayloadTooLargeError) { + return reply.code(413).send({ + code: "sensitive_data_suggestion_payload_too_large", + message: "The selected structural metadata cannot be divided into safe LLM requests.", + }); + } + if (error instanceof SensitiveDataSuggestionInvalidResponseError) { + return reply.code(502).send({ + code: "sensitive_data_suggestion_invalid_response", + message: "The LLM returned an incomplete or invalid classification. No suggestions were applied.", + }); + } + if (error instanceof ModelCompletionProviderError) { + return reply.code(502).send({ + code: "sensitive_data_suggestion_provider_unavailable", + message: "The selected LLM service could not complete the request. No suggestions were applied.", + }); + } + if (error instanceof z.ZodError) { + return reply.code(400).send({ + code: "sensitive_data_suggestion_request_invalid", + message: "Choose a database, one or more tables, or one or more columns to classify.", + }); + } + return reply.code(500).send({ + code: "sensitive_data_suggestion_failed", + message: "Sensitive-field suggestions failed before review. No changes were applied.", + }); +} + export function catalogDescriptionGenerationRoutes( app: FastifyInstance, deps: { @@ -206,11 +276,13 @@ export function catalogDescriptionGenerationRoutes( const suggestions = await deps.sensitiveDataSuggester.suggest( databaseId, input.modelId, + input.scope, + "targetIds" in input ? input.targetIds : [], new AbortController().signal, ); return { suggestions }; } catch (error) { - return safeError(reply, error); + return safeSuggestionError(reply, error); } }); diff --git a/backend/src/routes/catalog-schema.ts b/backend/src/routes/catalog-schema.ts index 261d66c9..3c17e8a5 100644 --- a/backend/src/routes/catalog-schema.ts +++ b/backend/src/routes/catalog-schema.ts @@ -16,10 +16,12 @@ import { const idSchema = z.uuid(); const metadataSchema = z.object({ version: z.number().int().positive(), - description: z.string().max(20_000).nullable(), - generatedDescription: z.string().max(20_000).nullable(), + description: z.string().max(20_000).nullable().optional(), + generatedDescription: z.string().max(20_000).nullable().optional(), sensitive: z.boolean().optional(), -}).strict(); +}).strict().refine((value) => ( + "description" in value || "generatedDescription" in value || "sensitive" in value +)); const createRunSchema = z.object({ version: z.number().int().positive(), scope: z.enum(["tables", "columns", "relationships", "all"]), @@ -116,8 +118,10 @@ export function catalogSchemaRoutes( tableId, columnId, input.version, - normalized(input.description), - normalized(input.generatedDescription), + "description" in input ? normalized(input.description ?? null) : current.description, + "generatedDescription" in input + ? normalized(input.generatedDescription ?? null) + : current.generatedDescription, input.sensitive, ); if (!updated) return reply.code(409).send({ code: "column_stale", message: "Column metadata changed. Reload and try again." }); diff --git a/backend/test/catalog-description-generation-routes.test.ts b/backend/test/catalog-description-generation-routes.test.ts index dbb47893..b10a18bb 100644 --- a/backend/test/catalog-description-generation-routes.test.ts +++ b/backend/test/catalog-description-generation-routes.test.ts @@ -147,7 +147,7 @@ test("suggests sensitive flags from structural metadata without persisting them" suggestions: [{ columnId: expect.any(String), sensitive: true }], })), }; - const { app, repository, database, column } = await setup(modelCompleter); + const { app, repository, database, table, column } = await setup(modelCompleter); modelCompleter.complete.mockResolvedValueOnce(JSON.stringify({ suggestions: [{ columnId: column.id, sensitive: true }], })); @@ -156,12 +156,20 @@ test("suggests sensitive flags from structural metadata without persisting them" const response = await app.inject({ method: "POST", url: `/catalog/databases/${database.id}/sensitive-data-suggestions`, - payload: { modelId: configuredModel.id }, + payload: { modelId: configuredModel.id, scope: "all" }, }); expect(response.statusCode).toBe(200); expect(response.json()).toEqual({ - suggestions: [{ columnId: column.id, sensitive: true }], + suggestions: [{ + columnId: column.id, + tableId: table.id, + tableName: table.name, + columnName: column.name, + version: column.version, + currentSensitive: false, + sensitive: true, + }], }); expect(await repository.getColumn(database.id, column.tableId, column.id)) .toMatchObject({ sensitive: false }); @@ -178,6 +186,262 @@ test("suggests sensitive flags from structural metadata without persisting them" } }); +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); + await repository.applySchemaSync(database.id, database.version, "all", [], { + schemaVersion: 1, + capabilities: { tables: "available", columns: "available", relationships: "available" }, + tables: [ + { name: "patients", sourceComment: null }, + { name: "visits", sourceComment: null }, + { name: "billing", sourceComment: null }, + ], + columns: [ + { tableName: "patients", name: "patient_name", ordinalPosition: 1, dataType: "text", isNullable: false, defaultExpression: null, primaryKeyPosition: null, sourceComment: null }, + { tableName: "patients", name: "status", ordinalPosition: 2, dataType: "text", isNullable: true, defaultExpression: null, primaryKeyPosition: null, sourceComment: null }, + { tableName: "visits", name: "clinical_note", ordinalPosition: 1, dataType: "text", isNullable: true, defaultExpression: null, primaryKeyPosition: null, sourceComment: null }, + { tableName: "billing", name: "invoice_total", ordinalPosition: 1, dataType: "numeric", isNullable: false, defaultExpression: null, primaryKeyPosition: null, sourceComment: null }, + ], + relationships: [], + }); + const tables = await repository.listTables(database.id); + const patients = tables.find((table) => table.name === "patients")!; + const visits = tables.find((table) => table.name === "visits")!; + const billing = tables.find((table) => table.name === "billing")!; + const patientColumns = await repository.listColumns(database.id, patients.id); + const visitColumns = await repository.listColumns(database.id, visits.id); + const billingColumns = await repository.listColumns(database.id, billing.id); + const status = patientColumns.find((column) => column.name === "status")!; + const clinicalNote = visitColumns.find((column) => column.name === "clinical_note")!; + + try { + const tableResponse = await app.inject({ + method: "POST", + url: `/catalog/databases/${database.id}/sensitive-data-suggestions`, + payload: { + modelId: configuredModel.id, + scope: "selected_tables", + targetIds: [visits.id, patients.id], + }, + }); + expect(tableResponse.statusCode).toBe(200); + expect(tableResponse.json().suggestions).toHaveLength(3); + expect(tableResponse.json().suggestions).toEqual(expect.arrayContaining([ + expect.objectContaining({ tableId: patients.id, columnName: "patient_name", sensitive: true }), + expect.objectContaining({ tableId: patients.id, columnName: "status", sensitive: false }), + expect.objectContaining({ tableId: visits.id, columnName: "clinical_note", sensitive: true }), + ])); + expect(tableResponse.json().suggestions).not.toEqual(expect.arrayContaining([ + expect.objectContaining({ tableId: billing.id }), + ])); + + const columnResponse = await app.inject({ + method: "POST", + url: `/catalog/databases/${database.id}/sensitive-data-suggestions`, + payload: { + modelId: configuredModel.id, + scope: "selected_columns", + targetIds: [clinicalNote.id, status.id], + }, + }); + expect(columnResponse.statusCode).toBe(200); + expect(columnResponse.json().suggestions).toHaveLength(2); + expect(columnResponse.json().suggestions).toEqual(expect.arrayContaining([ + expect.objectContaining({ tableId: patients.id, columnId: status.id, sensitive: false }), + 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(), + ); + } finally { + await app.close(); + } +}); + +test("explains invalid sensitive-data suggestion selections without calling the model", async () => { + const modelCompleter: ModelCompleter = { complete: vi.fn(async () => "unused") }; + const { app, database, table } = 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: [] }, + }); + 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.", + }); + + 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], + }, + }); + expect(duplicate.statusCode).toBe(400); + expect(duplicate.json()).toEqual({ + code: "sensitive_data_suggestion_target_ids_duplicate", + message: "Each selected table or column must appear only once.", + }); + + const missingTable = await app.inject({ + method: "POST", + url: `/catalog/databases/${database.id}/sensitive-data-suggestions`, + payload: { + modelId: configuredModel.id, + scope: "selected_tables", + targetIds: ["00000000-0000-4000-8000-000000000001"], + }, + }); + expect(missingTable.statusCode).toBe(404); + expect(missingTable.json()).toEqual({ + code: "catalog_table_not_found", + message: "One or more selected Catalog Tables were not found in this database.", + }); + + const missingColumn = await app.inject({ + method: "POST", + url: `/catalog/databases/${database.id}/sensitive-data-suggestions`, + payload: { + modelId: configuredModel.id, + scope: "selected_columns", + targetIds: ["00000000-0000-4000-8000-000000000002"], + }, + }); + expect(missingColumn.statusCode).toBe(404); + expect(missingColumn.json()).toEqual({ + code: "catalog_column_not_found", + 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) => { @@ -201,13 +465,13 @@ test.each(["malformed", "incomplete", "duplicate"] as const)( const response = await app.inject({ method: "POST", url: `/catalog/databases/${database.id}/sensitive-data-suggestions`, - payload: { modelId: configuredModel.id }, + payload: { modelId: configuredModel.id, scope: "all" }, }); expect(response.statusCode).toBe(502); expect(response.json()).toEqual({ - code: "sensitive_data_suggestion_failed", - message: "Sensitive-data suggestions could not be prepared.", + 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)) @@ -218,6 +482,34 @@ test.each(["malformed", "incomplete", "duplicate"] as const)( }, ); +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 }); + } finally { + await app.close(); + } +}); + interface SseFrame { id?: string; event?: string; diff --git a/backend/test/catalog-schema-routes.test.ts b/backend/test/catalog-schema-routes.test.ts index c4cb3103..d7a54aaf 100644 --- a/backend/test/catalog-schema-routes.test.ts +++ b/backend/test/catalog-schema-routes.test.ts @@ -241,14 +241,29 @@ test("keeps generated descriptions editable and preserves them across synchroniz version: idColumn.version, description: "Reviewed key", generatedDescription: "Generated key draft", - sensitive: true, + sensitive: false, }, }); expect(editedColumn.json()).toMatchObject({ + description: "Reviewed key", + generatedDescription: "Generated key draft", + sensitive: false, + }); + const sensitiveOnly = await app.inject({ + method: "PATCH", url: `/catalog/databases/${database.id}/tables/${patients.id}/columns/${idColumn.id}`, + payload: { version: editedColumn.json().version, sensitive: true }, + }); + expect(sensitiveOnly.statusCode).toBe(200); + expect(sensitiveOnly.json()).toMatchObject({ description: "Reviewed key", generatedDescription: "Generated key draft", sensitive: true, }); + const emptyPatch = await app.inject({ + method: "PATCH", url: `/catalog/databases/${database.id}/tables/${patients.id}/columns/${idColumn.id}`, + payload: { version: sensitiveOnly.json().version }, + }); + expect(emptyPatch.statusCode).toBe(400); const second = await app.inject({ method: "POST", url: `/catalog/databases/${database.id}/sync-runs`, payload: { version: database.version, scope: "all", tableIds: [] } }); await waitFor(repository, second.json().id, "succeeded"); diff --git a/docs/adr/0011-gate-source-samples-with-a-sensitive-data-flag.md b/docs/adr/0011-gate-source-samples-with-a-sensitive-data-flag.md index 818719e8..9a8563cb 100644 --- a/docs/adr/0011-gate-source-samples-with-a-sensitive-data-flag.md +++ b/docs/adr/0011-gate-source-samples-with-a-sensitive-data-flag.md @@ -9,6 +9,13 @@ draft suggestions from structural metadata only, but the user decides and only t persisted; there is no rationale, history, audit ledger, fingerprint, review state, or retroactive regeneration of existing descriptions. +The user starts a suggestion from an explicit selection at database, table, or column level. A +database request accepts exactly one selected database; a table or column request contains only the +selected tables' columns or the selected columns, respectively. The backend divides that structural +metadata into deterministic model requests of at most ten columns, also bounded by helper message +size, and returns one combined draft for human review. No flag changes until the user saves the +reviewed draft. + Description generation extends ADR-0010 by sending bounded real values when `sensitive` is false and deterministic plausible synthetic values when it is true, without identifying the synthetic values to the model. The false default deliberately favors the expected stable schemas and the diff --git a/frontend/src/api/catalog-databases.ts b/frontend/src/api/catalog-databases.ts index 68e6f893..b792550b 100644 --- a/frontend/src/api/catalog-databases.ts +++ b/frontend/src/api/catalog-databases.ts @@ -99,6 +99,11 @@ export interface CatalogColumn { export interface SensitiveDataSuggestion { columnId: string; + tableId: string; + tableName: string; + columnName: string; + version: number; + currentSensitive: boolean; sensitive: boolean; } @@ -106,6 +111,11 @@ export interface SensitiveDataSuggestions { suggestions: SensitiveDataSuggestion[]; } +export type SensitiveDataSuggestionRequest = + | { scope: "all" } + | { scope: "selected_tables"; targetIds: string[] } + | { scope: "selected_columns"; targetIds: string[] }; + export interface CatalogRelationshipColumn { position: number; sourceColumnId: string; @@ -357,10 +367,25 @@ export const updateCatalogColumnMetadata = ( { method: "PATCH", body: JSON.stringify({ version, description, generatedDescription, sensitive }) }, ); -export const suggestSensitiveFields = (databaseId: string, modelId: string) => +export const updateCatalogColumnSensitive = ( + databaseId: string, + tableId: string, + columnId: string, + version: number, + sensitive: boolean, +) => apiFetch( + `/catalog/databases/${encodeURIComponent(databaseId)}/tables/${encodeURIComponent(tableId)}/columns/${encodeURIComponent(columnId)}`, + { method: "PATCH", body: JSON.stringify({ version, sensitive }) }, +); + +export const suggestSensitiveFields = ( + databaseId: string, + modelId: string, + selection: SensitiveDataSuggestionRequest, +) => apiFetch( `/catalog/databases/${encodeURIComponent(databaseId)}/sensitive-data-suggestions`, - { method: "POST", body: JSON.stringify({ modelId }) }, + { method: "POST", body: JSON.stringify({ modelId, ...selection }) }, ); export const listCatalogRelationships = (databaseId: string) => diff --git a/frontend/src/api/client.test.ts b/frontend/src/api/client.test.ts index 0bafbb1a..301e529b 100644 --- a/frontend/src/api/client.test.ts +++ b/frontend/src/api/client.test.ts @@ -147,7 +147,9 @@ test.each([ ["description_generation_target_ids_duplicate", "Description generation target IDs must be unique."], ["description_generation_no_eligible_targets", "No eligible catalog tables or columns need description generation."], ["catalog_table_not_found", "One or more selected catalog tables were not found."], -])("maps the description-generation error code %s to safe local copy", async (code, message) => { + ["sensitive_data_suggestion_invalid_response", "The model returned an incomplete or invalid classification. No suggestions were applied."], + ["sensitive_data_suggestion_provider_unavailable", "The selected model could not complete the request. No suggestions were applied."], +])("maps the catalog error code %s to safe local copy", async (code, message) => { const fetchSpy = vi.spyOn(globalThis, "fetch").mockResolvedValue( new Response(JSON.stringify({ code, message: "provider detail must not be trusted" }), { status: 400, diff --git a/frontend/src/api/client.ts b/frontend/src/api/client.ts index 5dd93bfc..d9c6b04e 100644 --- a/frontend/src/api/client.ts +++ b/frontend/src/api/client.ts @@ -23,6 +23,13 @@ const safeErrorCodes = new Set([ "description_generation_target_ids_duplicate", "metadata_generation_model_unavailable", "catalog_column_not_found", "catalog_table_not_found", "workspace_configuration_unavailable", + "sensitive_data_suggestion_request_invalid", + "sensitive_data_suggestion_target_ids_duplicate", + "sensitive_data_suggestion_no_columns", + "sensitive_data_suggestion_payload_too_large", + "sensitive_data_suggestion_invalid_response", + "sensitive_data_suggestion_provider_unavailable", + "sensitive_data_suggestion_failed", "schema_sync_conflict", "schema_introspection_failed", "schema_request_invalid", "schema_operation_failed", "sync_run_not_found", "table_stale", "column_stale", ]); @@ -77,6 +84,13 @@ const localCodeMessages: Record = { catalog_column_not_found: "The selected catalog column was not found.", catalog_table_not_found: "One or more selected catalog tables were not found.", workspace_configuration_unavailable: "The database workspace configuration is unavailable.", + sensitive_data_suggestion_request_invalid: "Select a database, one or more tables, or one or more columns before requesting sensitive-field suggestions.", + sensitive_data_suggestion_target_ids_duplicate: "Each selected table or column can be included only once.", + sensitive_data_suggestion_no_columns: "The selected scope contains no catalog columns to classify.", + sensitive_data_suggestion_payload_too_large: "The selected structural metadata cannot be divided into safe model requests.", + sensitive_data_suggestion_invalid_response: "The model returned an incomplete or invalid classification. No suggestions were applied.", + sensitive_data_suggestion_provider_unavailable: "The selected model could not complete the request. No suggestions were applied.", + sensitive_data_suggestion_failed: "Sensitive-field suggestions failed before review. No changes were applied.", schema_sync_conflict: "A schema synchronization is already active or no longer current.", schema_introspection_failed: "The database schema could not be read safely.", schema_request_invalid: "The schema request is invalid.", diff --git a/frontend/src/shell/DatabaseManagementPage.test.tsx b/frontend/src/shell/DatabaseManagementPage.test.tsx index 68165082..750b5955 100644 --- a/frontend/src/shell/DatabaseManagementPage.test.tsx +++ b/frontend/src/shell/DatabaseManagementPage.test.tsx @@ -1571,6 +1571,125 @@ test("starts one selected column with the configured default model", async () => expect(await screen.findByText("Description generation started for 1 column")).toBeVisible(); }); +test("shows database sensitive suggestions only for a selection and rejects multiple databases clearly", async () => { + const user = userEvent.setup(); + let suggestionCalls = 0; + const radiology = makeDatabase({ + id: "44444444-4444-4444-8444-444444444444", + workspaceId: "radiology", + workspaceName: "Radiology", + }); + server.use( + http.get("/api/catalog/metadata-generation/models", () => HttpResponse.json({ + models: [{ id: "local-qwen", label: "Local Qwen" }], + default: "local-qwen", + })), + http.post("/api/catalog/databases/:databaseId/sensitive-data-suggestions", () => { + suggestionCalls += 1; + return HttpResponse.json({ suggestions: [] }); + }), + ); + renderPage({ rows: [makeDatabase(), radiology] }); + + expect(screen.queryByRole("button", { name: "Suggest sensitive fields" })).not.toBeInTheDocument(); + const psdRow = await screen.findByRole("row", { name: /Policlinico San Donato/ }); + const radiologyRow = await screen.findByRole("row", { name: /Radiology/ }); + await user.click(within(psdRow).getByRole("checkbox", { name: /toggle row selection/i })); + expect(screen.getByRole("button", { name: "Suggest sensitive fields" })).toBeVisible(); + await user.click(within(radiologyRow).getByRole("checkbox", { name: /toggle row selection/i })); + await user.click(screen.getByRole("button", { name: "Suggest sensitive fields" })); + + expect(await screen.findByText("Sensitive-field suggestions can be requested for only one database at a time. Select one database and try again.")).toBeVisible(); + expect(suggestionCalls).toBe(0); +}); + +test("requests database-level sensitive suggestions for the only selected database", async () => { + const user = userEvent.setup(); + let suggestionBody: unknown; + server.use( + http.get("/api/catalog/metadata-generation/models", () => HttpResponse.json({ + models: [{ id: "local-qwen", label: "Local Qwen" }], + default: "local-qwen", + })), + http.post("/api/catalog/databases/:databaseId/sensitive-data-suggestions", async ({ request }) => { + suggestionBody = await request.json(); + return HttpResponse.json({ + suggestions: [{ + columnId: patientIdColumn.id, + tableId: patientsTable.id, + tableName: patientsTable.name, + columnName: patientIdColumn.name, + version: patientIdColumn.version, + currentSensitive: false, + sensitive: true, + }], + }); + }), + ); + renderPage({ rows: [makeDatabase()] }); + + const databaseRow = await screen.findByRole("row", { name: /Policlinico San Donato/ }); + await user.click(within(databaseRow).getByRole("checkbox", { name: /toggle row selection/i })); + await user.click(screen.getByRole("button", { name: "Suggest sensitive fields" })); + + await waitFor(() => expect(suggestionBody).toEqual({ modelId: "local-qwen", scope: "all" })); + expect(await screen.findByRole("complementary", { name: "Sensitive field review" })).toBeVisible(); +}); + +test("requests sensitive suggestions only for selected tables", async () => { + const user = userEvent.setup(); + const visitsTable: CatalogTable = { + ...patientsTable, + id: "bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb", + name: "visits", + }; + const visitColumn = { + ...patientIdColumn, + id: "eeeeeeee-eeee-4eee-8eee-eeeeeeeeeeee", + tableId: visitsTable.id, + name: "clinical_note", + }; + let suggestionBody: unknown; + server.use( + http.get("/api/catalog/metadata-generation/models", () => HttpResponse.json({ + models: [{ id: "local-qwen", label: "Local Qwen" }], + default: "local-qwen", + })), + http.get("/api/catalog/databases/:databaseId/tables", () => HttpResponse.json([ + patientsTable, + visitsTable, + ])), + http.post("/api/catalog/databases/:databaseId/sensitive-data-suggestions", async ({ request }) => { + suggestionBody = await request.json(); + return HttpResponse.json({ + suggestions: [{ + columnId: visitColumn.id, + tableId: visitsTable.id, + tableName: visitsTable.name, + columnName: visitColumn.name, + version: visitColumn.version, + currentSensitive: false, + sensitive: true, + }], + }); + }), + ); + renderPage({ rows: [makeDatabase({ connectionStatus: "reachable", testedVersion: 3 })] }); + + await user.click(await screen.findByRole("button", { name: "View Policlinico San Donato" })); + await user.click(screen.getByRole("tab", { name: "Tables" })); + const visitsRow = await screen.findByRole("row", { name: /visits/ }); + await user.click(within(visitsRow).getByRole("checkbox", { name: /toggle row selection/i })); + await user.click(screen.getByRole("button", { name: "Suggest sensitive fields" })); + + await waitFor(() => expect(suggestionBody).toEqual({ + modelId: "local-qwen", + scope: "selected_tables", + targetIds: [visitsTable.id], + })); + expect(await screen.findByRole("complementary", { name: "Sensitive field review" })).toBeVisible(); +}); + test("reviews AI-sensitive-field suggestions as an editable draft and saves only changed columns", async () => { const user = userEvent.setup(); const idColumn = { ...patientIdColumn, sensitive: false }; @@ -1585,7 +1704,18 @@ test("reviews AI-sensitive-field suggestions as an editable draft and saves only generatedDescription: "Name of the patient", sensitive: false, }; - let columns = [idColumn, nameColumn]; + const unselectedColumn = { + ...patientIdColumn, + id: "ffffffff-ffff-4fff-8fff-ffffffffffff", + name: "address", + ordinalPosition: 3, + primaryKeyPosition: null, + isPrimaryKey: false, + description: "Patient address", + generatedDescription: "Address of the patient", + sensitive: false, + }; + let columns = [idColumn, nameColumn, unselectedColumn]; let suggestionBody: unknown; const patches: Array<{ columnId: string; body: unknown }> = []; server.use( @@ -1604,9 +1734,24 @@ test("reviews AI-sensitive-field suggestions as an editable draft and saves only suggestionBody = await request.json(); return HttpResponse.json({ suggestions: [ - { columnId: idColumn.id, sensitive: true }, - { columnId: nameColumn.id, sensitive: true }, - { columnId: "ffffffff-ffff-4fff-8fff-ffffffffffff", sensitive: true }, + { + columnId: idColumn.id, + tableId: patientsTable.id, + tableName: patientsTable.name, + columnName: idColumn.name, + version: idColumn.version, + currentSensitive: false, + sensitive: true, + }, + { + columnId: nameColumn.id, + tableId: patientsTable.id, + tableName: patientsTable.name, + columnName: nameColumn.name, + version: nameColumn.version, + currentSensitive: false, + sensitive: true, + }, ], }); }, @@ -1616,8 +1761,6 @@ test("reviews AI-sensitive-field suggestions as an editable draft and saves only async ({ params, request }) => { const body = await request.json() as { version: number; - description: string | null; - generatedDescription: string | null; sensitive: boolean; }; patches.push({ columnId: String(params.columnId), body }); @@ -1638,31 +1781,41 @@ test("reviews AI-sensitive-field suggestions as an editable draft and saves only const nameSensitive = await screen.findByRole("checkbox", { name: "Sensitive data for name" }); expect(idSensitive).not.toBeChecked(); expect(nameSensitive).not.toBeChecked(); + expect(screen.queryByRole("button", { name: "Suggest sensitive fields" })).not.toBeInTheDocument(); + const selectableRow = async (name: RegExp) => { + const rows = await screen.findAllByRole("row", { name }); + return rows.find((row) => within(row).queryByRole("checkbox", { name: /toggle row selection/i }))!; + }; + await user.click(within(await selectableRow(/Patient identifier/)) + .getByRole("checkbox", { name: /toggle row selection/i })); + await user.click(within(await selectableRow(/Patient name/)) + .getByRole("checkbox", { name: /toggle row selection/i })); await user.click(screen.getByRole("button", { name: "Suggest sensitive fields" })); - await waitFor(() => expect(suggestionBody).toEqual({ modelId: "local-qwen" })); - await waitFor(() => { - expect(screen.getByRole("checkbox", { name: "Sensitive data for id" })).toBeChecked(); - expect(screen.getByRole("checkbox", { name: "Sensitive data for name" })).toBeChecked(); - }); + await waitFor(() => expect(suggestionBody).toEqual({ + modelId: "local-qwen", + scope: "selected_columns", + targetIds: [idColumn.id, nameColumn.id], + })); + const review = await screen.findByRole("complementary", { name: "Sensitive field review" }); + expect(within(review).getByRole("checkbox", { name: "Protect patients.id" })).toBeChecked(); + expect(within(review).getByRole("checkbox", { name: "Protect patients.name" })).toBeChecked(); expect(patches).toHaveLength(0); - await user.click(screen.getByRole("checkbox", { name: "Sensitive data for name" })); - expect(screen.getByRole("checkbox", { name: "Sensitive data for name" })).not.toBeChecked(); - await user.click(screen.getByRole("button", { name: "Save sensitive fields" })); + await user.click(within(review).getByRole("checkbox", { name: "Protect patients.name" })); + await user.click(within(review).getByRole("button", { name: "Save 1" })); await waitFor(() => expect(patches).toEqual([{ columnId: idColumn.id, body: { version: idColumn.version, - description: idColumn.description, - generatedDescription: idColumn.generatedDescription, sensitive: true, }, }])); - expect(await screen.findByText("Sensitive fields saved")).toBeVisible(); - expect(screen.getByRole("checkbox", { name: "Sensitive data for id" })).toBeChecked(); + expect(await screen.findByText("Saved 1 sensitive flag")).toBeVisible(); + await waitFor(() => expect(screen.queryByRole("complementary", { name: "Sensitive field review" })).not.toBeInTheDocument()); + await waitFor(() => expect(screen.getByRole("checkbox", { name: "Sensitive data for id" })).toBeChecked()); expect(screen.getByRole("checkbox", { name: "Sensitive data for name" })).not.toBeChecked(); }); diff --git a/frontend/src/shell/DatabaseManagementPage.tsx b/frontend/src/shell/DatabaseManagementPage.tsx index b9fc2909..9bc32ece 100644 --- a/frontend/src/shell/DatabaseManagementPage.tsx +++ b/frontend/src/shell/DatabaseManagementPage.tsx @@ -20,9 +20,11 @@ import { replaceCatalogDatabaseSecrets, startCatalogSync, startDescriptionGenerationRun, + suggestSensitiveFields, testCatalogDatabase, updateCatalogDatabase, type CatalogDatabase, + type CatalogColumn, type CatalogDatabaseMetadataDeleteTarget, type CatalogSecretName, type CatalogSyncScope, @@ -30,6 +32,8 @@ import { type DatabaseBinding, type DatabaseTransport, type DescriptionGenerationRun, + type SensitiveDataSuggestion, + type SensitiveDataSuggestionRequest, } from "../api/catalog-databases"; import { DatabaseGrid } from "./database-management/DatabaseGrid"; import { DatabaseForm } from "./database-management/DatabaseForm"; @@ -38,6 +42,7 @@ import { DatabaseRelationships } from "./database-management/DatabaseRelationshi import { CatalogSyncDrawer } from "./database-management/CatalogSyncDrawer"; import { MetadataGenerationModelSelector } from "./database-management/MetadataGenerationModelSelector"; import { DescriptionGenerationDrawer } from "./database-management/DescriptionGenerationDrawer"; +import { SensitiveDataReviewDrawer } from "./database-management/SensitiveDataReviewDrawer"; import { configurationFingerprint, configurationFromDraft, @@ -141,6 +146,11 @@ export function DatabaseManagementPage({ const [syncDrawerOpen, setSyncDrawerOpen] = useState(false); const [activeDescriptionGenerationRun, setActiveDescriptionGenerationRun] = useState(null); const [descriptionGenerationDrawerOpen, setDescriptionGenerationDrawerOpen] = useState(false); + const [sensitiveReview, setSensitiveReview] = useState<{ + databaseId: string; + scopeLabel: string; + suggestions: SensitiveDataSuggestion[]; + } | null>(null); const originRef = useRef(null); const searchInputRef = useRef(null); @@ -683,6 +693,67 @@ export function DatabaseManagementPage({ } }, [rememberDescriptionGenerationRun, selectedMetadataModel]); + const requestSensitiveSuggestions = useCallback(async ( + database: CatalogDatabase, + selection: SensitiveDataSuggestionRequest, + scopeLabel: string, + ) => { + if (!database.id) { + toast.error("The selected database is not configured, so sensitive-field suggestions were not requested."); + throw new Error("database is not configured"); + } + if (!selectedMetadataModel) { + toast.error("Select a metadata-generation model before requesting sensitive-field suggestions."); + throw new Error("metadata-generation model is not selected"); + } + try { + const result = await suggestSensitiveFields(database.id, selectedMetadataModel, selection); + setSensitiveReview({ databaseId: database.id, scopeLabel, suggestions: result.suggestions }); + toast.success(`Prepared ${result.suggestions.length} sensitive-field suggestion${result.suggestions.length === 1 ? "" : "s"} for review`); + } catch (error) { + toast.error(apiErrorMessage(error)); + throw error; + } + }, [selectedMetadataModel]); + + const suggestDatabaseSensitiveFields = useCallback(async (selected: CatalogDatabase[]) => { + if (selected.length !== 1) { + toast.error("Sensitive-field suggestions can be requested for only one database at a time. Select one database and try again."); + throw new Error("more than one database selected"); + } + const database = selected[0]!; + await requestSensitiveSuggestions(database, { scope: "all" }, `Entire database ${database.workspaceName}`); + }, [requestSensitiveSuggestions]); + + const suggestActiveDatabaseSensitiveFields = useCallback(async ( + selection: SensitiveDataSuggestionRequest, + scopeLabel: string, + ) => { + if (!activeRow) { + toast.error("The database is no longer available, so sensitive-field suggestions were not requested."); + throw new Error("database is no longer available"); + } + await requestSensitiveSuggestions(activeRow, selection, scopeLabel); + }, [activeRow, requestSensitiveSuggestions]); + + const sensitiveColumnsSaved = useCallback((columns: CatalogColumn[]) => { + const savedById = new Map(columns.map((column) => [column.id, column])); + setSensitiveReview((current) => current ? { + ...current, + suggestions: current.suggestions.map((suggestion) => { + const saved = savedById.get(suggestion.columnId); + return saved ? { + ...suggestion, + version: saved.version, + currentSensitive: saved.sensitive, + sensitive: saved.sensitive, + } : suggestion; + }), + } : null); + const databaseId = sensitiveReview?.databaseId; + if (databaseId) void queryClient.invalidateQueries({ queryKey: ["catalog-columns", databaseId] }); + }, [queryClient, sensitiveReview?.databaseId]); + const deleteSelectedMetadata = useCallback(async ( selected: CatalogDatabase[], target: CatalogDatabaseMetadataDeleteTarget, @@ -832,6 +903,7 @@ export function DatabaseManagementPage({ selectedMetadataModel={selectedMetadataModelAvailable ? selectedMetadataModel : null} descriptionGenerationActive={descriptionGenerationActive} onGenerateDescriptions={generateDatabaseDescriptions} + onSuggestSensitive={suggestDatabaseSensitiveFields} onDeleteMetadataSelected={deleteSelectedMetadata} /> )} @@ -886,6 +958,7 @@ export function DatabaseManagementPage({ onRunStarted={rememberSyncRun} onOpenSync={() => openSync(activeRow)} onDescriptionGenerationRunStarted={rememberDescriptionGenerationRun} + onSuggestSensitive={suggestActiveDatabaseSensitiveFields} /> ) : null} @@ -920,6 +993,15 @@ export function DatabaseManagementPage({ onRunUpdate={setActiveDescriptionGenerationRun} onTerminal={(run) => void descriptionGenerationTerminated(run)} /> + setSensitiveReview(null)} + onSaved={sensitiveColumnsSaved} + /> ); } diff --git a/frontend/src/shell/database-management/DatabaseColumns.tsx b/frontend/src/shell/database-management/DatabaseColumns.tsx index 6d335e9d..dbac8bfd 100644 --- a/frontend/src/shell/database-management/DatabaseColumns.tsx +++ b/frontend/src/shell/database-management/DatabaseColumns.tsx @@ -11,11 +11,11 @@ import { consolidateCatalogDescriptions, listCatalogColumns, startDescriptionGenerationRun, - suggestSensitiveFields, updateCatalogColumnMetadata, type CatalogColumn, type CatalogTable, type DescriptionGenerationRun, + type SensitiveDataSuggestionRequest, } from "../../api/catalog-databases"; import type { DatabaseNavigationState } from "./model"; @@ -28,6 +28,7 @@ interface Props { onDescriptionGenerationRunStarted: (run: DescriptionGenerationRun) => void; onNavigationStateChange: (state: DatabaseNavigationState) => void; onSync: () => void; + onSuggestSensitive: (selection: SensitiveDataSuggestionRequest, scopeLabel: string) => Promise; } interface GridContext { @@ -84,6 +85,7 @@ export function DatabaseColumns({ onDescriptionGenerationRunStarted, onNavigationStateChange, onSync, + onSuggestSensitive, }: Props) { const queryClient = useQueryClient(); const queryKey = ["catalog-columns", databaseId, table.id] as const; @@ -222,23 +224,16 @@ export function DatabaseColumns({ }; const suggestSensitive = async () => { - if (!selectedMetadataModel) return; + if (selectedIds.length === 0) return; setBusy(true); setSensitiveAction("suggest"); try { - const result = await suggestSensitiveFields(databaseId, selectedMetadataModel); - const currentById = new Map(data.map((column) => [column.id, column])); - const next: Record = {}; - for (const suggestion of result.suggestions) { - const column = currentById.get(suggestion.columnId); - if (column && suggestion.sensitive !== column.sensitive) { - next[column.id] = suggestion.sensitive; - } - } - setSensitiveDrafts(next); - toast.success("Sensitive field suggestions ready for review"); - } catch (error) { - toast.error(apiErrorMessage(error)); + await onSuggestSensitive( + { scope: "selected_columns", targetIds: selectedIds }, + `${selectedIds.length} selected column${selectedIds.length === 1 ? "" : "s"} in ${table.name}`, + ); + } catch { + // The page-level request reports the safe, specific reason and preserves the selection. } finally { setSensitiveAction(null); setBusy(false); @@ -351,6 +346,20 @@ export function DatabaseColumns({ + + {changedSensitiveColumns.length > 0 ? ( + + ) : null} ) : ( @@ -358,14 +367,6 @@ export function DatabaseColumns({ Catalog columns setSearch(event.target.value)} /> {data.length} - {changedSensitiveColumns.length > 0 ? ( ) diff --git a/frontend/src/shell/database-management/DatabaseTables.tsx b/frontend/src/shell/database-management/DatabaseTables.tsx index 12d9c80a..35589dc0 100644 --- a/frontend/src/shell/database-management/DatabaseTables.tsx +++ b/frontend/src/shell/database-management/DatabaseTables.tsx @@ -3,7 +3,7 @@ import { Menu } from "@base-ui/react/menu"; import { useQuery, useQueryClient } from "@tanstack/react-query"; import { AgGridReact } from "ag-grid-react"; import type { ColDef, ICellRendererParams } from "ag-grid-community"; -import { ArrowLeft, ChevronDown, Columns3, Pencil, RefreshCw, Save, Trash2, X } from "lucide-react"; +import { ArrowLeft, ChevronDown, Columns3, Pencil, RefreshCw, Save, Sparkles, Trash2, X } from "lucide-react"; import { toast } from "sonner"; import { Button } from "../../components/ui/button"; import { ApiError, apiErrorMessage } from "../../api/client"; @@ -19,6 +19,7 @@ import { type CatalogTable, type CatalogTableMetadataDeleteTarget, type DescriptionGenerationRun, + type SensitiveDataSuggestionRequest, } from "../../api/catalog-databases"; import type { DatabaseNavigationState } from "./model"; import { DatabaseColumns } from "./DatabaseColumns"; @@ -36,6 +37,7 @@ interface Props { onRunStarted: (run: CatalogSyncRun) => void; onOpenSync: () => void; onDescriptionGenerationRunStarted: (run: DescriptionGenerationRun) => void; + onSuggestSensitive: (selection: SensitiveDataSuggestionRequest, scopeLabel: string) => Promise; } interface TableGridContext { @@ -70,6 +72,7 @@ export function DatabaseTables({ onRunStarted, onOpenSync, onDescriptionGenerationRunStarted, + onSuggestSensitive, }: Props) { const databaseId = database.id!; const queryClient = useQueryClient(); @@ -86,7 +89,7 @@ export function DatabaseTables({ const [editorVersion, setEditorVersion] = useState(null); const [stale, setStale] = useState(false); const [staleBannerOpen, setStaleBannerOpen] = useState(true); - const [busy, setBusy] = useState<"sync" | "save" | "delete" | "consolidate" | "generate" | null>(null); + const [busy, setBusy] = useState<"sync" | "save" | "delete" | "consolidate" | "generate" | "suggest" | null>(null); const [pendingDelete, setPendingDelete] = useState(null); const [columnNavigation, setColumnNavigation] = useState({ dirty: false, busy: false }); const gridRef = useRef>(null); @@ -240,6 +243,20 @@ export function DatabaseTables({ toast.error(apiErrorMessage(error)); } finally { setBusy(null); } }; + const suggestSensitive = async () => { + if (selectedIds.length === 0) return; + setBusy("suggest"); + try { + await onSuggestSensitive( + { scope: "selected_tables", targetIds: selectedIds }, + `${selectedIds.length} selected table${selectedIds.length === 1 ? "" : "s"}`, + ); + } catch { + // The page-level request reports the safe, specific reason and preserves the selection. + } finally { + setBusy(null); + } + }; const columns = useMemo[]>(() => [ { field: "name", headerName: "Name", minWidth: 250, flex: 1, cellClass: "font-mono text-xs" }, @@ -303,6 +320,7 @@ export function DatabaseTables({ onDescriptionGenerationRunStarted={onDescriptionGenerationRunStarted} onNavigationStateChange={setColumnNavigation} onSync={() => void synchronize("columns", [activeTable.id])} + onSuggestSensitive={onSuggestSensitive} /> ) : (
@@ -420,6 +438,15 @@ export function DatabaseTables({ + ) diff --git a/frontend/src/shell/database-management/SensitiveDataReviewDrawer.tsx b/frontend/src/shell/database-management/SensitiveDataReviewDrawer.tsx new file mode 100644 index 00000000..9ce964fb --- /dev/null +++ b/frontend/src/shell/database-management/SensitiveDataReviewDrawer.tsx @@ -0,0 +1,188 @@ +import { useEffect, useMemo, useState } from "react"; +import { Save, X } from "lucide-react"; +import { toast } from "sonner"; +import { Button } from "../../components/ui/button"; +import { apiErrorMessage } from "../../api/client"; +import { + updateCatalogColumnSensitive, + type CatalogColumn, + type SensitiveDataSuggestion, +} from "../../api/catalog-databases"; + +interface Props { + open: boolean; + databaseId: string | null; + scopeLabel: string; + suggestions: SensitiveDataSuggestion[]; + canManage: boolean; + onClose: () => void; + onSaved: (columns: CatalogColumn[]) => void; +} + +export function SensitiveDataReviewDrawer({ + open, + databaseId, + scopeLabel, + suggestions, + canManage, + onClose, + onSaved, +}: Props) { + const [drafts, setDrafts] = useState>({}); + const [search, setSearch] = useState(""); + const [showAll, setShowAll] = useState(false); + const [saving, setSaving] = useState(false); + + useEffect(() => { + if (!open) return; + setDrafts(Object.fromEntries(suggestions.map((suggestion) => [ + suggestion.columnId, + suggestion.sensitive, + ]))); + setSearch(""); + setShowAll(false); + }, [open, suggestions]); + + const changed = useMemo(() => suggestions.filter((suggestion) => ( + drafts[suggestion.columnId] !== undefined + && drafts[suggestion.columnId] !== suggestion.currentSensitive + )), [drafts, suggestions]); + const visible = useMemo(() => { + const term = search.trim().toLocaleLowerCase(); + return suggestions.filter((suggestion) => ( + (showAll || drafts[suggestion.columnId] !== suggestion.currentSensitive) + && (!term || `${suggestion.tableName}.${suggestion.columnName}`.toLocaleLowerCase().includes(term)) + )); + }, [drafts, search, showAll, suggestions]); + + const close = () => { + if (saving) return; + if (changed.length > 0 && !window.confirm("Discard the sensitive-field review?")) return; + onClose(); + }; + + const save = async () => { + if (!databaseId || changed.length === 0) return; + setSaving(true); + const results = await Promise.allSettled(changed.map((suggestion) => ( + updateCatalogColumnSensitive( + databaseId, + suggestion.tableId, + suggestion.columnId, + suggestion.version, + drafts[suggestion.columnId]!, + ) + ))); + const saved = results.flatMap((result) => result.status === "fulfilled" ? [result.value] : []); + const failed = results.flatMap((result) => result.status === "rejected" ? [result.reason] : []); + if (saved.length > 0) onSaved(saved); + if (failed.length > 0) { + toast.error(`${failed.length} sensitive flag${failed.length === 1 ? "" : "s"} could not be saved: ${apiErrorMessage(failed[0])}`); + } else { + toast.success(`Saved ${saved.length} sensitive flag${saved.length === 1 ? "" : "s"}`); + onClose(); + } + setSaving(false); + }; + + if (!open || !databaseId) return null; + return ( + + ); +}