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

import { z } from "zod";
import type { MetadataGenerationModels } from "./metadata-generation-models.js";
import type { ModelCompleter, ModelCompletionMessage } from "./model-completer.js";
import type {
  CatalogColumn,
  CatalogRepository,
  CatalogTable,
  SensitiveDataSuggestionScope,
} from "./types.js";

export type { SensitiveDataSuggestionScope } from "./types.js";

// 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()),
}).strict();

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(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");
    this.name = "SensitiveDataSuggestionInvalidResponseError";
  }
}

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,
    private readonly models: MetadataGenerationModels,
    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,
    onPrepared?: (total: number) => void | Promise<void>,
  ): Promise<readonly SensitiveDataSuggestion[]> {
    const database = await this.repository.get(databaseId);
    if (!database) throw new SensitiveDataSuggestionTargetNotFoundError("database");
    const columns = await this.selectColumns(databaseId, scope, targetIds);
    await onPrepared?.(columns.length);
    const model = this.models.resolve(modelId);
    const suggestions: SensitiveDataSuggestion[] = [];

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