fix: indicate docs and tests edit changes (#1557)

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Co-authored-by: Michiel De Smet <[email protected]>

Author: 3 people

webview_panels/src/lib/altimate/main.css

@@ -1695,6 +1695,7 @@
   --gray--gray-02: #3a3d41;
   --gray--blue--gray-01: #082247;
   --gray--blue--gray-02: #374c6a;
+  --background-orange: #32240c;
 }
 .altimate-component.vscode-dark {
   --primary-color: #247efe;
@@ -1730,6 +1731,7 @@
   --gray--gray-02: #3a3d41;
   --gray--blue--gray-01: #082247;
   --gray--blue--gray-02: #374c6a;
+  --background-orange: #32240c;
 }
 .altimate-component .card {
   padding: 16px;

webview_panels/src/main.scss

@@ -43,6 +43,7 @@ body.vscode-high-contrast-light {
   --gray--gray-02: #3a3d41;
   --gray--blue--gray-01: #082247;
   --gray--blue--gray-02: #374c6a;
+  --background-orange: #32240c;
 }
 
 body.vscode-dark,
@@ -80,6 +81,7 @@ body.vscode-high-contrast:not(.vscode-high-contrast-light) {
   --gray--gray-02: #3a3d41;
   --gray--blue--gray-01: #082247;
   --gray--blue--gray-02: #374c6a;
+  --background-orange: #32240c;
 }
 
 body {

webview_panels/src/modules/documentationEditor/DocumentationEditor.tsx

@@ -17,7 +17,6 @@ import useDocumentationContext from "./state/useDocumentationContext";
 import classes from "./styles.module.scss";
 import { addDefaultActions } from "./utils";
 import ConversationsRightPanel from "./components/conversation/ConversationsRightPanel";
-import useIncomingDocsDataHandler from "./useIncomingDocsDataHandler";
 import CoachAiIfModified from "./components/docGenerator/CoachAiIfModified";
 import Citations from "./components/docGenerator/Citations";
 import { Citation } from "@lib";
@@ -29,7 +28,6 @@ const DocumentationEditor = (): JSX.Element => {
     dispatch,
   } = useDocumentationContext();
   const { postMessageToDataPilot } = useAppContext();
-  useIncomingDocsDataHandler();
 
   const modelTests = useMemo(() => {
     return currentDocsTests?.filter((test) => !test.column_name);
@@ -130,6 +128,7 @@ const DocumentationEditor = (): JSX.Element => {
                   onSubmit={onModelDocSubmit}
                   placeholder="Describe your model"
                   title={`Model: ${currentDocsData.name}`}
+                  tests={modelTests}
                 />
                 <EntityWithTests
                   title={currentDocsData.name}

webview_panels/src/modules/documentationEditor/DocumentationProvider.tsx

@@ -1,4 +1,7 @@
-import { executeRequestInAsync } from "@modules/app/requestExecutor";
+import {
+  executeRequestInAsync,
+  executeRequestInSync,
+} from "@modules/app/requestExecutor";
 import { IncomingMessageProps } from "@modules/app/types";
 import { panelLogger } from "@modules/logger";
 import {
@@ -7,6 +10,7 @@ import {
   useEffect,
   useMemo,
   useReducer,
+  useRef,
 } from "react";
 import documentationSlice, {
   initialState,
@@ -21,6 +25,7 @@ import documentationSlice, {
   updateColumnsInCurrentDocsData,
   updateConversationsRightPanelState,
   updateCurrentDocsData,
+  updateCurrentDocsTests,
   updateSelectedConversationGroup,
   updateUserInstructions,
 } from "./state/documentationSlice";
@@ -31,7 +36,7 @@ import {
   MetadataColumn,
 } from "./state/types";
 import { ContextProps } from "./types";
-import { getGenerationsInModel } from "./utils";
+import { getGenerationsInModel, isStateDirty } from "./utils";
 import DocumentationEditor from "./DocumentationEditor";
 import { ConversationGroup, DbtDocsShareDetails } from "@lib";
 import { TelemetryEvents } from "@telemetryEvents";
@@ -43,14 +48,38 @@ export const DocumentationContext = createContext<ContextProps>({
   dispatch: () => null,
 });
 
+type IncomingMessageEvent = MessageEvent<
+  IncomingMessageProps & {
+    docs?: DBTDocumentation;
+    tests?: DBTModelTest[];
+    project?: string;
+    columns?: DBTDocumentation["columns"];
+    model?: string;
+    name?: string;
+    description?: string;
+    collaborationEnabled?: boolean;
+    missingDocumentationMessage?: {
+      message: string;
+      type: "error" | "warning";
+    };
+  }
+>;
+
+enum ActionState {
+  CANCEL_STAY = "Stay",
+  DISCARD_PROCEED = "Discard",
+  SAVE_PROCEED = "Save changes",
+}
+
 const DocumentationProvider = (): JSX.Element => {
   const {
     state: { isComponentsApiInitialized },
   } = useAppContext();
   const [state, dispatch] = useReducer(
     documentationSlice.reducer,
-    documentationSlice.getInitialState()
+    documentationSlice.getInitialState(),
   );
+  const stateRef = useRef(state);
 
   const updateFocus = (name?: string) => {
     dispatch(setInsertedEntityName(name));
@@ -84,107 +113,147 @@ const DocumentationProvider = (): JSX.Element => {
       updateSelectedConversationGroup({
         shareId,
         conversationGroupId: conversation_group_id,
-      })
+      }),
     );
   };
 
-  const onMessage = useCallback(
-    (
-      event: MessageEvent<
-        IncomingMessageProps & {
-          docs?: DBTDocumentation;
-          tests?: DBTModelTest[];
-          project?: string;
-          columns?: DBTDocumentation["columns"];
-          model?: string;
-          name?: string;
-          description?: string;
-          collaborationEnabled?: boolean;
-          missingDocumentationMessage?: {
-            message: string;
-            type: "error" | "warning";
-          };
-        }
-      >
-    ) => {
-      const { command, ...params } = event.data;
-      switch (command) {
-        case "viewConversation":
-          handleViewConversation(
-            params as unknown as Parameters<typeof handleViewConversation>["0"]
-          );
+  const renderDocumentation = (event: IncomingMessageEvent) => {
+    dispatch(
+      setIncomingDocsData({
+        docs: event.data.docs,
+        tests: event.data.tests,
+      }),
+    );
+    dispatch(setProject(event.data.project));
+    dispatch(
+      updateCollaborationEnabled(Boolean(event.data.collaborationEnabled)),
+    );
+    dispatch(
+      setMissingDocumentationMessage(event.data.missingDocumentationMessage),
+    );
+  };
+
+  const onMessage = useCallback((event: IncomingMessageEvent) => {
+    const { command, ...params } = event.data;
+    switch (command) {
+      case "viewConversation":
+        handleViewConversation(
+          params as unknown as Parameters<typeof handleViewConversation>["0"],
+        );
+        break;
+      case "conversations:updates":
+        handleConversationUpdates(
+          params as unknown as Parameters<
+            typeof handleConversationUpdates
+          >["0"],
+        );
+        break;
+      case "renderDocumentation": {
+        if (
+          event.data.docs?.uniqueId ===
+          stateRef.current.currentDocsData?.uniqueId
+        ) {
           break;
-        case "conversations:updates":
-          handleConversationUpdates(
-            params as unknown as Parameters<
-              typeof handleConversationUpdates
-            >["0"]
-          );
+        }
+        if (!isStateDirty(stateRef.current)) {
+          renderDocumentation(event);
           break;
-        case "renderDocumentation":
-          dispatch(
-            setIncomingDocsData({
-              docs: event.data.docs,
-              tests: event.data.tests,
-            })
-          );
-          dispatch(setProject(event.data.project));
-          dispatch(
-            updateCollaborationEnabled(Boolean(event.data.collaborationEnabled))
-          );
+        }
+        const { currentDocsData, currentDocsTests } = stateRef.current;
+        executeRequestInSync("showWarningMessage", {
+          infoMessage: `You have unsaved changes in model: ‘${currentDocsData?.name}’. Would you
+          like to discard the changes, save them and proceed, or remain in the
+          current state?`,
+          items: [
+            ActionState.DISCARD_PROCEED,
+            ActionState.CANCEL_STAY,
+            ActionState.SAVE_PROCEED,
+          ],
+        })
+          .then(async (action) => {
+            switch (action) {
+              case ActionState.SAVE_PROCEED: {
+                const result = (await executeRequestInSync(
+                  "saveDocumentation",
+                  {
+                    ...currentDocsData,
+                    updatedTests: currentDocsTests,
+                    dialogType: "Existing file",
+                  },
+                )) as { saved: boolean };
+                if (result.saved) {
+                  dispatch(updateCurrentDocsData(event.data.docs));
+                  dispatch(updateCurrentDocsTests(event.data.tests));
+                }
+                renderDocumentation(event);
+                break;
+              }
+              case ActionState.DISCARD_PROCEED: {
+                dispatch(updateCurrentDocsData(event.data.docs));
+                dispatch(updateCurrentDocsTests(event.data.tests));
+                renderDocumentation(event);
+                break;
+              }
+              case ActionState.CANCEL_STAY: {
+                break;
+              }
+              default:
+                break;
+            }
+          })
+          .catch((err) => {
+            panelLogger.error(
+              "error while showing unsaved changes dialog",
+              err,
+            );
+          });
+        break;
+      }
+      case "renderColumnsFromMetadataFetch":
+        if (event.data.columns) {
           dispatch(
-            setMissingDocumentationMessage(
-              event.data.missingDocumentationMessage
-            )
+            updateColumnsAfterSync({
+              columns: event.data.columns,
+            }),
           );
-          break;
-        case "renderColumnsFromMetadataFetch":
-          if (event.data.columns) {
-            dispatch(
-              updateColumnsAfterSync({
-                columns: event.data.columns,
-              })
-            );
-          }
-          break;
-        case "docgen:insert":
-          panelLogger.info("received new doc gen", event.data);
-          // insert model desc
-          if (params.model) {
-            dispatch(
-              updateCurrentDocsData({
-                description: params.description,
-                name: params.model,
-                isNewGeneration: true,
-              })
-            );
-            updateFocus(params.model);
-            return;
-          }
-          // insert column desc
+        }
+        break;
+      case "docgen:insert":
+        panelLogger.info("received new doc gen", event.data);
+        // insert model desc
+        if (params.model) {
           dispatch(
-            updateColumnsInCurrentDocsData({
-              columns: [params as Partial<MetadataColumn>],
+            updateCurrentDocsData({
+              description: params.description,
+              name: params.model,
               isNewGeneration: true,
-            })
+            }),
           );
-          updateFocus((params as Partial<MetadataColumn>).name);
+          updateFocus(params.model);
+          return;
+        }
+        // insert column desc
+        dispatch(
+          updateColumnsInCurrentDocsData({
+            columns: [params as Partial<MetadataColumn>],
+            isNewGeneration: true,
+          }),
+        );
+        updateFocus((params as Partial<MetadataColumn>).name);
 
-          break;
-        default:
-          break;
-      }
-    },
-    []
-  );
+        break;
+      default:
+        break;
+    }
+  }, []);
 
   const loadGenerationsHistory = (project: string, model: string) => {
     getGenerationsInModel(project, model)
       .then((data) => {
         dispatch(setGenerationsHistory(data));
       })
       .catch((err) =>
-        panelLogger.error("error while loading generations history", err)
+        panelLogger.error("error while loading generations history", err),
       );
   };
 
@@ -197,8 +266,8 @@ const DocumentationProvider = (): JSX.Element => {
     if (userInstructions) {
       dispatch(
         updateUserInstructions(
-          JSON.parse(userInstructions) as DocsGenerateUserInstructions
-        )
+          JSON.parse(userInstructions) as DocsGenerateUserInstructions,
+        ),
       );
     }
     loadGenerationsHistory(state.project, state.currentDocsData.name);
@@ -220,9 +289,14 @@ const DocumentationProvider = (): JSX.Element => {
       state,
       dispatch,
     }),
-    [state, dispatch]
+    [state, dispatch],
   );
 
+  // hack to get latest state in onMessage
+  useEffect(() => {
+    stateRef.current = state;
+  }, [state]);
+
   if (!isComponentsApiInitialized) {
     return <div>Loading...</div>;
   }