fix: scope and batch sensitive suggestions

frontend/src/shell/DatabaseManagementPage.tsx

@@ -20,9 +20,11 @@ import {
   replaceCatalogDatabaseSecrets,
   startCatalogSync,
   startDescriptionGenerationRun,
+  suggestSensitiveFields,
   testCatalogDatabase,
   updateCatalogDatabase,
   type CatalogDatabase,
+  type CatalogColumn,
   type CatalogDatabaseMetadataDeleteTarget,
   type CatalogSecretName,
   type CatalogSyncScope,
@@ -30,6 +32,8 @@ import {
   type DatabaseBinding,
   type DatabaseTransport,
   type DescriptionGenerationRun,
+  type SensitiveDataSuggestion,
+  type SensitiveDataSuggestionRequest,
 } from "../api/catalog-databases";
 import { DatabaseGrid } from "./database-management/DatabaseGrid";
 import { DatabaseForm } from "./database-management/DatabaseForm";
@@ -38,6 +42,7 @@ import { DatabaseRelationships } from "./database-management/DatabaseRelationshi
 import { CatalogSyncDrawer } from "./database-management/CatalogSyncDrawer";
 import { MetadataGenerationModelSelector } from "./database-management/MetadataGenerationModelSelector";
 import { DescriptionGenerationDrawer } from "./database-management/DescriptionGenerationDrawer";
+import { SensitiveDataReviewDrawer } from "./database-management/SensitiveDataReviewDrawer";
 import {
   configurationFingerprint,
   configurationFromDraft,
@@ -141,6 +146,11 @@ export function DatabaseManagementPage({
   const [syncDrawerOpen, setSyncDrawerOpen] = useState(false);
   const [activeDescriptionGenerationRun, setActiveDescriptionGenerationRun] = useState<DescriptionGenerationRun | null>(null);
   const [descriptionGenerationDrawerOpen, setDescriptionGenerationDrawerOpen] = useState(false);
+  const [sensitiveReview, setSensitiveReview] = useState<{
+    databaseId: string;
+    scopeLabel: string;
+    suggestions: SensitiveDataSuggestion[];
+  } | null>(null);
 
   const originRef = useRef<HTMLElement | null>(null);
   const searchInputRef = useRef<HTMLInputElement>(null);
@@ -683,6 +693,67 @@ export function DatabaseManagementPage({
     }
   }, [rememberDescriptionGenerationRun, selectedMetadataModel]);
 
+  const requestSensitiveSuggestions = useCallback(async (
+    database: CatalogDatabase,
+    selection: SensitiveDataSuggestionRequest,
+    scopeLabel: string,
+  ) => {
+    if (!database.id) {
+      toast.error("The selected database is not configured, so sensitive-field suggestions were not requested.");
+      throw new Error("database is not configured");
+    }
+    if (!selectedMetadataModel) {
+      toast.error("Select a metadata-generation model before requesting sensitive-field suggestions.");
+      throw new Error("metadata-generation model is not selected");
+    }
+    try {
+      const result = await suggestSensitiveFields(database.id, selectedMetadataModel, selection);
+      setSensitiveReview({ databaseId: database.id, scopeLabel, suggestions: result.suggestions });
+      toast.success(`Prepared ${result.suggestions.length} sensitive-field suggestion${result.suggestions.length === 1 ? "" : "s"} for review`);
+    } catch (error) {
+      toast.error(apiErrorMessage(error));
+      throw error;
+    }
+  }, [selectedMetadataModel]);
+
+  const suggestDatabaseSensitiveFields = useCallback(async (selected: CatalogDatabase[]) => {
+    if (selected.length !== 1) {
+      toast.error("Sensitive-field suggestions can be requested for only one database at a time. Select one database and try again.");
+      throw new Error("more than one database selected");
+    }
+    const database = selected[0]!;
+    await requestSensitiveSuggestions(database, { scope: "all" }, `Entire database ${database.workspaceName}`);
+  }, [requestSensitiveSuggestions]);
+
+  const suggestActiveDatabaseSensitiveFields = useCallback(async (
+    selection: SensitiveDataSuggestionRequest,
+    scopeLabel: string,
+  ) => {
+    if (!activeRow) {
+      toast.error("The database is no longer available, so sensitive-field suggestions were not requested.");
+      throw new Error("database is no longer available");
+    }
+    await requestSensitiveSuggestions(activeRow, selection, scopeLabel);
+  }, [activeRow, requestSensitiveSuggestions]);
+
+  const sensitiveColumnsSaved = useCallback((columns: CatalogColumn[]) => {
+    const savedById = new Map(columns.map((column) => [column.id, column]));
+    setSensitiveReview((current) => current ? {
+      ...current,
+      suggestions: current.suggestions.map((suggestion) => {
+        const saved = savedById.get(suggestion.columnId);
+        return saved ? {
+          ...suggestion,
+          version: saved.version,
+          currentSensitive: saved.sensitive,
+          sensitive: saved.sensitive,
+        } : suggestion;
+      }),
+    } : null);
+    const databaseId = sensitiveReview?.databaseId;
+    if (databaseId) void queryClient.invalidateQueries({ queryKey: ["catalog-columns", databaseId] });
+  }, [queryClient, sensitiveReview?.databaseId]);
+
   const deleteSelectedMetadata = useCallback(async (
     selected: CatalogDatabase[],
     target: CatalogDatabaseMetadataDeleteTarget,
@@ -832,6 +903,7 @@ export function DatabaseManagementPage({
               selectedMetadataModel={selectedMetadataModelAvailable ? selectedMetadataModel : null}
               descriptionGenerationActive={descriptionGenerationActive}
               onGenerateDescriptions={generateDatabaseDescriptions}
+              onSuggestSensitive={suggestDatabaseSensitiveFields}
               onDeleteMetadataSelected={deleteSelectedMetadata}
             />
           )}
@@ -886,6 +958,7 @@ export function DatabaseManagementPage({
             onRunStarted={rememberSyncRun}
             onOpenSync={() => openSync(activeRow)}
             onDescriptionGenerationRunStarted={rememberDescriptionGenerationRun}
+            onSuggestSensitive={suggestActiveDatabaseSensitiveFields}
           />
         ) : null}
 
@@ -920,6 +993,15 @@ export function DatabaseManagementPage({
         onRunUpdate={setActiveDescriptionGenerationRun}
         onTerminal={(run) => void descriptionGenerationTerminated(run)}
       />
+      <SensitiveDataReviewDrawer
+        open={Boolean(sensitiveReview)}
+        databaseId={sensitiveReview?.databaseId ?? null}
+        scopeLabel={sensitiveReview?.scopeLabel ?? ""}
+        suggestions={sensitiveReview?.suggestions ?? []}
+        canManage={canManage}
+        onClose={() => setSensitiveReview(null)}
+        onSaved={sensitiveColumnsSaved}
+      />
     </main>
   );
 }