feat: refine metadata catalog workflows

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

@@ -10,6 +10,7 @@ import type {
   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.";
@@ -73,6 +74,31 @@ export class SensitiveDataSuggestionRunner {
           });
           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;