background asset wipe

dagster-io · Nov 26, 2024 · e5924b0 · e5924b0
1 parent 323d5a3
commit e5924b0
Show file tree

Hide file tree

Showing 13 changed files with 515 additions and 30 deletions.
diff --git a/js_modules/dagster-ui/packages/ui-core/client.json b/js_modules/dagster-ui/packages/ui-core/client.json
diff --git a/js_modules/dagster-ui/packages/ui-core/src/assets/AssetWipeDialog.oss.tsx b/js_modules/dagster-ui/packages/ui-core/src/assets/AssetWipeDialog.oss.tsx
@@ -53,7 +53,7 @@ export const AssetWipeDialogInner = memo(
     onComplete?: () => void;
     requery?: RefetchQueriesFunction;
   }) => {
-    const {wipeAssets, isWiping, isDone, wipedCount, failedCount} = useWipeAssets({
+    const {wipeAssets, backgroundWipeAssets, isWiping, isDone, wipedCount, failedCount} = useWipeAssets({
       refetchQueries: requery,
       onClose,
       onComplete,
@@ -104,7 +104,7 @@ export const AssetWipeDialogInner = memo(
           {isDone ? null : (
             <Button
               intent="danger"
-              onClick={() => wipeAssets(assetKeys.map((key) => asAssetPartitionRangeInput(key)))}
+              onClick={() => backgroundWipeAssets(assetKeys.map((key) => asAssetPartitionRangeInput(key)))}
               disabled={isWiping}
               loading={isWiping}
             >

diff --git a/js_modules/dagster-ui/packages/ui-core/src/assets/types/useWipeAssets.types.ts b/js_modules/dagster-ui/packages/ui-core/src/assets/types/useWipeAssets.types.ts
diff --git a/js_modules/dagster-ui/packages/ui-core/src/assets/useWipeAssets.tsx b/js_modules/dagster-ui/packages/ui-core/src/assets/useWipeAssets.tsx
@@ -1,7 +1,7 @@
 import {useLayoutEffect, useRef, useState} from 'react';
 
-import {RefetchQueriesFunction, gql, useMutation} from '../apollo-client';
-import {AssetWipeMutation, AssetWipeMutationVariables} from './types/useWipeAssets.types';
+import {RefetchQueriesFunction, gql, useMutation, useLazyQuery} 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 +21,11 @@ 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 +35,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 +96,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 +155,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
+      }
+    }
+  }
+`;