Poll classification results for datasets

CHANGELOG.md

@@ -21,6 +21,7 @@ The types of changes are:
 
 * Added the connection key to the execution log [#2100](https://github.com/ethyca/fides/pull/2100)
 * Added endpoints to retrieve DSR `Rule`s and `Rule Target`s [#2116](https://github.com/ethyca/fides/pull/2116)
+* Dataset classification UI now polls for results [#2123](https://github.com/ethyca/fides/pull/2123)
 
 ### Changed
 

clients/admin-ui/src/features/common/classifications/helpers.tsx

@@ -0,0 +1,48 @@
+import { useEffect, useState } from "react";
+
+import { useGetAllClassifyInstancesQuery } from "~/features/plus/plus.slice";
+import { ClassificationStatus, GenerateTypes } from "~/types/api";
+
+const POLL_INTERVAL_SECONDS = 3;
+
+/**
+ * Poll for updates to classification until all classifications are finished
+ */
+export const usePollForClassifications = ({
+  resourceType,
+  fidesKeys,
+  skip,
+}: {
+  resourceType: GenerateTypes;
+  fidesKeys?: string[];
+  skip?: boolean;
+}) => {
+  const [shouldPoll, setShouldPoll] = useState(true);
+  const result = useGetAllClassifyInstancesQuery(
+    {
+      resource_type: resourceType,
+      fides_keys: fidesKeys,
+    },
+    {
+      skip,
+      pollingInterval: shouldPoll ? POLL_INTERVAL_SECONDS * 1000 : undefined,
+    }
+  );
+
+  const isClassificationFinished = result.data
+    ? result.data.every(
+        (c) =>
+          c.status === ClassificationStatus.COMPLETE ||
+          c.status === ClassificationStatus.FAILED ||
+          c.status === ClassificationStatus.REVIEWED
+      )
+    : false;
+
+  useEffect(() => {
+    if (isClassificationFinished) {
+      setShouldPoll(false);
+    }
+  }, [isClassificationFinished]);
+
+  return { ...result, isClassificationFinished };
+};

clients/admin-ui/src/features/common/classifications/index.tsx

@@ -0,0 +1 @@
+export * from "./helpers";

clients/admin-ui/src/features/dataset/DatasetTable.tsx

@@ -1,12 +1,10 @@
 import { Table, Tbody, Td, Th, Thead, Tr } from "@fidesui/react";
 import { useDispatch, useSelector } from "react-redux";
 
+import { usePollForClassifications } from "~/features/common/classifications";
 import { useFeatures } from "~/features/common/features";
 import ClassificationStatusBadge from "~/features/plus/ClassificationStatusBadge";
-import {
-  selectDatasetClassifyInstanceMap,
-  useGetAllClassifyInstancesQuery,
-} from "~/features/plus/plus.slice";
+import { selectDatasetClassifyInstanceMap } from "~/features/plus/plus.slice";
 import { Dataset, GenerateTypes } from "~/types/api";
 
 import {
@@ -21,12 +19,10 @@ const DatasetsTable = () => {
 
   const { data: datasets } = useGetAllDatasetsQuery();
   const features = useFeatures();
-  useGetAllClassifyInstancesQuery(
-    { resource_type: GenerateTypes.DATASETS },
-    {
-      skip: !features.plus,
-    }
-  );
+  usePollForClassifications({
+    resourceType: GenerateTypes.DATASETS,
+    skip: !features.plus,
+  });
   const classifyInstanceMap = useSelector(selectDatasetClassifyInstanceMap);
 
   const handleRowClick = (dataset: Dataset) => {

clients/admin-ui/src/pages/classify-systems/index.tsx

@@ -2,25 +2,21 @@ import { Button, Heading, HStack, Spinner, Stack, Text } from "@fidesui/react";
 import type { NextPage } from "next";
 import NextLink from "next/link";
 import { useRouter } from "next/router";
-import { ReactNode, useEffect, useState } from "react";
+import { ReactNode, useEffect } from "react";
 import { useDispatch } from "react-redux";
 
 import { useAppSelector } from "~/app/hooks";
+import { usePollForClassifications } from "~/features/common/classifications";
 import { useInterzoneNav } from "~/features/common/hooks/useInterzoneNav";
 import Layout from "~/features/common/Layout";
-import {
-  useGetAllClassifyInstancesQuery,
-  useGetHealthQuery,
-} from "~/features/plus/plus.slice";
+import { useGetHealthQuery } from "~/features/plus/plus.slice";
 import {
   selectSystemsToClassify,
   setSystemsToClassify,
   useGetAllSystemsQuery,
 } from "~/features/system";
 import ClassifySystemsTable from "~/features/system/ClassifySystemsTable";
-import { ClassificationStatus, GenerateTypes } from "~/types/api";
-
-const POLL_INTERVAL_SECONDS = 3;
+import { GenerateTypes } from "~/types/api";
 
 const ClassifySystemsLayout = ({ children }: { children: ReactNode }) => (
   <Layout title="Classify Systems">
@@ -58,19 +54,15 @@ const ClassifySystems: NextPage = () => {
     }
   }, [dispatch, allSystems]);
 
-  // Poll for updates to classification until all classifications are finished
-  const [shouldPoll, setShouldPoll] = useState(true);
-  const { isLoading: isLoadingClassifications, data: classifications } =
-    useGetAllClassifyInstancesQuery(
-      {
-        resource_type: GenerateTypes.SYSTEMS,
-        fides_keys: systems?.map((s) => s.fides_key),
-      },
-      {
-        skip: !hasPlus,
-        pollingInterval: shouldPoll ? POLL_INTERVAL_SECONDS * 1000 : undefined,
-      }
-    );
+  const {
+    isLoading: isLoadingClassifications,
+    data: classifications,
+    isClassificationFinished,
+  } = usePollForClassifications({
+    resourceType: GenerateTypes.SYSTEMS,
+    fidesKeys: systems?.map((s) => s.fides_key),
+    skip: !hasPlus,
+  });
 
   useEffect(() => {
     if (!isLoadingPlus && !hasPlus) {
@@ -80,21 +72,6 @@ const ClassifySystems: NextPage = () => {
 
   const isLoading = isLoadingSystems || isLoadingClassifications;
 
-  const isClassificationFinished = classifications
-    ? classifications.every(
-        (c) =>
-          c.status === ClassificationStatus.COMPLETE ||
-          c.status === ClassificationStatus.FAILED ||
-          c.status === ClassificationStatus.REVIEWED
-      )
-    : false;
-
-  useEffect(() => {
-    if (isClassificationFinished) {
-      setShouldPoll(false);
-    }
-  }, [isClassificationFinished]);
-
   if (isLoading) {
     return (
       <ClassifySystemsLayout>