ThothII/backend/src/catalog/sensitive-data-suggestion-runner.ts

import type {
  SensitiveDataSuggestion,
} from "./sensitive-data-suggester.js";
import {
  SensitiveDataSuggester,
  SensitiveDataSuggestionTargetNotFoundError,
} from "./sensitive-data-suggester.js";
import type {
  CatalogRepository,
  SensitiveDataSuggestionRun,
  SensitiveDataSuggestionScope,
} from "./types.js";
import type { ModelCompletionUsage } from "./model-completer.js";

const interruptedMessage = "Sensitive-field suggestion generation was interrupted by backend restart.";
const failedMessage = "Sensitive-field suggestion generation failed.";

export interface SensitiveDataSuggestionRunResult {
  suggestions: readonly SensitiveDataSuggestion[];
  run: SensitiveDataSuggestionRun;
}

export class SensitiveDataSuggestionRunner {
  constructor(
    private readonly repository: CatalogRepository,
    private readonly suggester: SensitiveDataSuggester,
  ) {}

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

  async run(
    databaseId: string,
    modelId: string,
    scope: SensitiveDataSuggestionScope,
    targetIds: readonly string[],
    signal: AbortSignal,
  ): Promise<SensitiveDataSuggestionRunResult> {
    if (!(await this.repository.get(databaseId))) {
      throw new SensitiveDataSuggestionTargetNotFoundError("database");
    }
    const started = await this.repository.createSensitiveDataSuggestionRun(
      databaseId,
      scope,
      modelId,
    );

    try {
      await this.repository.appendSensitiveDataSuggestionEvent(
        started.id,
        "info",
        "Sensitive-field suggestion generation started.",
      );
      const suggestions = await this.suggester.suggest(
        databaseId,
        modelId,
        scope,
        targetIds,
        signal,
        async (total) => {
          const prepared = await this.repository.updateSensitiveDataSuggestionRun(started.id, {
            total,
          });
          if (!prepared) throw new Error("Sensitive Data Suggestion Run disappeared");
        },
        async (processed, batch) => {
          const suggestedSensitive = batch.filter((suggestion) => suggestion.sensitive).length;
          const suggestedNonSensitive = batch.length - suggestedSensitive;
          const current = await this.repository.getSensitiveDataSuggestionRun(started.id);
          if (!current) throw new Error("Sensitive Data Suggestion Run disappeared");
          const progress = await this.repository.updateSensitiveDataSuggestionRun(started.id, {
            suggestedSensitive: current.suggestedSensitive + suggestedSensitive,
            suggestedNonSensitive: current.suggestedNonSensitive + suggestedNonSensitive,
          });
          if (!progress) throw new Error("Sensitive Data Suggestion Run disappeared");
          await this.repository.appendSensitiveDataSuggestionEvent(
            started.id,
            "info",
            `Classified ${processed} of ${progress.total} columns.`,
          );
        },
        async (usage: ModelCompletionUsage) => {
          const current = await this.repository.getSensitiveDataSuggestionRun(started.id);
          if (!current) throw new Error("Sensitive Data Suggestion Run disappeared");
          await this.repository.updateSensitiveDataSuggestionRun(started.id, {
            inputTokens: current.inputTokens + usage.input,
            cacheReadTokens: current.cacheReadTokens + usage.cacheRead,
            outputTokens: current.outputTokens + usage.output,
          });
        },
      );
      const suggestedSensitive = suggestions.filter((suggestion) => suggestion.sensitive).length;
      const suggestedNonSensitive = suggestions.length - suggestedSensitive;
      await this.repository.appendSensitiveDataSuggestionEvent(
        started.id,
        "info",
        `Sensitive-field suggestion generation completed for ${suggestions.length} column${
          suggestions.length === 1 ? "" : "s"
        }.`,
      );
      const completed = await this.repository.updateSensitiveDataSuggestionRun(started.id, {
        status: "completed",
        total: suggestions.length,
        suggestedSensitive,
        suggestedNonSensitive,
        finishedAt: new Date().toISOString(),
        errorSummary: null,
      });
      if (!completed) throw new Error("Sensitive Data Suggestion Run disappeared");
      return { suggestions, run: completed };
    } catch (error) {
      await this.repository.updateSensitiveDataSuggestionRun(started.id, {
        status: "failed",
        finishedAt: new Date().toISOString(),
        errorSummary: failedMessage,
      }).catch(() => undefined);
      await this.repository.appendSensitiveDataSuggestionEvent(
        started.id,
        "error",
        failedMessage,
      ).catch(() => undefined);
      throw error;
    }
  }
}