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;
  }
 }