feat: complete catalog fleet management workflow

backend/src/catalog/memory-repository.ts

@@ -1,4 +1,9 @@
 import { randomUUID } from "node:crypto";
+import {
+  createCatalogMetrics,
+  hasCatalogDescription,
+  latestCatalogTimestamp,
+} from "./metrics.js";
 import {
   CatalogConflictError,
   CatalogConnectorError,
@@ -8,6 +13,7 @@ import {
   type CatalogDescriptionTarget,
   type CatalogDatabaseMetadataDeleteTarget,
   type CatalogMetadataDeleteCounts,
+  type CatalogMetrics,
   type CatalogRelationship,
   type CatalogSchemaDiff,
   type CatalogSyncCounts,
@@ -26,6 +32,10 @@ import {
   type DescriptionGenerationRun,
   type DescriptionGenerationRunUpdate,
   type DescriptionGenerationScope,
+  type SensitiveDataSuggestionEvent,
+  type SensitiveDataSuggestionRun,
+  type SensitiveDataSuggestionRunUpdate,
+  type SensitiveDataSuggestionScope,
   type TableSyncRepositoryResult,
   type WorkspaceDatabase,
 } from "./types.js";
@@ -42,6 +52,8 @@ export class MemoryCatalogRepository implements CatalogRepository {
   private readonly relationships = new Map<string, CatalogRelationship>();
   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 syncRuns = new Map<string, CatalogSyncRun>();
   private readonly syncEvents = new Map<string, CatalogSyncEvent[]>();
 
@@ -56,6 +68,34 @@ export class MemoryCatalogRepository implements CatalogRepository {
     const value = [...this.records.values()].find((record) => record.workspaceId === workspaceId);
     return value ? clone(value) : undefined;
   }
+  async getCatalogMetrics(databaseId?: string): Promise<CatalogMetrics | undefined> {
+    if (databaseId !== undefined && !this.records.has(databaseId)) return undefined;
+
+    const selectedDatabaseIds = new Set(
+      databaseId === undefined ? this.records.keys() : [databaseId],
+    );
+    const tables = [...this.tables.values()]
+      .filter((table) => selectedDatabaseIds.has(table.databaseId));
+    const tableIds = new Set(tables.map((table) => table.id));
+    const columns = [...this.columns.values()]
+      .filter((column) => tableIds.has(column.tableId));
+    const relationships = [...this.relationships.values()]
+      .filter((relationship) => selectedDatabaseIds.has(relationship.databaseId));
+
+    return createCatalogMetrics(databaseId, {
+      tables: tables.length,
+      columns: columns.length,
+      sensitiveColumns: columns.filter((column) => column.sensitive).length,
+      relationships: relationships.length,
+      describedTables: tables.filter(hasCatalogDescription).length,
+      describedColumns: columns.filter(hasCatalogDescription).length,
+    }, latestCatalogTimestamp([
+      ...[...selectedDatabaseIds].map((id) => this.records.get(id)?.schemaSyncedAt),
+      ...tables.map((table) => table.updatedAt),
+      ...columns.map((column) => column.updatedAt),
+      ...relationships.map((relationship) => relationship.updatedAt),
+    ]));
+  }
   async create(input: DatabaseConfigurationInput): Promise<WorkspaceDatabase> {
     if ([...this.records.values()].some((record) => record.workspaceId === input.workspaceId)) {
       throw new CatalogConflictError("Workspace database already exists");
@@ -136,6 +176,11 @@ export class MemoryCatalogRepository implements CatalogRepository {
       this.descriptionGenerationRuns.delete(runId);
       this.descriptionGenerationEvents.delete(runId);
     }
+    for (const [runId, run] of this.sensitiveDataSuggestionRuns) {
+      if (run.databaseId !== id) continue;
+      this.sensitiveDataSuggestionRuns.delete(runId);
+      this.sensitiveDataSuggestionEvents.delete(runId);
+    }
     return this.records.delete(id);
   }
   async listTables(databaseId: string): Promise<CatalogTable[]> {
@@ -378,6 +423,110 @@ export class MemoryCatalogRepository implements CatalogRepository {
       .map((event) => structuredClone(event));
   }
 
+  async createSensitiveDataSuggestionRun(
+    databaseId: string,
+    scope: SensitiveDataSuggestionScope,
+    modelId: string,
+  ): Promise<SensitiveDataSuggestionRun> {
+    const now = new Date().toISOString();
+    const run: SensitiveDataSuggestionRun = {
+      id: randomUUID(),
+      databaseId,
+      scope,
+      modelId,
+      status: "running",
+      total: 0,
+      suggestedSensitive: 0,
+      suggestedNonSensitive: 0,
+      createdAt: now,
+      startedAt: now,
+      updatedAt: now,
+      finishedAt: null,
+      errorSummary: null,
+    };
+    this.sensitiveDataSuggestionRuns.set(run.id, run);
+    return structuredClone(run);
+  }
+
+  async getSensitiveDataSuggestionRun(
+    runId: string,
+  ): Promise<SensitiveDataSuggestionRun | undefined> {
+    const run = this.sensitiveDataSuggestionRuns.get(runId);
+    return run ? structuredClone(run) : undefined;
+  }
+
+  async listSensitiveDataSuggestionRuns(limit = 50): Promise<SensitiveDataSuggestionRun[]> {
+    return [...this.sensitiveDataSuggestionRuns.values()]
+      .sort((a, b) => b.createdAt.localeCompare(a.createdAt) || b.id.localeCompare(a.id))
+      .slice(0, limit)
+      .map((run) => structuredClone(run));
+  }
+
+  async interruptActiveSensitiveDataSuggestionRuns(
+    errorSummary: string,
+  ): Promise<SensitiveDataSuggestionRun[]> {
+    const interrupted: SensitiveDataSuggestionRun[] = [];
+    for (const run of this.sensitiveDataSuggestionRuns.values()) {
+      if (run.status !== "running") continue;
+      const now = new Date().toISOString();
+      const updated: SensitiveDataSuggestionRun = {
+        ...run,
+        status: "interrupted",
+        updatedAt: now,
+        finishedAt: now,
+        errorSummary,
+      };
+      this.sensitiveDataSuggestionRuns.set(run.id, updated);
+      interrupted.push(structuredClone(updated));
+    }
+    return interrupted;
+  }
+
+  async updateSensitiveDataSuggestionRun(
+    runId: string,
+    update: SensitiveDataSuggestionRunUpdate,
+  ): Promise<SensitiveDataSuggestionRun | undefined> {
+    const current = this.sensitiveDataSuggestionRuns.get(runId);
+    if (!current) return undefined;
+    const updated = {
+      ...current,
+      ...structuredClone(update),
+      updatedAt: new Date().toISOString(),
+    };
+    this.sensitiveDataSuggestionRuns.set(runId, updated);
+    return structuredClone(updated);
+  }
+
+  async appendSensitiveDataSuggestionEvent(
+    runId: string,
+    level: SensitiveDataSuggestionEvent["level"],
+    message: string,
+  ): Promise<SensitiveDataSuggestionEvent> {
+    if (!this.sensitiveDataSuggestionRuns.has(runId)) {
+      throw new CatalogConflictError("Sensitive Data Suggestion Run does not exist");
+    }
+    const events = this.sensitiveDataSuggestionEvents.get(runId) ?? [];
+    const event: SensitiveDataSuggestionEvent = {
+      runId,
+      sequence: events.length + 1,
+      level,
+      message,
+      createdAt: new Date().toISOString(),
+    };
+    events.push(event);
+    this.sensitiveDataSuggestionEvents.set(runId, events);
+    return structuredClone(event);
+  }
+
+  async listSensitiveDataSuggestionEvents(
+    runId: string,
+    afterSequence = 0,
+  ): Promise<SensitiveDataSuggestionEvent[]> {
+    return (this.sensitiveDataSuggestionEvents.get(runId) ?? [])
+      .filter((event) => event.sequence > afterSequence)
+      .map((event) => structuredClone(event));
+  }
+
   async listRelationships(databaseId: string): Promise<CatalogRelationship[]> {
     return [...this.relationships.values()].filter((relationship) => relationship.databaseId === databaseId)
       .sort((a, b) => `${a.sourceTableName}.${a.constraintName}`.localeCompare(`${b.sourceTableName}.${b.constraintName}`))