feat: protect sensitive catalog samples

backend/src/catalog/sensitive-data-suggester.ts

@@ -0,0 +1,103 @@
+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";
+
+const MAX_COLUMNS = 10_000;
+const responseSchema = z.object({
+  suggestions: z.array(z.object({
+    columnId: z.uuid(),
+    sensitive: z.boolean(),
+  }).strict()).max(MAX_COLUMNS),
+}).strict();
+
+export interface SensitiveDataSuggestion {
+  columnId: string;
+  sensitive: boolean;
+}
+
+export class SensitiveDataSuggestionTargetNotFoundError extends Error {
+  constructor() {
+    super("database not found");
+    this.name = "SensitiveDataSuggestionTargetNotFoundError";
+  }
+}
+
+export class SensitiveDataSuggestionInvalidResponseError extends Error {
+  constructor() {
+    super("sensitive-data suggestion response is invalid");
+    this.name = "SensitiveDataSuggestionInvalidResponseError";
+  }
+}
+
+export class SensitiveDataSuggester {
+  constructor(
+    private readonly repository: CatalogRepository,
+    private readonly models: MetadataGenerationModels,
+    private readonly completer: ModelCompleter,
+  ) {}
+
+  async suggest(
+    databaseId: string,
+    modelId: string,
+    signal: AbortSignal,
+  ): Promise<readonly SensitiveDataSuggestion[]> {
+    const database = await this.repository.get(databaseId);
+    if (!database) throw new SensitiveDataSuggestionTargetNotFoundError();
+
+    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();
+      }
+      return parsed.suggestions;
+    } catch (error) {
+      if (error instanceof SensitiveDataSuggestionInvalidResponseError) throw error;
+      throw new SensitiveDataSuggestionInvalidResponseError();
+    }
+  }
+}