background asset wipe

js_modules/dagster-ui/packages/ui-core/src/assets/useWipeAssets.tsx

@@ -1,7 +1,13 @@
 import {useLayoutEffect, useRef, useState} from 'react';
 
-import {RefetchQueriesFunction, gql, useMutation} from '../apollo-client';
-import {AssetWipeMutation, AssetWipeMutationVariables} from './types/useWipeAssets.types';
+import {RefetchQueriesFunction, gql, useLazyQuery, useMutation} from '../apollo-client';
+import {
+  AssetWipeMutation,
+  AssetWipeMutationVariables,
+  BackgroundAssetWipeMutation,
+  BackgroundAssetWipeMutationVariables,
+  BackgroundAssetWipeStatusQuery,
+} from './types/useWipeAssets.types';
 import {showCustomAlert} from '../app/CustomAlertProvider';
 import {PYTHON_ERROR_FRAGMENT} from '../app/PythonErrorFragment';
 import {PartitionsByAssetSelector} from '../graphql/types';
@@ -21,6 +27,14 @@ export function useWipeAssets({
     ASSET_WIPE_MUTATION,
     {refetchQueries},
   );
+  const [requestBackgroundAssetWipe] = useMutation<
+    BackgroundAssetWipeMutation,
+    BackgroundAssetWipeMutationVariables
+  >(BACKGROUND_ASSET_WIPE_MUTATION, {refetchQueries});
+  const [
+    getBackgroundWipeStatus,
+    {error: backgroundWipeStatusError, data: backgroundWipeStatusData},
+  ] = useLazyQuery<BackgroundAssetWipeStatusQuery>(BACKGROUND_ASSET_WIPE_STATUS);
 
   const [isWiping, setIsWiping] = useState(false);
   const [wipedCount, setWipedCount] = useState(0);
@@ -30,6 +44,31 @@ export function useWipeAssets({
 
   const didCancel = useRef(false);
 
+  if (isWiping && backgroundWipeStatusError) {
+    setFailedCount(1);
+    onComplete?.();
+    setIsWiping(false);
+  }
+
+  if (isWiping && backgroundWipeStatusData) {
+    const data = backgroundWipeStatusData.backgroundAssetWipeStatus;
+    switch (data.__typename) {
+      case 'BackgroundAssetWipeInProgress':
+        console.log('Background asset wipe in progress');
+        break;
+      case 'BackgroundAssetWipeSuccess':
+        setWipedCount(1);
+        onComplete?.();
+        setIsWiping(false);
+        break;
+      case 'BackgroundAssetWipeFailed':
+        setFailedCount(1);
+        onComplete?.();
+        setIsWiping(false);
+        break;
+    }
+  }
+
   const wipeAssets = async (assetPartitionRanges: PartitionsByAssetSelector[]) => {
     if (!assetPartitionRanges.length) {
       return;
@@ -66,13 +105,43 @@ export function useWipeAssets({
     setIsWiping(false);
   };
 
+  const backgroundWipeAssets = async (assetPartitionRanges: PartitionsByAssetSelector[]) => {
+    if (!assetPartitionRanges.length) {
+      return;
+    }
+    setIsWiping(true);
+    const result = await requestBackgroundAssetWipe({
+      variables: {assetPartitionRanges},
+      refetchQueries,
+    });
+    const data = result.data?.backgroundWipeAssets;
+    switch (data?.__typename) {
+      case 'AssetNotFoundError':
+      case 'PythonError':
+        setFailedCount(assetPartitionRanges.length);
+        onComplete?.();
+        setIsWiping(false);
+        return;
+      case 'AssetWipeInProgress':
+        getBackgroundWipeStatus({variables: {workToken: data.workToken}, pollInterval: 1000});
+        break;
+      case 'UnauthorizedError':
+        showCustomAlert({
+          title: 'Could not wipe asset materializations',
+          body: 'You do not have permission to do this.',
+        });
+        onClose();
+        return;
+    }
+  };
+
   useLayoutEffect(() => {
     return () => {
       didCancel.current = true;
     };
   }, []);
 
-  return {wipeAssets, isWiping, isDone, wipedCount, failedCount};
+  return {backgroundWipeAssets, wipeAssets, isWiping, isDone, wipedCount, failedCount};
 }
 
 export const ASSET_WIPE_MUTATION = gql`
@@ -95,3 +164,33 @@ export const ASSET_WIPE_MUTATION = gql`
 
   ${PYTHON_ERROR_FRAGMENT}
 `;
+
+export const BACKGROUND_ASSET_WIPE_MUTATION = gql`
+  mutation BackgroundAssetWipeMutation($assetPartitionRanges: [PartitionsByAssetSelector!]!) {
+    backgroundWipeAssets(assetPartitionRanges: $assetPartitionRanges) {
+      ... on AssetWipeInProgress {
+        workToken
+      }
+      ...PythonErrorFragment
+    }
+  }
+
+  ${PYTHON_ERROR_FRAGMENT}
+`;
+
+export const BACKGROUND_ASSET_WIPE_STATUS = gql`
+  query BackgroundAssetWipeStatus($workToken: String!) {
+    backgroundAssetWipeStatus(workToken: $workToken) {
+      ... on BackgroundAssetWipeSuccess {
+        completedAt
+      }
+      ... on BackgroundAssetWipeInProgress {
+        startedAt
+      }
+      ... on BackgroundAssetWipeFailed {
+        failedAt
+        message
+      }
+    }
+  }
+`;