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

import type {
  SensitivityReviewItem,
} from "./sensitivity-analysis-service.js";
import {
  SENSITIVITY_POLICY_VERSION,
  SensitivityAnalysisInterruptedError,
  SensitivityAnalysisService,
  SensitivityAnalysisTargetNotFoundError,
} from "./sensitivity-analysis-service.js";
import type {
  CatalogRepository,
  SensitivityAnalysisRun,
  SensitivityAnalysisScope,
} from "./types.js";

const interruptedMessage = "Local sensitivity analysis was interrupted by backend restart.";
const interruptedDuringRunMessage = "Local sensitivity analysis was interrupted before completion.";
const failedMessage = "Local sensitivity analysis failed.";

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

export interface SensitivityAnalysisRunResult {
  suggestions: readonly SensitivityReviewItem[];
  run: SensitivityAnalysisRun;
}

export class SensitivityAnalysisRunner {
  constructor(
    private readonly repository: CatalogRepository,
    private readonly analysis: SensitivityAnalysisService,
  ) {}

  async initialize(): Promise<void> {
    if (!(await this.repository.available())) return;
    const interrupted = await this.repository.interruptActiveSensitivityAnalysisRuns(
      interruptedMessage,
    );
    for (const run of interrupted) {
      await this.repository.appendSensitivityAnalysisEvent(
        run.id,
        "warning",
        interruptedMessage,
      );
    }
  }

  async run(
    databaseId: string,
    scope: SensitivityAnalysisScope,
    targetIds: readonly string[],
    signal: AbortSignal,
  ): Promise<SensitivityAnalysisRunResult> {
    ensureActive(signal);
    const database = await this.repository.get(databaseId);
    ensureActive(signal);
    if (!database) {
      throw new SensitivityAnalysisTargetNotFoundError("database");
    }
    ensureActive(signal);
    const started = await this.repository.createSensitivityAnalysisRun(
      databaseId,
      scope,
      { engine: "local", policyVersion: SENSITIVITY_POLICY_VERSION },
    );
    let preparedTotal = 0;
    let processedSensitive = 0;
    let processedNonSensitive = 0;

    try {
      ensureActive(signal);
      await this.repository.appendSensitivityAnalysisEvent(
        started.id,
        "info",
        "Local sensitivity analysis started.",
      );
      ensureActive(signal);
      const suggestions = await this.analysis.analyze(
        databaseId,
        scope,
        targetIds,
        signal,
        async (total) => {
          ensureActive(signal);
          preparedTotal = total;
          const prepared = await this.repository.updateSensitivityAnalysisRun(started.id, {
            total,
          });
          ensureActive(signal);
          if (!prepared) throw new Error("Sensitivity Analysis Run disappeared");
        },
        async (processed, batch) => {
          ensureActive(signal);
          const suggestedSensitive = batch.filter(
            (suggestion) => suggestion.assessment === "sensitive",
          ).length;
          const suggestedNonSensitive = batch.filter(
            (suggestion) => suggestion.assessment === "non_sensitive",
          ).length;
          const current = await this.repository.getSensitivityAnalysisRun(started.id);
          ensureActive(signal);
          if (!current) throw new Error("Sensitivity Analysis Run disappeared");
          const progress = await this.repository.updateSensitivityAnalysisRun(started.id, {
            suggestedSensitive: current.suggestedSensitive + suggestedSensitive,
            suggestedNonSensitive: current.suggestedNonSensitive + suggestedNonSensitive,
          });
          if (!progress) throw new Error("Sensitivity Analysis Run disappeared");
          processedSensitive += suggestedSensitive;
          processedNonSensitive += suggestedNonSensitive;
          ensureActive(signal);
          await this.repository.appendSensitivityAnalysisEvent(
            started.id,
            "info",
            `Assessed ${processed} of ${progress.total} columns locally.`,
          );
          ensureActive(signal);
        },
        async (message) => {
          ensureActive(signal);
          await this.repository.appendSensitivityAnalysisEvent(
            started.id,
            "info",
            message,
          );
          ensureActive(signal);
        },
      );
      ensureActive(signal);
      const suggestedSensitive = suggestions.filter(
        (suggestion) => suggestion.assessment === "sensitive",
      ).length;
      const suggestedNonSensitive = suggestions.filter(
        (suggestion) => suggestion.assessment === "non_sensitive",
      ).length;
      await this.repository.appendSensitivityAnalysisEvent(
        started.id,
        "info",
        `Local sensitivity analysis completed for ${suggestions.length} column${
          suggestions.length === 1 ? "" : "s"
        }.`,
      );
      ensureActive(signal);
      const completed = await this.repository.updateSensitivityAnalysisRun(started.id, {
        status: "completed",
        total: suggestions.length,
        suggestedSensitive,
        suggestedNonSensitive,
        unknown: 0,
        finishedAt: new Date().toISOString(),
        errorSummary: null,
      });
      ensureActive(signal);
      if (!completed) throw new Error("Sensitivity Analysis Run disappeared");
      return { suggestions, run: completed };
    } catch (error) {
      const interrupted = signal.aborted || error instanceof SensitivityAnalysisInterruptedError;
      const message = interrupted ? interruptedDuringRunMessage : failedMessage;
      await this.repository.updateSensitivityAnalysisRun(started.id, {
        status: interrupted ? "interrupted" : "failed",
        ...(interrupted ? {
          total: preparedTotal,
          suggestedSensitive: processedSensitive,
          suggestedNonSensitive: processedNonSensitive,
          unknown: Math.max(0, preparedTotal - processedSensitive - processedNonSensitive),
        } : {}),
        finishedAt: new Date().toISOString(),
        errorSummary: message,
      }).catch(() => undefined);
      await this.repository.appendSensitivityAnalysisEvent(
        started.id,
        interrupted ? "warning" : "error",
        message,
      ).catch(() => undefined);
      throw error;
    }
  }
}