persist json value in document

packages/compass-crud/src/components/json-editor.tsx

@@ -76,10 +76,22 @@ const JSONEditor: React.FunctionComponent<JSONEditorProps> = ({
   const [expanded, setExpanded] = useState<boolean>(doc.expanded);
   const [editing, setEditing] = useState<boolean>(doc.editing);
   const [deleting, setDeleting] = useState<boolean>(doc.markedForDeletion);
-  const [value, setValue] = useState<string>(() => doc.toEJSON());
+  const [value, setValue] = useState<string>(
+    () => doc.modifiedEJSONString ?? doc.toEJSON()
+  );
   const [initialValue] = useState<string>(() => doc.toEJSON());
   const [containsErrors, setContainsErrors] = useState<boolean>(false);
 
+  useEffect(() => {
+    return () => {
+      // When this component is used in virtualized list, the editor is
+      // unmounted on scroll and if the user is editing the document, the
+      // editor value is lost. This is a way to keep track of the editor
+      // value when the it's unmounted and is restored on next mount.
+      doc.setModifiedEJSONString(editing ? value : null);
+    };
+  }, [value, editing, doc]);
+
   const handleCopy = useCallback(() => {
     copyToClipboard?.(doc);
   }, [copyToClipboard, doc]);

packages/compass-crud/src/components/virtualized-document-json-view.spec.tsx

@@ -12,6 +12,10 @@ import {
 import { type VirtualListRef } from '@mongodb-js/compass-components';
 
 import VirtualizedDocumentJsonView from './virtualized-document-json-view';
+import {
+  getCodemirrorEditorValue,
+  setCodemirrorEditorValue,
+} from '@mongodb-js/compass-editor';
 
 const createBigDocument = (variable: number) =>
   new HadronDocument({
@@ -168,4 +172,44 @@ describe('VirtualizedDocumentJsonView', function () {
     expect(() => within(documentElement).getByText('Cancel')).to.throw;
     expect(() => within(documentElement).getByText('Replace')).to.throw;
   });
+
+  it('preserves the edit state of document when a document goes out of visible viewport when scrolling', async function () {
+    const bigDocuments = Array.from({ length: 20 }, (_, idx) =>
+      createBigDocument(idx)
+    );
+    const listRef: VirtualListRef = React.createRef();
+    render(
+      <VirtualizedDocumentJsonView
+        namespace="x.y"
+        docs={bigDocuments}
+        isEditable={true}
+        __TEST_OVERSCAN_COUNT={0}
+        __TEST_LIST_HEIGHT={178}
+        __TEST_LIST_REF={listRef}
+      />
+    );
+
+    let [firstDocumentElement] = screen.getAllByTestId('editable-json');
+    // Trigger the edit state (we only set the edited value if the document is being edited)
+    userEvent.click(within(firstDocumentElement).getByLabelText('Edit'));
+
+    let cmEditor = firstDocumentElement.querySelector(
+      '[data-codemirror="true"]'
+    );
+    await setCodemirrorEditorValue(cmEditor, '{value: "edited"}');
+
+    // Scroll down and then scroll back up
+    act(() => {
+      listRef.current?.scrollToItem(15);
+    });
+    act(() => {
+      listRef.current?.scrollToItem(0);
+    });
+
+    [firstDocumentElement] = screen.getAllByTestId('editable-json');
+    cmEditor = firstDocumentElement.querySelector('[data-codemirror="true"]');
+
+    const value = getCodemirrorEditorValue(cmEditor);
+    expect(value).to.equal('{value: "edited"}');
+  });
 });

packages/hadron-document/src/document.ts

@@ -50,6 +50,9 @@ export class Document extends EventEmitter {
   maxVisibleElementsCount = DEFAULT_VISIBLE_ELEMENTS;
   editing = false;
   markedForDeletion = false;
+  // This is used to store the changed EJSON string when the document is modified
+  // via the JSONEditor.
+  modifiedEJSONString: string | null = null;
 
   /**
    * Send cancel event.
@@ -456,6 +459,7 @@ export class Document extends EventEmitter {
   finishEditing() {
     if (this.editing) {
       this.editing = false;
+      this.setModifiedEJSONString(null);
       this.emit(DocumentEvents.EditingFinished);
     }
   }
@@ -503,6 +507,10 @@ export class Document extends EventEmitter {
   onRemoveError(error: Error) {
     this.emit('remove-error', error.message);
   }
+
+  setModifiedEJSONString(ejson: string | null) {
+    this.modifiedEJSONString = ejson;
+  }
 }
 
 export default Document;