feat: sample sensitive columns progressively

backend/src/routes/catalog-description-generation.ts

@@ -261,9 +261,9 @@ function safeSuggestionError(reply: FastifyReply, error: unknown) {
     });
   }
   if (error instanceof SensitivityAnalysisInterruptedError) {
-    return reply.code(504).send({
-      code: "sensitivity_analysis_timeout",
-      message: "Sensitivity analysis reached its time limit. No assessments were applied.",
+    return reply.code(499).send({
+      code: "sensitivity_analysis_interrupted",
+      message: "Sensitivity analysis was interrupted before completion. No assessments were applied.",
     });
   }
   if (error instanceof CatalogConnectorError) {
@@ -284,32 +284,6 @@ function safeSuggestionError(reply: FastifyReply, error: unknown) {
   });
 }
 
-function untilAborted<T>(operation: Promise<T>, signal: AbortSignal): Promise<T> {
-  if (signal.aborted) {
-    void operation.catch(() => undefined);
-    return Promise.reject(new SensitivityAnalysisInterruptedError());
-  }
-  return new Promise<T>((resolve, reject) => {
-    const abort = () => reject(new SensitivityAnalysisInterruptedError());
-    signal.addEventListener("abort", abort, { once: true });
-    if (signal.aborted) {
-      void operation.catch(() => undefined);
-      abort();
-      return;
-    }
-    operation.then(
-      (value) => {
-        signal.removeEventListener("abort", abort);
-        resolve(value);
-      },
-      (error: unknown) => {
-        signal.removeEventListener("abort", abort);
-        reject(error);
-      },
-    );
-  });
-}
-
 function safeSuggestionHistoryError(reply: FastifyReply, error: unknown) {
   if (error instanceof CatalogUnavailableError) {
     return reply.code(503).send({
@@ -342,13 +316,22 @@ export function catalogDescriptionGenerationRoutes(
     try {
       const databaseId = idSchema.parse((request.params as { databaseId?: unknown }).databaseId);
       const input = suggestionSchema.parse(request.body);
-      const signal = AbortSignal.timeout(60_000);
-      const result = await untilAborted(deps.sensitivityAnalysisRunner.run(
-        databaseId,
-        input.scope,
-        "targetIds" in input ? input.targetIds : [],
-        signal,
-      ), signal);
+      const controller = new AbortController();
+      const abort = () => controller.abort();
+      request.raw.once("aborted", abort);
+      reply.raw.once("close", abort);
+      let result;
+      try {
+        result = await deps.sensitivityAnalysisRunner.run(
+          databaseId,
+          input.scope,
+          "targetIds" in input ? input.targetIds : [],
+          controller.signal,
+        );
+      } finally {
+        request.raw.off("aborted", abort);
+        reply.raw.off("close", abort);
+      }
       return {
         suggestions: result.suggestions,
         run: publicSensitivityAnalysisRun(result.run),