Merge pull request DataRecce#256 from DataRecce/feature/drc-332-change-the-behavior-of-state-loading-and-exporting

[Feature] Change the behavior of state loading and exporting

Author: wcchang1115

js/app/page.tsx

@@ -30,7 +30,7 @@ import { InfoIcon } from "@chakra-ui/icons";
 import { RunPage } from "@/components/run/RunPage";
 import { ErrorBoundary } from "@/components/errorboundary/ErrorBoundary";
 import { StateExporter } from "@/components/check/StateExporter";
-import { StateLoader } from "@/components/check/StateLoader";
+import { StateImporter } from "@/components/check/StateImporter";
 
 function getCookie(key: string) {
   var b = document.cookie.match("(^|;)\\s*" + key + "\\s*=\\s*([^;]+)");
@@ -120,8 +120,8 @@ function NavBar() {
         <Spacer />
         {!isDemoSite && (
           <>
+            <StateImporter />
             <StateExporter />
-            <StateLoader />
           </>
         )}
         <Box p="8px" mr="10px" color="gray.500">

js/src/components/check/StateExporter.tsx

@@ -1,8 +1,8 @@
 import { exportState } from "@/lib/api/state";
-import { DownloadIcon } from "@chakra-ui/icons";
-import { IconButton, Tooltip, useToast } from "@chakra-ui/react";
+import { Icon, IconButton, Tooltip, useToast } from "@chakra-ui/react";
 import { format } from "date-fns";
 import saveAs from "file-saver";
+import { TfiExport } from "react-icons/tfi";
 
 export function StateExporter() {
   const toast = useToast();
@@ -37,7 +37,7 @@ export function StateExporter() {
         variant="unstyled"
         aria-label="Export state"
         onClick={handleExport}
-        icon={<DownloadIcon />}
+        icon={<Icon as={TfiExport} boxSize={"1.2em"} />}
       />
     </Tooltip>
   );

js/src/components/check/StateLoader.tsx js/src/components/check/StateImporter.tsx

@@ -18,12 +18,12 @@ import {
   useToast,
 } from "@chakra-ui/react";
 import { InfoIcon } from "@chakra-ui/icons";
-import { loadState } from "@/lib/api/state";
-import { IoFolderOpenOutline } from "react-icons/io5";
+import { importState } from "@/lib/api/state";
 import { useLocation } from "wouter";
 import { useRunsAggregated } from "@/lib/hooks/LineageGraphContext";
+import { TfiImport } from "react-icons/tfi";
 
-export function StateLoader() {
+export function StateImporter() {
   const toast = useToast();
   const queryClient = useQueryClient();
   const hiddenFileInput = useRef<HTMLInputElement>(null);
@@ -33,30 +33,30 @@ export function StateLoader() {
   const [location, setLocation] = useLocation();
   const [, refetchRunsAggregated] = useRunsAggregated();
 
-  const handleLoad = useCallback(async () => {
+  const handleImport = useCallback(async () => {
     if (!selectedFile) {
       return;
     }
 
     try {
-      const { runs, checks } = await loadState(selectedFile);
+      const { runs, checks } = await importState(selectedFile);
       refetchRunsAggregated();
       await queryClient.invalidateQueries({ queryKey: cacheKeys.checks() });
       if (location.includes("/checks")) {
         setLocation("/checks");
       }
       toast({
-        description: `${runs} runs and ${checks} checks loaded successfully`,
+        description: `${runs} runs and ${checks} checks imported successfully`,
         status: "info",
         variant: "left-accent",
         position: "bottom",
         duration: 5000,
         isClosable: true,
       });
     } catch (error) {
-      console.error("Load failed", error);
+      console.error("Import failed", error);
       toast({
-        title: "Load failed",
+        title: "Import failed",
         description: `${error}`,
         status: "error",
         variant: "left-accent",
@@ -92,12 +92,12 @@ export function StateLoader() {
 
   return (
     <>
-      <Tooltip label="Load">
+      <Tooltip label="Import">
         <IconButton
           variant="unstyled"
-          aria-label="Load state"
+          aria-label="Import state"
           onClick={handleClick}
-          icon={<Icon pt="10px" as={IoFolderOpenOutline} boxSize={"2em"} />}
+          icon={<Icon as={TfiImport} boxSize={"1.2em"} />}
         />
       </Tooltip>
       <input
@@ -110,12 +110,12 @@ export function StateLoader() {
         isOpen={isOpen}
         leastDestructiveRef={cancelRef}
         onClose={onClose}
-        size={"lg"}
+        size={"xl"}
       >
         <AlertDialogOverlay>
           <AlertDialogContent>
             <AlertDialogHeader fontSize="lg" fontWeight="bold">
-              Load state
+              Import state
             </AlertDialogHeader>
 
             <AlertDialogBody>
@@ -128,11 +128,11 @@ export function StateLoader() {
                 </Flex>
                 <Flex>
                   <Text>
-                    All runs and checks will be{" "}
+                    The current runs and checks will be{" "}
                     <Text as="span" fontWeight="600">
-                      overwritten
+                      merged
                     </Text>{" "}
-                    by the loaded state
+                    with the imported state
                   </Text>
                 </Flex>
               </Flex>
@@ -142,8 +142,8 @@ export function StateLoader() {
               <Button ref={cancelRef} onClick={onClose}>
                 Cancel
               </Button>
-              <Button colorScheme="blue" onClick={handleLoad} ml="5px">
-                Load
+              <Button colorScheme="blue" onClick={handleImport} ml="5px">
+                Import
               </Button>
             </AlertDialogFooter>
           </AlertDialogContent>

js/src/lib/api/state.ts

@@ -1,6 +1,6 @@
 import { axiosClient } from "./axiosClient";
 
-export interface LoadedState {
+export interface ImportedState {
   runs: number;
   checks: number;
 }
@@ -10,10 +10,10 @@ export async function exportState(): Promise<string> {
   return response.data;
 }
 
-export async function loadState(file: File): Promise<LoadedState> {
+export async function importState(file: File): Promise<ImportedState> {
   const formData = new FormData();
   formData.append("file", file);
 
-  const response = await axiosClient.post("/api/load", formData);
+  const response = await axiosClient.post("/api/import", formData);
   return response.data;
 }

recce/models/state.py

@@ -62,6 +62,33 @@ def load(cls, file_path):
 
         return state
 
+    def export_state(self):
+        self.metadata = RecceStateMetadata()
+        return pydantic_model_json_dump(self)
+
+    def import_state(self, content):
+        import_state = RecceState.model_validate_json(content)
+        current_check_ids = [str(c.check_id) for c in self.checks]
+        current_run_ids = [str(r.run_id) for r in self.runs]
+
+        # merge checks
+        import_checks = 0
+        for check in import_state.checks:
+            if str(check.check_id) not in current_check_ids:
+                self.checks.append(check)
+                import_checks += 1
+
+        # merge runs
+        import_runs = 0
+        for run in import_state.runs:
+            if str(run.run_id) not in current_run_ids:
+                self.runs.append(run)
+                import_runs += 1
+
+        self.runs.sort(key=lambda x: x.run_at)
+
+        return import_runs, import_checks
+
 
 recce_state: Optional[RecceState] = None
 

recce/server.py

@@ -20,8 +20,9 @@
 from .apis.run_api import run_router
 from .dbt import load_dbt_context, default_dbt_context
 from .exceptions import RecceException
+from .models.lineage import LineageDAO
 from .models.state import load_default_state, default_recce_state, RecceState
-from .models.util import pydantic_model_json_dump
+from .models.types import Lineage
 
 logger = logging.getLogger('uvicorn')
 
@@ -124,21 +125,27 @@ async def get_lineage(base: Optional[bool] = False):
 @app.post("/api/export", response_class=PlainTextResponse, status_code=200)
 async def export_handler():
     try:
-        return pydantic_model_json_dump(default_recce_state())
+        ctx = default_dbt_context()
+        base = ctx.get_lineage(base=True)
+        current = ctx.get_lineage(base=False)
+        lineage = Lineage(
+            base=base,
+            current=current,
+        )
+        LineageDAO().set(lineage)
+
+        return default_recce_state().export_state()
     except RecceException as e:
         raise HTTPException(status_code=400, detail=e.message)
 
 
-@app.post("/api/load", status_code=200)
-async def load_handler(file: UploadFile):
+@app.post("/api/import", status_code=200)
+async def import_handler(file: UploadFile):
     try:
         content = await file.read()
-        load_state = RecceState().model_validate_json(content)
-        current_state = default_recce_state()
-        current_state.checks = load_state.checks
-        current_state.runs = load_state.runs
+        import_runs, import_checks = default_recce_state().import_state(content)
 
-        return {"runs": len(current_state.runs), "checks": len(current_state.checks)}
+        return {"runs": import_runs, "checks": import_checks}
     except ValidationError as e:
         raise HTTPException(status_code=400, detail=str(e))
     except RecceException as e: