feat: classify sensitive columns locally

backend/src/catalog/memory-repository.ts

@@ -35,10 +35,10 @@ import {
   type DescriptionGenerationRun,
   type DescriptionGenerationRunUpdate,
   type DescriptionGenerationScope,
-  type SensitiveDataSuggestionEvent,
-  type SensitiveDataSuggestionRun,
-  type SensitiveDataSuggestionRunUpdate,
-  type SensitiveDataSuggestionScope,
+  type SensitivityAnalysisEvent,
+  type SensitivityAnalysisRun,
+  type SensitivityAnalysisRunUpdate,
+  type SensitivityAnalysisScope,
   type TableSyncRepositoryResult,
   type WorkspaceDatabase,
 } from "./types.js";
@@ -56,8 +56,8 @@ export class MemoryCatalogRepository implements CatalogRepository {
   private readonly logicalRelationships = new Map<string, CatalogLogicalRelationship>();
   private readonly descriptionGenerationRuns = new Map<string, DescriptionGenerationRun>();
   private readonly descriptionGenerationEvents = new Map<string, DescriptionGenerationEvent[]>();
-  private readonly sensitiveDataSuggestionRuns = new Map<string, SensitiveDataSuggestionRun>();
-  private readonly sensitiveDataSuggestionEvents = new Map<string, SensitiveDataSuggestionEvent[]>();
+  private readonly sensitivityAnalysisRuns = new Map<string, SensitivityAnalysisRun>();
+  private readonly sensitivityAnalysisEvents = new Map<string, SensitivityAnalysisEvent[]>();
   private readonly syncRuns = new Map<string, CatalogSyncRun>();
   private readonly syncEvents = new Map<string, CatalogSyncEvent[]>();
 
@@ -183,10 +183,10 @@ export class MemoryCatalogRepository implements CatalogRepository {
       this.descriptionGenerationRuns.delete(runId);
       this.descriptionGenerationEvents.delete(runId);
     }
-    for (const [runId, run] of this.sensitiveDataSuggestionRuns) {
+    for (const [runId, run] of this.sensitivityAnalysisRuns) {
       if (run.databaseId !== id) continue;
-      this.sensitiveDataSuggestionRuns.delete(runId);
-      this.sensitiveDataSuggestionEvents.delete(runId);
+      this.sensitivityAnalysisRuns.delete(runId);
+      this.sensitivityAnalysisEvents.delete(runId);
     }
     return this.records.delete(id);
   }
@@ -433,93 +433,96 @@ export class MemoryCatalogRepository implements CatalogRepository {
       .map((event) => structuredClone(event));
   }
 
-  async createSensitiveDataSuggestionRun(
+  async createSensitivityAnalysisRun(
     databaseId: string,
-    scope: SensitiveDataSuggestionScope,
-    modelId: string,
-  ): Promise<SensitiveDataSuggestionRun> {
+    scope: SensitivityAnalysisScope,
+    origin: { engine: "llm"; modelId: string } | { engine: "local"; policyVersion: string },
+  ): Promise<SensitivityAnalysisRun> {
     const now = new Date().toISOString();
-    const run: SensitiveDataSuggestionRun = {
+    const run: SensitivityAnalysisRun = {
       id: randomUUID(),
       databaseId,
       scope,
-      modelId,
+      engine: origin.engine,
+      modelId: origin.engine === "llm" ? origin.modelId : null,
+      policyVersion: origin.engine === "local" ? origin.policyVersion : null,
       status: "running",
       total: 0,
       suggestedSensitive: 0,
-        suggestedNonSensitive: 0,
-        inputTokens: 0,
-        cacheReadTokens: 0,
-        outputTokens: 0,
+      suggestedNonSensitive: 0,
+      unknown: 0,
+      inputTokens: 0,
+      cacheReadTokens: 0,
+      outputTokens: 0,
       createdAt: now,
       startedAt: now,
       updatedAt: now,
       finishedAt: null,
       errorSummary: null,
     };
-    this.sensitiveDataSuggestionRuns.set(run.id, run);
+    this.sensitivityAnalysisRuns.set(run.id, run);
     return structuredClone(run);
   }
 
-  async getSensitiveDataSuggestionRun(
+  async getSensitivityAnalysisRun(
     runId: string,
-  ): Promise<SensitiveDataSuggestionRun | undefined> {
-    const run = this.sensitiveDataSuggestionRuns.get(runId);
+  ): Promise<SensitivityAnalysisRun | undefined> {
+    const run = this.sensitivityAnalysisRuns.get(runId);
     return run ? structuredClone(run) : undefined;
   }
 
-  async listSensitiveDataSuggestionRuns(limit = 50): Promise<SensitiveDataSuggestionRun[]> {
-    return [...this.sensitiveDataSuggestionRuns.values()]
+  async listSensitivityAnalysisRuns(limit = 50): Promise<SensitivityAnalysisRun[]> {
+    return [...this.sensitivityAnalysisRuns.values()]
       .sort((a, b) => b.createdAt.localeCompare(a.createdAt) || b.id.localeCompare(a.id))
       .slice(0, limit)
       .map((run) => structuredClone(run));
   }
 
-  async interruptActiveSensitiveDataSuggestionRuns(
+  async interruptActiveSensitivityAnalysisRuns(
     errorSummary: string,
-  ): Promise<SensitiveDataSuggestionRun[]> {
-    const interrupted: SensitiveDataSuggestionRun[] = [];
-    for (const run of this.sensitiveDataSuggestionRuns.values()) {
+  ): Promise<SensitivityAnalysisRun[]> {
+    const interrupted: SensitivityAnalysisRun[] = [];
+    for (const run of this.sensitivityAnalysisRuns.values()) {
       if (run.status !== "running") continue;
       const now = new Date().toISOString();
-      const updated: SensitiveDataSuggestionRun = {
+      const updated: SensitivityAnalysisRun = {
         ...run,
         status: "interrupted",
         updatedAt: now,
         finishedAt: now,
         errorSummary,
       };
-      this.sensitiveDataSuggestionRuns.set(run.id, updated);
+      this.sensitivityAnalysisRuns.set(run.id, updated);
       interrupted.push(structuredClone(updated));
     }
     return interrupted;
   }
 
-  async updateSensitiveDataSuggestionRun(
+  async updateSensitivityAnalysisRun(
     runId: string,
-    update: SensitiveDataSuggestionRunUpdate,
-  ): Promise<SensitiveDataSuggestionRun | undefined> {
-    const current = this.sensitiveDataSuggestionRuns.get(runId);
+    update: SensitivityAnalysisRunUpdate,
+  ): Promise<SensitivityAnalysisRun | undefined> {
+    const current = this.sensitivityAnalysisRuns.get(runId);
     if (!current) return undefined;
     const updated = {
       ...current,
       ...structuredClone(update),
       updatedAt: new Date().toISOString(),
     };
-    this.sensitiveDataSuggestionRuns.set(runId, updated);
+    this.sensitivityAnalysisRuns.set(runId, updated);
     return structuredClone(updated);
   }
 
-  async appendSensitiveDataSuggestionEvent(
+  async appendSensitivityAnalysisEvent(
     runId: string,
-    level: SensitiveDataSuggestionEvent["level"],
+    level: SensitivityAnalysisEvent["level"],
     message: string,
-  ): Promise<SensitiveDataSuggestionEvent> {
-    if (!this.sensitiveDataSuggestionRuns.has(runId)) {
-      throw new CatalogConflictError("Sensitive Data Suggestion Run does not exist");
+  ): Promise<SensitivityAnalysisEvent> {
+    if (!this.sensitivityAnalysisRuns.has(runId)) {
+      throw new CatalogConflictError("Sensitivity Analysis Run does not exist");
     }
-    const events = this.sensitiveDataSuggestionEvents.get(runId) ?? [];
-    const event: SensitiveDataSuggestionEvent = {
+    const events = this.sensitivityAnalysisEvents.get(runId) ?? [];
+    const event: SensitivityAnalysisEvent = {
       runId,
       sequence: events.length + 1,
       level,
@@ -527,15 +530,15 @@ export class MemoryCatalogRepository implements CatalogRepository {
       createdAt: new Date().toISOString(),
     };
     events.push(event);
-    this.sensitiveDataSuggestionEvents.set(runId, events);
+    this.sensitivityAnalysisEvents.set(runId, events);
     return structuredClone(event);
   }
 
-  async listSensitiveDataSuggestionEvents(
+  async listSensitivityAnalysisEvents(
     runId: string,
     afterSequence = 0,
-  ): Promise<SensitiveDataSuggestionEvent[]> {
-    return (this.sensitiveDataSuggestionEvents.get(runId) ?? [])
+  ): Promise<SensitivityAnalysisEvent[]> {
+    return (this.sensitivityAnalysisEvents.get(runId) ?? [])
       .filter((event) => event.sequence > afterSequence)
       .map((event) => structuredClone(event));
   }