Refactor version fetching and refetch

coder · BrunoQuaresma · Nov 3, 2023 · Oct 31, 2023 · Oct 31, 2023 · Oct 31, 2023
commit e7713feacd23eca4705f530214a91a890a7dd5fb
diff --git a/site/src/pages/TemplateVersionEditorPage/TemplateVersionEditorPage.tsx b/site/src/pages/TemplateVersionEditorPage/TemplateVersionEditorPage.tsx
@@ -6,7 +6,7 @@ import { Helmet } from "react-helmet-async";
 import { useNavigate, useParams } from "react-router-dom";
 import { pageTitle } from "utils/page";
 import { templateVersionEditorMachine } from "xServices/templateVersionEditor/templateVersionEditorXService";
-import { useMutation, useQuery } from "react-query";
+import { useMutation, useQuery, useQueryClient } from "react-query";
 import {
   createTemplateVersion,
   templateByName,
@@ -20,22 +20,30 @@ import {
   isAllowedFile,
 } from "utils/templateVersion";
 import { watchBuildLogsByTemplateVersionId } from "api/api";
-import { ProvisionerJobLog, TemplateVersion } from "api/typesGenerated";
+import { ProvisionerJobLog } from "api/typesGenerated";
 
 type Params = {
   version: string;
   template: string;
 };
 
 export const TemplateVersionEditorPage: FC = () => {
+  const queryClient = useQueryClient();
   const navigate = useNavigate();
   const { version: versionName, template: templateName } =
     useParams() as Params;
   const orgId = useOrganizationId();
+  const [currentVersionName, setCurrentVersionName] = useState(versionName);
   const templateQuery = useQuery(templateByName(orgId, templateName));
-  const templateVersionQuery = useQuery(
-    templateVersionByName(orgId, templateName, versionName),
+  const templateVersionOptions = templateVersionByName(
+    orgId,
+    templateName,
+    currentVersionName,
   );
+  const templateVersionQuery = useQuery({
+    ...templateVersionOptions,
+    keepPreviousData: true,
+  });
   const fileQuery = useQuery({
     ...file(templateVersionQuery.data?.job.file_id ?? ""),
     enabled: templateVersionQuery.isSuccess,
@@ -49,8 +57,6 @@ export const TemplateVersionEditorPage: FC = () => {
   const createTemplateVersionMutation = useMutation(
     createTemplateVersion(orgId),
   );
-  const [currentVersionOnEditor, setCurrentVersionOnEditor] =
-    useState<TemplateVersion>();
 
   // Initialize file tree
   useEffect(() => {
@@ -69,50 +75,47 @@ export const TemplateVersionEditorPage: FC = () => {
     }
   }, [fileQuery.data, sendEvent]);
 
-  // Initialize current version used on editor
-  useEffect(() => {
-    if (templateVersionQuery.data) {
-      setCurrentVersionOnEditor(templateVersionQuery.data);
-    }
-  }, [templateVersionQuery.data]);
-
   // Watch version logs
   const [logs, setLogs] = useState<ProvisionerJobLog[]>([]);
+  const templateVersionId = templateVersionQuery.data?.id;
+  const refetchTemplateVersion = templateVersionQuery.refetch;
+  const templateVersionStatus = templateVersionQuery.data?.job.status;
   useEffect(() => {
-    if (!currentVersionOnEditor) {
+    if (!templateVersionId || !templateVersionStatus) {
+      return;
+    }
+
+    if (templateVersionStatus !== "running") {
       return;
     }
 
-    const socket = watchBuildLogsByTemplateVersionId(
-      currentVersionOnEditor.id,
-      {
-        onMessage: (log) => {
-          setLogs((logs) => [...logs, log]);
-        },
-        onDone: () => {
-          sendEvent({ type: "BUILD_DONE" });
-        },
-        onError: (error) => {
-          console.error(error);
-        },
+    const socket = watchBuildLogsByTemplateVersionId(templateVersionId, {
+      onMessage: (log) => {
+        setLogs((logs) => [...logs, log]);
+      },
+      onDone: async () => {
+        await refetchTemplateVersion();
       },
-    );
+      onError: (error) => {
+        console.error(error);
+      },
+    });
 
     return () => {
       socket.close();
     };
-  }, [currentVersionOnEditor, sendEvent]);
+  }, [refetchTemplateVersion, templateVersionId, templateVersionStatus]);
 
   return (
     <>
       <Helmet>
         <title>{pageTitle(`${templateName} · Template Editor`)}</title>
       </Helmet>
 
-      {templateQuery.data && currentVersionOnEditor && fileTree && (
+      {templateQuery.data && templateVersionQuery.data && fileTree && (
         <TemplateVersionEditor
           template={templateQuery.data}
-          templateVersion={currentVersionOnEditor}
+          templateVersion={templateVersionQuery.data}
           isBuildingNewVersion={Boolean(editorState.context.version)}
           defaultFileTree={fileTree}
           onPreview={async (newFileTree) => {
@@ -134,7 +137,11 @@ export const TemplateVersionEditorPage: FC = () => {
               file_id: serverFile.hash,
             });
             setLogs([]);
-            setCurrentVersionOnEditor(newVersion);
+            setCurrentVersionName(newVersion.name);
+            queryClient.setQueryData(
+              templateVersionOptions.queryKey,
+              newVersion,
+            );
             sendEvent({
               type: "CREATED_VERSION",
               data: newVersion,

diff --git a/site/src/xServices/templateVersionEditor/templateVersionEditorXService.ts b/site/src/xServices/templateVersionEditor/templateVersionEditorXService.ts
@@ -37,7 +37,6 @@ export const templateVersionEditorMachine = createMachine(
             fileId: string;
           }
         | { type: "CANCEL_MISSING_VARIABLE_VALUES" }
-        | { type: "BUILD_DONE" }
         | { type: "PUBLISH" }
         | ({ type: "CONFIRM_PUBLISH" } & PublishVersionData)
         | { type: "CANCEL_PUBLISH" },
@@ -67,7 +66,7 @@ export const templateVersionEditorMachine = createMachine(
         on: {
           CREATED_VERSION: {
             actions: ["resetCreateBuildData", "assignBuild"],
-            target: "watchingBuildLogs",
+            target: "fetchingVersion",
           },
           PUBLISH: {
             target: "askPublishParameters",
@@ -107,18 +106,11 @@ export const templateVersionEditorMachine = createMachine(
           src: "createBuild",
           onDone: {
             actions: "assignBuild",
-            target: "watchingBuildLogs",
+            target: "fetchingVersion",
           },
         },
       },
 
-      watchingBuildLogs: {
-        tags: "loading",
-        on: {
-          BUILD_DONE: "fetchingVersion",
-        },
-      },
-
       fetchingVersion: {
         tags: "loading",
         invoke: {