feat: classify sensitive columns locally

backend/src/catalog/repository.ts

@@ -45,10 +45,10 @@ import {
   type DescriptionGenerationScope,
   type ObservedCatalogTable,
   type ObservedSchemaSnapshot,
-  type SensitiveDataSuggestionEvent,
-  type SensitiveDataSuggestionRun,
-  type SensitiveDataSuggestionRunUpdate,
-  type SensitiveDataSuggestionScope,
+  type SensitivityAnalysisEvent,
+  type SensitivityAnalysisRun,
+  type SensitivityAnalysisRunUpdate,
+  type SensitivityAnalysisScope,
   type TableSyncRepositoryResult,
   type WorkspaceDatabase,
 } from "./types.js";
@@ -189,15 +189,18 @@ interface DescriptionGenerationEventTable {
   createdAt: Timestamp;
 }
 
-interface SensitiveDataSuggestionRunTable {
+interface SensitivityAnalysisRunTable {
   id: string;
   databaseId: string;
-  scope: SensitiveDataSuggestionScope;
-  modelId: string;
-  status: SensitiveDataSuggestionRun["status"];
+  scope: SensitivityAnalysisScope;
+  engine: SensitivityAnalysisRun["engine"];
+  modelId: string | null;
+  policyVersion: string | null;
+  status: SensitivityAnalysisRun["status"];
   total: number;
   suggestedSensitive: number;
   suggestedNonSensitive: number;
+  unknown: number;
   inputTokens: number;
   cacheReadTokens: number;
   outputTokens: number;
@@ -208,10 +211,10 @@ interface SensitiveDataSuggestionRunTable {
   errorSummary: string | null;
 }
 
-interface SensitiveDataSuggestionEventTable {
+interface SensitivityAnalysisEventTable {
   runId: string;
   sequence: number;
-  level: SensitiveDataSuggestionEvent["level"];
+  level: SensitivityAnalysisEvent["level"];
   message: string;
   createdAt: Timestamp;
 }
@@ -264,8 +267,9 @@ export interface CatalogDatabase {
   catalogLogicalRelationships: CatalogLogicalRelationshipTable;
   descriptionGenerationRuns: DescriptionGenerationRunTable;
   descriptionGenerationEvents: DescriptionGenerationEventTable;
-  sensitiveDataSuggestionRuns: SensitiveDataSuggestionRunTable;
-  sensitiveDataSuggestionEvents: SensitiveDataSuggestionEventTable;
+  // Legacy physical table names retained for migration and storage compatibility.
+  sensitiveDataSuggestionRuns: SensitivityAnalysisRunTable;
+  sensitiveDataSuggestionEvents: SensitivityAnalysisEventTable;
   catalogSyncRuns: CatalogSyncRunTable;
   catalogSyncEvents: CatalogSyncEventTable;
 }
@@ -402,9 +406,9 @@ function serializeDescriptionGenerationEvent(
   return { ...row, createdAt: new Date(row.createdAt).toISOString() };
 }
 
-function serializeSensitiveDataSuggestionRun(
-  row: Selectable<SensitiveDataSuggestionRunTable>,
-): SensitiveDataSuggestionRun {
+function serializeSensitivityAnalysisRun(
+  row: Selectable<SensitivityAnalysisRunTable>,
+): SensitivityAnalysisRun {
   const stamp = (value: Date | string | null) => value === null ? null : new Date(value).toISOString();
   return {
     ...row,
@@ -415,9 +419,9 @@ function serializeSensitiveDataSuggestionRun(
   };
 }
 
-function serializeSensitiveDataSuggestionEvent(
-  row: Selectable<SensitiveDataSuggestionEventTable>,
-): SensitiveDataSuggestionEvent {
+function serializeSensitivityAnalysisEvent(
+  row: Selectable<SensitivityAnalysisEventTable>,
+): SensitivityAnalysisEvent {
   return { ...row, createdAt: new Date(row.createdAt).toISOString() };
 }
 
@@ -938,52 +942,55 @@ export class KyselyCatalogRepository implements CatalogRepository {
     return rows.map(serializeDescriptionGenerationEvent);
   }
 
-  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 row = await this.db.insertInto("sensitiveDataSuggestionRuns").values({
       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,
+      unknown: 0,
       inputTokens: 0,
       cacheReadTokens: 0,
       outputTokens: 0,
       finishedAt: null,
       errorSummary: null,
     }).returningAll().executeTakeFirstOrThrow();
-    return serializeSensitiveDataSuggestionRun(row);
+    return serializeSensitivityAnalysisRun(row);
   }
 
-  async getSensitiveDataSuggestionRun(
+  async getSensitivityAnalysisRun(
     runId: string,
-  ): Promise<SensitiveDataSuggestionRun | undefined> {
+  ): Promise<SensitivityAnalysisRun | undefined> {
     const row = await this.db.selectFrom("sensitiveDataSuggestionRuns")
       .selectAll()
       .where("id", "=", runId)
       .executeTakeFirst();
-    return row ? serializeSensitiveDataSuggestionRun(row) : undefined;
+    return row ? serializeSensitivityAnalysisRun(row) : undefined;
   }
 
-  async listSensitiveDataSuggestionRuns(limit = 50): Promise<SensitiveDataSuggestionRun[]> {
+  async listSensitivityAnalysisRuns(limit = 50): Promise<SensitivityAnalysisRun[]> {
     const rows = await this.db.selectFrom("sensitiveDataSuggestionRuns")
       .selectAll()
       .orderBy("createdAt", "desc")
       .orderBy("id", "desc")
       .limit(limit)
       .execute();
-    return rows.map(serializeSensitiveDataSuggestionRun);
+    return rows.map(serializeSensitivityAnalysisRun);
   }
 
-  async interruptActiveSensitiveDataSuggestionRuns(
+  async interruptActiveSensitivityAnalysisRuns(
     errorSummary: string,
-  ): Promise<SensitiveDataSuggestionRun[]> {
+  ): Promise<SensitivityAnalysisRun[]> {
     const rows = await this.db.updateTable("sensitiveDataSuggestionRuns")
       .set({
         status: "interrupted",
@@ -994,34 +1001,34 @@ export class KyselyCatalogRepository implements CatalogRepository {
       .where("status", "=", "running")
       .returningAll()
       .execute();
-    return rows.map(serializeSensitiveDataSuggestionRun);
+    return rows.map(serializeSensitivityAnalysisRun);
   }
 
-  async updateSensitiveDataSuggestionRun(
+  async updateSensitivityAnalysisRun(
     runId: string,
-    update: SensitiveDataSuggestionRunUpdate,
-  ): Promise<SensitiveDataSuggestionRun | undefined> {
+    update: SensitivityAnalysisRunUpdate,
+  ): Promise<SensitivityAnalysisRun | undefined> {
     const values: any = { ...update, updatedAt: sql`now()` };
     const row = await this.db.updateTable("sensitiveDataSuggestionRuns")
       .set(values)
       .where("id", "=", runId)
       .returningAll()
       .executeTakeFirst();
-    return row ? serializeSensitiveDataSuggestionRun(row) : undefined;
+    return row ? serializeSensitivityAnalysisRun(row) : undefined;
   }
 
-  async appendSensitiveDataSuggestionEvent(
+  async appendSensitivityAnalysisEvent(
     runId: string,
-    level: SensitiveDataSuggestionEvent["level"],
+    level: SensitivityAnalysisEvent["level"],
     message: string,
-  ): Promise<SensitiveDataSuggestionEvent> {
+  ): Promise<SensitivityAnalysisEvent> {
     return await this.db.transaction().execute(async (trx) => {
       const run = await trx.selectFrom("sensitiveDataSuggestionRuns")
         .select("id")
         .where("id", "=", runId)
         .forUpdate()
         .executeTakeFirst();
-      if (!run) throw new CatalogConflictError("Sensitive Data Suggestion Run does not exist");
+      if (!run) throw new CatalogConflictError("Sensitivity Analysis Run does not exist");
       const current = await trx.selectFrom("sensitiveDataSuggestionEvents")
         .select(sql<number>`coalesce(max(sequence), 0)::int`.as("sequence"))
         .where("runId", "=", runId)
@@ -1032,21 +1039,21 @@ export class KyselyCatalogRepository implements CatalogRepository {
         level,
         message,
       }).returningAll().executeTakeFirstOrThrow();
-      return serializeSensitiveDataSuggestionEvent(row);
+      return serializeSensitivityAnalysisEvent(row);
     });
   }
 
-  async listSensitiveDataSuggestionEvents(
+  async listSensitivityAnalysisEvents(
     runId: string,
     afterSequence = 0,
-  ): Promise<SensitiveDataSuggestionEvent[]> {
+  ): Promise<SensitivityAnalysisEvent[]> {
     const rows = await this.db.selectFrom("sensitiveDataSuggestionEvents")
       .selectAll()
       .where("runId", "=", runId)
       .where("sequence", ">", afterSequence)
       .orderBy("sequence")
       .execute();
-    return rows.map(serializeSensitiveDataSuggestionEvent);
+    return rows.map(serializeSensitivityAnalysisEvent);
   }
 
   async listRelationships(databaseId: string): Promise<CatalogPhysicalRelationship[]> {
@@ -1792,13 +1799,13 @@ export class UnavailableCatalogRepository implements CatalogRepository {
   async updateDescriptionGenerationRun(): Promise<DescriptionGenerationRun | undefined> { return this.fail(); }
   async appendDescriptionGenerationEvent(): Promise<DescriptionGenerationEvent> { return this.fail(); }
   async listDescriptionGenerationEvents(): Promise<DescriptionGenerationEvent[]> { return this.fail(); }
-  async createSensitiveDataSuggestionRun(): Promise<SensitiveDataSuggestionRun> { return this.fail(); }
-  async getSensitiveDataSuggestionRun(): Promise<SensitiveDataSuggestionRun | undefined> { return this.fail(); }
-  async listSensitiveDataSuggestionRuns(): Promise<SensitiveDataSuggestionRun[]> { return this.fail(); }
-  async interruptActiveSensitiveDataSuggestionRuns(): Promise<SensitiveDataSuggestionRun[]> { return this.fail(); }
-  async updateSensitiveDataSuggestionRun(): Promise<SensitiveDataSuggestionRun | undefined> { return this.fail(); }
-  async appendSensitiveDataSuggestionEvent(): Promise<SensitiveDataSuggestionEvent> { return this.fail(); }
-  async listSensitiveDataSuggestionEvents(): Promise<SensitiveDataSuggestionEvent[]> { return this.fail(); }
+  async createSensitivityAnalysisRun(): Promise<SensitivityAnalysisRun> { return this.fail(); }
+  async getSensitivityAnalysisRun(): Promise<SensitivityAnalysisRun | undefined> { return this.fail(); }
+  async listSensitivityAnalysisRuns(): Promise<SensitivityAnalysisRun[]> { return this.fail(); }
+  async interruptActiveSensitivityAnalysisRuns(): Promise<SensitivityAnalysisRun[]> { return this.fail(); }
+  async updateSensitivityAnalysisRun(): Promise<SensitivityAnalysisRun | undefined> { return this.fail(); }
+  async appendSensitivityAnalysisEvent(): Promise<SensitivityAnalysisEvent> { return this.fail(); }
+  async listSensitivityAnalysisEvents(): Promise<SensitivityAnalysisEvent[]> { return this.fail(); }
   async listRelationships(): Promise<CatalogPhysicalRelationship[]> { return this.fail(); }
   async listLogicalRelationships(): Promise<CatalogLogicalRelationship[]> { return this.fail(); }
   async getLogicalRelationshipContext(): Promise<CatalogLogicalRelationshipContext | undefined> { return this.fail(); }