ThothII/backend/src/catalog/sensitivity-analysis-service.ts

import type {
  SensitivityClassifier,
  SensitivityColumnAssessment,
  SensitivityEvidence,
  SensitivityNerBudget,
} from "./sensitivity-classifier.js";
import type {
  CatalogColumn,
  CatalogRepository,
  CatalogTable,
  SensitivityAnalysisScope,
} from "./types.js";

export type { SensitivityAnalysisScope } from "./types.js";
export const SENSITIVITY_POLICY_VERSION = "sensitivity-v4";

interface SelectedColumn {
  table: CatalogTable;
  column: CatalogColumn;
}

export interface SensitivityReviewItem {
  columnId: string;
  tableId: string;
  tableName: string;
  columnName: string;
  version: number;
  currentSensitive: boolean;
  sensitive: boolean;
  assessment: SensitivityColumnAssessment["assessment"];
  evidence: readonly SensitivityEvidence[];
  observedValues: number;
  coverage: SensitivityColumnAssessment["coverage"];
}

export class SensitivityAnalysisTargetNotFoundError extends Error {
  constructor(readonly target: "database" | "table" | "column") {
    super(`${target} not found`);
    this.name = "SensitivityAnalysisTargetNotFoundError";
  }
}

export class SensitivityAnalysisDuplicateTargetIdsError extends Error {
  constructor() {
    super("sensitivity analysis target IDs must be unique");
    this.name = "SensitivityAnalysisDuplicateTargetIdsError";
  }
}

export class SensitivityAnalysisInterruptedError extends Error {
  constructor() {
    super("sensitivity analysis interrupted");
    this.name = "SensitivityAnalysisInterruptedError";
  }
}

function ensureActive(signal: AbortSignal): void {
  if (signal.aborted) throw new SensitivityAnalysisInterruptedError();
}

export class SensitivityAnalysisNoEligibleColumnsError extends Error {
  constructor(readonly scope: SensitivityAnalysisScope) {
    super("selected scope has no catalog columns");
    this.name = "SensitivityAnalysisNoEligibleColumnsError";
  }
}

/** Selection and table orchestration around the single SensitivityClassifier decision module. */
export class SensitivityAnalysisService {
  constructor(
    private readonly repository: CatalogRepository,
    private readonly classifier: SensitivityClassifier,
    private readonly options: { nerBudgetMs?: number } = {},
  ) {}

  private async selectColumns(
    databaseId: string,
    scope: SensitivityAnalysisScope,
    targetIds: readonly string[],
    signal: AbortSignal,
  ): Promise<readonly SelectedColumn[]> {
    ensureActive(signal);
    if (new Set(targetIds).size !== targetIds.length) {
      throw new SensitivityAnalysisDuplicateTargetIdsError();
    }
    const tables = await this.repository.listTables(databaseId);
    ensureActive(signal);
    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 SensitivityAnalysisTargetNotFoundError("table");
    }
    const columns = (await Promise.all(selectedTables.map(async (table) => (
      (await this.repository.listColumns(databaseId, table.id)).map((column) => ({ table, column }))
    )))).flat();
    ensureActive(signal);
    const columnIds = new Set(targetIds);
    const selectedColumns = scope === "selected_columns"
      ? columns.filter(({ column }) => columnIds.has(column.id))
      : columns;
    if (scope === "selected_columns" && selectedColumns.length !== targetIds.length) {
      throw new SensitivityAnalysisTargetNotFoundError("column");
    }
    if (selectedColumns.length === 0) {
      throw new SensitivityAnalysisNoEligibleColumnsError(scope);
    }
    return selectedColumns;
  }

  async analyze(
    databaseId: string,
    scope: SensitivityAnalysisScope,
    targetIds: readonly string[],
    signal: AbortSignal,
    onPrepared?: (total: number) => void | Promise<void>,
    onProgress?: (processed: number, suggestions: readonly SensitivityReviewItem[]) => void | Promise<void>,
    onActivity?: (message: string) => void | Promise<void>,
  ): Promise<readonly SensitivityReviewItem[]> {
    const configuredNerBudget = this.options.nerBudgetMs ?? 10_000;
    const nerBudget: SensitivityNerBudget = {
      remainingMs: Number.isFinite(configuredNerBudget) && configuredNerBudget >= 0
        ? configuredNerBudget
        : 10_000,
    };
    ensureActive(signal);
    const database = await this.repository.get(databaseId);
    ensureActive(signal);
    if (!database) throw new SensitivityAnalysisTargetNotFoundError("database");
    const selected = await this.selectColumns(databaseId, scope, targetIds, signal);
    await onPrepared?.(selected.length);
    ensureActive(signal);
    const byTable = new Map<string, SelectedColumn[]>();
    for (const item of selected) {
      const items = byTable.get(item.table.id) ?? [];
      items.push(item);
      byTable.set(item.table.id, items);
    }
    const tableTargets = [...byTable.values()].map((items) => {
      const first = items[0]!;
      return {
        database,
        table: first.table,
        columns: items.map(({ column }) => column),
      };
    });
    const assessments = await this.classifier.assess(
      tableTargets,
      signal,
      nerBudget,
      onActivity,
    );
    ensureActive(signal);
    const assessmentById = new Map(assessments.map((assessment) => [
      assessment.columnId,
      assessment,
    ]));
    const suggestions: SensitivityReviewItem[] = [];
    for (const items of byTable.values()) {
      ensureActive(signal);
      const batch = items.map(({ table, column }) => {
        const assessment = assessmentById.get(column.id)!;
        return {
          columnId: column.id,
          tableId: table.id,
          tableName: table.name,
          columnName: column.name,
          version: column.version,
          currentSensitive: column.sensitive,
          sensitive: assessment.proposedSensitive,
          assessment: assessment.assessment,
          evidence: assessment.evidence,
          observedValues: assessment.observedValues,
          coverage: assessment.coverage,
        };
      });
      suggestions.push(...batch);
      await onProgress?.(suggestions.length, batch);
      ensureActive(signal);
    }
    return suggestions;
  }
}