Merge pull request #470 from DataRecce/feature/drc-781-bug-should-upd…

…ate-to-the-latest-run-after-rerun

[Bug] Should update to the latest run after rerun a check

js/src/components/check/CheckDetail.tsx

@@ -1,9 +1,4 @@
 import {
-  Accordion,
-  AccordionButton,
-  AccordionIcon,
-  AccordionItem,
-  AccordionPanel,
   Box,
   Button,
   Center,
@@ -55,16 +50,17 @@ import { stripIndents } from "common-tags";
 import { useClipBoardToast } from "@/lib/hooks/useClipBoardToast";
 import { buildTitle, buildDescription, buildQuery } from "./check";
 import SqlEditor, { DualSqlEditor } from "../query/SqlEditor";
-import { useCallback, useEffect, useState } from "react";
-import { cancelRun, submitRunFromCheck, waitRun } from "@/lib/api/runs";
+import { useCallback, useState } from "react";
+import { cancelRun, submitRunFromCheck } from "@/lib/api/runs";
 import { Run } from "@/lib/api/types";
 import { RunView } from "../run/RunView";
-import { formatDistanceToNow } from "date-fns";
+import { formatDistanceToNow, sub } from "date-fns";
 import { LineageDiffView } from "./LineageDiffView";
 import { findByRunType } from "../run/registry";
 import { PresetCheckTemplateView } from "./PresetCheckTemplateView";
 import { VSplit } from "../split/Split";
 import { useCopyToClipboardButton } from "@/lib/hooks/ScreenShot";
+import { useRun } from "@/lib/hooks/useRun";
 
 interface CheckDetailProps {
   checkId: string;
@@ -74,9 +70,9 @@ export const CheckDetail = ({ checkId }: CheckDetailProps) => {
   const queryClient = useQueryClient();
   const [, setLocation] = useLocation();
   const { successToast, failToast } = useClipBoardToast();
-  const [runId, setRunId] = useState<string>();
+  const [submittedRunId, setSubmittedRunId] = useState<string>();
   const [progress, setProgress] = useState<Run["progress"]>();
-  const [abort, setAborting] = useState(false);
+  const [isAborting, setAborting] = useState(false);
   const {
     isOpen: isPresetCheckTemplateOpen,
     onOpen: onPresetCheckTemplateOpen,
@@ -90,15 +86,19 @@ export const CheckDetail = ({ checkId }: CheckDetailProps) => {
   const {
     isLoading,
     error,
-    refetch,
     data: check,
   } = useQuery({
     queryKey: cacheKeys.check(checkId),
     queryFn: async () => getCheck(checkId),
-    refetchOnMount: false,
-    staleTime: 5 * 60 * 1000,
+    refetchOnMount: true,
   });
 
+  const trackedRunId = submittedRunId || check?.last_run?.run_id;
+  const { run, error: rerunError } = useRun(trackedRunId);
+  const isRunning = submittedRunId
+    ? !run || run.status === "running"
+    : run?.status === "running";
+
   const runTypeEntry = check?.type ? findByRunType(check?.type) : undefined;
   const isPresetCheck = check?.is_preset || false;
 
@@ -118,52 +118,24 @@ export const CheckDetail = ({ checkId }: CheckDetailProps) => {
     },
   });
 
-  const submitRunFn = async () => {
+  const handleRerun = useCallback(async () => {
     const type = check?.type;
     if (!type) {
       return;
     }
 
-    const { run_id } = await submitRunFromCheck(checkId, { nowait: true });
-
-    setRunId(run_id);
-
-    while (true) {
-      const run = await waitRun(run_id, 2);
-      setProgress(run.progress);
-      if (run.result || run.error) {
-        setAborting(false);
-        setProgress(undefined);
-        return run;
-      }
-    }
-  };
-
-  const {
-    data: rerunRun,
-    mutate: rerun,
-    error: rerunError,
-    isIdle: rerunIdle,
-    isPending: rerunPending,
-  } = useMutation({
-    mutationFn: submitRunFn,
-    onSuccess: (run) => {
-      refetch();
-    },
-  });
-
-  const handleRerun = async () => {
-    rerun();
-  };
+    const submittedRun = await submitRunFromCheck(checkId, { nowait: true });
+    setSubmittedRunId(submittedRun.run_id);
+  }, [check, checkId, setSubmittedRunId]);
 
   const handleCancel = useCallback(async () => {
     setAborting(true);
-    if (!runId) {
+    if (!trackedRunId) {
       return;
     }
 
-    return await cancelRun(runId);
-  }, [runId]);
+    return await cancelRun(trackedRunId);
+  }, [trackedRunId]);
 
   const handleCopy = async () => {
     if (!check) {
@@ -213,7 +185,6 @@ export const CheckDetail = ({ checkId }: CheckDetailProps) => {
     return <Center h="100%">Error: {error.message}</Center>;
   }
 
-  const run = rerunIdle ? check?.last_run : rerunRun;
   const relativeTime = run?.run_at
     ? formatDistanceToNow(new Date(run.run_at), { addSuffix: true })
     : null;
@@ -284,7 +255,7 @@ export const CheckDetail = ({ checkId }: CheckDetailProps) => {
             <Tooltip label="Rerun">
               <IconButton
                 isRound={true}
-                isLoading={rerunPending}
+                isLoading={isRunning}
                 variant="ghost"
                 aria-label="Rerun"
                 icon={<RepeatIcon />}
@@ -352,22 +323,28 @@ export const CheckDetail = ({ checkId }: CheckDetailProps) => {
           </TabList>
           <TabPanels height="100%" flex="1" style={{ contain: "strict" }}>
             <TabPanel p={0} width="100%" height="100%">
-              {runTypeEntry?.RunResultView && (
-                <RunView
-                  ref={ref}
-                  isPending={rerunPending}
-                  isAborting={abort}
-                  isCheckDetail={true}
-                  run={run}
-                  error={rerunError}
-                  progress={progress}
-                  RunResultView={runTypeEntry.RunResultView}
-                  viewOptions={check?.view_options}
-                  onViewOptionsChanged={handelUpdateViewOptions}
-                  onCancel={handleCancel}
-                  onExecuteRun={handleRerun}
-                />
-              )}
+              {runTypeEntry?.RunResultView &&
+                (check?.last_run || trackedRunId ? (
+                  <RunView
+                    ref={ref}
+                    isRunning={isRunning}
+                    isAborting={isAborting}
+                    run={trackedRunId ? run : check?.last_run}
+                    error={rerunError}
+                    progress={progress}
+                    RunResultView={runTypeEntry.RunResultView}
+                    viewOptions={check?.view_options}
+                    onViewOptionsChanged={handelUpdateViewOptions}
+                    onCancel={handleCancel}
+                    onExecuteRun={handleRerun}
+                  />
+                ) : (
+                  <Center bg="rgb(249,249,249)" height="100%">
+                    <Button onClick={handleRerun} colorScheme="blue" size="sm">
+                      Run Query
+                    </Button>
+                  </Center>
+                ))}
               {check && check.type === "schema_diff" && (
                 <SchemaDiffView check={check} />
               )}

js/src/components/run/RunResultPane.tsx

@@ -59,21 +59,7 @@ export const _LoadableRunView = ({
   onClose?: () => void;
 }) => {
   const { runAction } = useRecceActionContext();
-
-  const { isPending, error, run, onCancel } = useRun(runId);
-  const isPendingRef = useRef(false);
-  isPendingRef.current = isPending;
-
-  useEffect(() => {
-    return () => {
-      if (isPendingRef.current) {
-        onCancel();
-      }
-    };
-
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [isPendingRef]);
-
+  const { error, run, onCancel, isRunning } = useRun(runId);
   const [viewOptions, setViewOptions] = useState();
   const queryClient = useQueryClient();
   const [, setLocation] = useLocation();
@@ -133,7 +119,7 @@ export const _LoadableRunView = ({
             <Button
               leftIcon={<RepeatIcon />}
               variant="outline"
-              isDisabled={!runId || isPending}
+              isDisabled={!runId || isRunning}
               size="sm"
               onClick={handleRerun}
             >
@@ -187,7 +173,6 @@ export const _LoadableRunView = ({
       {tabIndex === 0 && (
         <RunView
           ref={ref}
-          isPending={isPending}
           error={error}
           run={run}
           onCancel={onCancel}

js/src/components/run/RunView.tsx

@@ -16,7 +16,7 @@ import { ErrorBoundary } from "@sentry/react";
 import { forwardRef } from "@chakra-ui/react";
 
 interface RunViewProps<PT, RT, VO = any> {
-  isPending?: boolean;
+  isRunning?: boolean;
   run?: Run<PT, RT>;
   error?: Error | null;
   progress?: Run["progress"];
@@ -35,9 +35,8 @@ interface RunViewProps<PT, RT, VO = any> {
 export const RunView = forwardRef(
   <PT, RT>(
     {
-      isPending,
+      isRunning,
       isAborting,
-      isCheckDetail,
       progress,
       error,
       run,
@@ -61,7 +60,7 @@ export const RunView = forwardRef(
       );
     }
 
-    if (isPending) {
+    if (isRunning !== undefined ? isRunning : run?.status === "running") {
       let loadingMessage = progress?.message
         ? progress?.message
         : run?.progress?.message
@@ -93,19 +92,12 @@ export const RunView = forwardRef(
           </VStack>
         </Center>
       );
-    } else if (!run) {
-      if (isCheckDetail) {
-        return (
-          <Center bg="rgb(249,249,249)" height="100%">
-            <Button onClick={onExecuteRun} colorScheme="blue" size="sm">
-              Run Query
-            </Button>
-          </Center>
-        );
-      }
+    }
+
+    if (!run) {
       return (
         <Center bg="rgb(249,249,249)" height="100%">
-          No Data
+          Loading...
         </Center>
       );
     }

js/src/lib/hooks/useRun.tsx

@@ -9,39 +9,41 @@ import { useRunsAggregated } from "./LineageGraphContext";
 interface UseRunResult {
   run?: Run;
   aborting: boolean;
-  isPending: boolean;
+  isRunning: boolean;
   error: Error | null;
   onCancel: () => void;
   RunResultView?: React.ComponentType<any>;
 }
 
 export const useRun = (runId?: string): UseRunResult => {
-  const [isPolling, setIsPolling] = useState(false);
+  const [isRunning, setIsRunning] = useState(false);
   const [aborting, setAborting] = useState(false);
   const [, refetchRunsAggregated] = useRunsAggregated();
 
   const { error, data: run } = useQuery({
     queryKey: cacheKeys.run(runId || ""),
     queryFn: async () => {
-      return waitRun(runId || "", 2);
+      return waitRun(runId || "", isRunning ? 2 : 0);
     },
     enabled: !!runId,
-    refetchInterval: isPolling ? 50 : false,
+    refetchInterval: isRunning ? 50 : false,
     retry: false,
   });
 
   useEffect(() => {
     if (error || run?.result || run?.error) {
-      if (isPolling) {
-        setIsPolling(false);
-        if (run?.type === "row_count_diff") {
-          refetchRunsAggregated();
-        }
+      if (isRunning) {
+        setIsRunning(false);
       }
-    } else {
-      setIsPolling(true);
+      if (run?.type === "row_count_diff") {
+        refetchRunsAggregated();
+      }
+    }
+
+    if (run?.status === "running") {
+      setIsRunning(true);
     }
-  }, [run, error, isPolling, refetchRunsAggregated]);
+  }, [run, error, isRunning, refetchRunsAggregated]);
 
   const onCancel = useCallback(async () => {
     setAborting(true);
@@ -58,7 +60,7 @@ export const useRun = (runId?: string): UseRunResult => {
 
   return {
     run,
-    isPending: isPolling,
+    isRunning,
     aborting,
     error,
     onCancel,

recce/apis/check_api.py

@@ -120,8 +120,6 @@ async def get_check_handler(check_id: UUID):
         raise HTTPException(status_code=404, detail='Not Found')
 
     runs = RunDAO().list_by_check_id(check_id)
-    # only get the last with successful result
-    runs = [run for run in runs if run.result is not None]
     last_run = runs[-1] if len(runs) > 0 else None
 
     out = CheckOut.from_check(check)