fix: scope and batch sensitive suggestions

2026-08-30 17:12:20 +02:00
parent 0736983bc5
commit cb40c09d9a
17 changed files with 1179 additions and 139 deletions
@@ -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<StructuralColumn[]> {
+    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<readonly SensitiveDataSuggestion[]> {
    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<string, { columnId: string; sensitive: boolean }> | 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;
  }
 }
@@ -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);
    }
  });

@@ -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." });