feat(admin): Add undo/redo to query editors

snuba/admin/static/common/components/sql_editor.tsx

@@ -1,11 +1,20 @@
 import React, { useLayoutEffect, useMemo, useRef } from "react";
-import { EditorState } from "@codemirror/state";
+import { Annotation, EditorState } from "@codemirror/state";
 import { EditorView, lineNumbers, keymap } from "@codemirror/view";
-import { insertNewline } from "@codemirror/commands";
+import {
+  history,
+  historyKeymap,
+  insertNewline,
+  isolateHistory,
+} from "@codemirror/commands";
 
 import { sql } from "@codemirror/lang-sql";
 import { theme, highlighting } from "./theme";
 
+// Marks changes that came from the `value` prop, so they are not echoed back
+// through `onChange`.
+const fromProps = Annotation.define<boolean>();
+
 export interface Props {
   value: string;
   onChange: (newValue: string) => void;
@@ -14,14 +23,20 @@ export interface Props {
 export function SQLEditor({ value, onChange }: Props) {
   const elementRef = useRef<HTMLDivElement | null>(null);
   const viewRef = useRef<EditorView | null>(null);
+  // The editor is only created once, so read the latest callback from a ref
+  const onChangeRef = useRef(onChange);
+  onChangeRef.current = onChange;
 
   const updateListener = useMemo(() => {
     return EditorView.updateListener.of((update) => {
-      if (update.docChanged) {
-        onChange(update.state.doc.toString());
+      if (
+        update.docChanged &&
+        !update.transactions.some((tr) => tr.annotation(fromProps))
+      ) {
+        onChangeRef.current(update.state.doc.toString());
       }
     });
-  }, [onChange]);
+  }, []);
 
   const multiNewLine = useMemo(() => {
     return keymap.of([{ key: "Enter", run: insertNewline }]);
@@ -30,6 +45,8 @@ export function SQLEditor({ value, onChange }: Props) {
   const extensions = [
     updateListener,
     multiNewLine,
+    history(),
+    keymap.of(historyKeymap),
     theme,
     highlighting,
     EditorView.lineWrapping,
@@ -41,6 +58,7 @@ export function SQLEditor({ value, onChange }: Props) {
     if (!elementRef.current) return;
 
     const state = EditorState.create({
+      doc: value,
       extensions,
     });
 
@@ -51,17 +69,18 @@ export function SQLEditor({ value, onChange }: Props) {
   }, []);
 
   useLayoutEffect(() => {
-    if (!viewRef.current) return;
+    const view = viewRef.current;
+    if (!view) return;
 
-    const state = viewRef.current.state;
+    const currentValue = view.state.doc.toString();
 
-    if (value !== state.doc.toString()) {
-      viewRef.current.setState(
-        EditorState.create({
-          doc: value,
-          extensions,
-        }),
-      );
+    // Dispatch rather than replace the state, so the undo history is kept and
+    // the new value (e.g. a predefined query) is its own undo step
+    if (value !== currentValue) {
+      view.dispatch({
+        changes: { from: 0, to: currentValue.length, insert: value },
+        annotations: [fromProps.of(true), isolateHistory.of("full")],
+      });
     }
   }, [value]);
 

snuba/admin/static/query_editor.tsx

@@ -123,6 +123,9 @@ function QueryEditor(props: {
   const [queryTemplate, setQueryTemplate] = useLocalStorage<string>({
     key: `${hash}-query-editor-query`,
     defaultValue: "",
+    // Read the saved query on the first render rather than in an effect, so
+    // the editor starts with it instead of receiving it as an undoable change
+    getInitialValueInEffect: false,
   });
 
   const [queryParamValues, setQueryParamValues] = useState<QueryParamValues>(

snuba/admin/static/sql_shell/shell_input.tsx

@@ -1,6 +1,7 @@
 import React, { useLayoutEffect, useMemo, useRef, useEffect } from "react";
 import { EditorState, Compartment } from "@codemirror/state";
 import { EditorView, keymap, placeholder as cmPlaceholder } from "@codemirror/view";
+import { history, historyKeymap, isolateHistory } from "@codemirror/commands";
 import { sql } from "@codemirror/lang-sql";
 import { HighlightStyle, syntaxHighlighting } from "@codemirror/language";
 import { tags } from "@lezer/highlight";
@@ -257,6 +258,8 @@ export function ShellInput({
 
   const extensions = useMemo(() => [
     customKeymap,
+    history(),
+    keymap.of(historyKeymap),
     updateListener,
     shellTheme,
     shellHighlighting,
@@ -303,6 +306,7 @@ export function ShellInput({
       view.dispatch({
         changes: { from: 0, to: currentValue.length, insert: value },
         selection: { anchor: value.length },
+        annotations: isolateHistory.of("full"),
       });
     }
   }, [value]);

snuba/admin/static/tests/query_editor_undo.spec.tsx

@@ -0,0 +1,33 @@
+import React from "react";
+import { it, expect, describe, jest, afterEach, beforeEach } from "@jest/globals";
+import { act, cleanup, fireEvent, render } from "@testing-library/react";
+import { EditorView } from "@codemirror/view";
+import QueryEditor from "SnubaAdmin/query_editor";
+
+// Unlike query_editor.spec.tsx, this uses the real CodeMirror SQLEditor
+
+// QueryEditor namespaces its storage by the URL hash, which is empty in jsdom
+const STORAGE_KEY = "-query-editor-query";
+
+describe("QueryEditor undo", () => {
+  global.ResizeObserver = require("resize-observer-polyfill");
+  beforeEach(() => {
+    localStorage.setItem(STORAGE_KEY, JSON.stringify("SELECT saved"));
+  });
+  afterEach(cleanup);
+
+  it("does not undo the saved query loaded on mount", async () => {
+    const { container } = render(<QueryEditor onQueryUpdate={jest.fn()} />);
+    const view = EditorView.findFromDOM(
+      container.querySelector(".cm-editor") as HTMLElement
+    )!;
+    expect(view.state.doc.toString()).toBe("SELECT saved");
+
+    await act(async () => {
+      fireEvent.keyDown(view.contentDOM, { key: "z", ctrlKey: true });
+    });
+
+    expect(view.state.doc.toString()).toBe("SELECT saved");
+    expect(localStorage.getItem(STORAGE_KEY)).toBe(JSON.stringify("SELECT saved"));
+  });
+});

snuba/admin/static/tests/sql_editor.spec.tsx

@@ -0,0 +1,72 @@
+import React from "react";
+import { it, expect, describe, jest, afterEach } from "@jest/globals";
+import { cleanup, fireEvent, render } from "@testing-library/react";
+import { EditorView } from "@codemirror/view";
+import { SQLEditor } from "SnubaAdmin/common/components/sql_editor";
+
+function getView(container: HTMLElement): EditorView {
+  return EditorView.findFromDOM(
+    container.querySelector(".cm-editor") as HTMLElement
+  )!;
+}
+
+function type(view: EditorView, text: string) {
+  const end = view.state.doc.length;
+  view.dispatch({ changes: { from: end, insert: text }, userEvent: "input.type" });
+}
+
+// jsdom is not a Mac, so CodeMirror's "Mod" key is Ctrl
+function pressUndo(view: EditorView) {
+  fireEvent.keyDown(view.contentDOM, { key: "z", ctrlKey: true });
+}
+
+function pressRedo(view: EditorView) {
+  fireEvent.keyDown(view.contentDOM, { key: "y", ctrlKey: true });
+}
+
+describe("SQLEditor", () => {
+  afterEach(cleanup);
+
+  it("undoes and redoes typed text with the keyboard", () => {
+    const onChange = jest.fn<(value: string) => void>();
+    const { container } = render(<SQLEditor value="" onChange={onChange} />);
+    const view = getView(container);
+
+    type(view, "SELECT 1");
+    expect(onChange).toHaveBeenLastCalledWith("SELECT 1");
+
+    pressUndo(view);
+    expect(view.state.doc.toString()).toBe("");
+    expect(onChange).toHaveBeenLastCalledWith("");
+
+    pressRedo(view);
+    expect(view.state.doc.toString()).toBe("SELECT 1");
+    expect(onChange).toHaveBeenLastCalledWith("SELECT 1");
+  });
+
+  it("can undo a value replaced from props", () => {
+    const onChange = jest.fn<(value: string) => void>();
+    const { container, rerender } = render(
+      <SQLEditor value="" onChange={onChange} />
+    );
+    const view = getView(container);
+
+    type(view, "SELECT 1");
+    rerender(<SQLEditor value="SELECT 1" onChange={onChange} />);
+    // e.g. the user picks a predefined query
+    rerender(<SQLEditor value="SELECT 2" onChange={onChange} />);
+    expect(view.state.doc.toString()).toBe("SELECT 2");
+
+    pressUndo(view);
+    expect(view.state.doc.toString()).toBe("SELECT 1");
+    expect(onChange).toHaveBeenLastCalledWith("SELECT 1");
+  });
+
+  it("does not report a value set from props as a change", () => {
+    const onChange = jest.fn<(value: string) => void>();
+    const { rerender } = render(<SQLEditor value="" onChange={onChange} />);
+
+    rerender(<SQLEditor value="SELECT 2" onChange={onChange} />);
+    expect(onChange).not.toHaveBeenCalled();
+  });
+});

snuba/admin/static/tests/sql_shell/shell_input.spec.tsx

@@ -0,0 +1,47 @@
+import React from "react";
+import { it, expect, describe, jest, afterEach } from "@jest/globals";
+import { cleanup, fireEvent, render } from "@testing-library/react";
+import { EditorView } from "@codemirror/view";
+import { ShellInput, ShellInputProps } from "SnubaAdmin/sql_shell/shell_input";
+
+function renderShellInput(overrides: Partial<ShellInputProps> = {}) {
+  const props: ShellInputProps = {
+    value: "",
+    onChange: jest.fn(),
+    onExecute: jest.fn(),
+    onTab: jest.fn(() => false),
+    onHistoryUp: jest.fn(),
+    onHistoryDown: jest.fn(),
+    onClear: jest.fn(),
+    onClearScreen: jest.fn(),
+    onDeleteToStart: jest.fn(),
+    onDeleteToEnd: jest.fn(),
+    onDeleteWord: jest.fn(),
+    ...overrides,
+  };
+  const { container } = render(<ShellInput {...props} />);
+  const view = EditorView.findFromDOM(
+    container.querySelector(".cm-editor") as HTMLElement
+  )!;
+  return { props, view };
+}
+
+describe("ShellInput", () => {
+  afterEach(cleanup);
+
+  it("undoes and redoes typed text with the keyboard", () => {
+    const onChange = jest.fn<(value: string) => void>();
+    const { view } = renderShellInput({ onChange });
+
+    view.dispatch({ changes: { from: 0, insert: "SELECT 1" }, userEvent: "input.type" });
+    expect(onChange).toHaveBeenLastCalledWith("SELECT 1");
+
+    // jsdom is not a Mac, so CodeMirror's "Mod" key is Ctrl
+    fireEvent.keyDown(view.contentDOM, { key: "z", ctrlKey: true });
+    expect(view.state.doc.toString()).toBe("");
+    expect(onChange).toHaveBeenLastCalledWith("");
+
+    fireEvent.keyDown(view.contentDOM, { key: "y", ctrlKey: true });
+    expect(view.state.doc.toString()).toBe("SELECT 1");
+  });
+});

snuba/admin/yarn.lock

@@ -474,7 +474,7 @@
   resolved "https://registry.yarnpkg.com/@codemirror/state/-/state-6.4.1.tgz#da57143695c056d9a3c38705ed34136e2b68171b"
   integrity sha512-QkEyUiLhsJoZkbumGZlswmAhA7CBU02Wrz7zvH4SrcifbsqwlXShVXg65f3v/ts57W3dqyamEriMhij1Z3Zz4A==
 
-"@codemirror/view@^6.17.0", "@codemirror/view@^6.23.0", "@codemirror/view@^6.32.0":
+"@codemirror/view@^6.17.0", "@codemirror/view@^6.23.0", "@codemirror/view@^6.27.0", "@codemirror/view@^6.32.0":
   version "6.32.0"
   resolved "https://registry.yarnpkg.com/@codemirror/view/-/view-6.32.0.tgz#94be8aa18b60044471200ee80b0339513036680e"
   integrity sha512-AgVNvED2QTsZp5e3syoHLsrWtwJFYWdx1Vr/m3f4h1ATQz0ax60CfXF3Htdmk69k2MlYZw8gXesnQdHtzyVmAw==
@@ -483,15 +483,6 @@
     style-mod "^4.1.0"
     w3c-keyname "^2.2.4"
 
-"@codemirror/view@^6.27.0":
-  version "6.33.0"
-  resolved "https://registry.yarnpkg.com/@codemirror/view/-/view-6.33.0.tgz#51e270410fc3af92a6e38798e80ebf8add7dc3ec"
-  integrity sha512-AroaR3BvnjRW8fiZBalAaK+ZzB5usGgI014YKElYZvQdNH5ZIidHlO+cyf/2rWzyBFRkvG6VhiXeAEbC53P2YQ==
-  dependencies:
-    "@codemirror/state" "^6.4.0"
-    style-mod "^4.1.0"
-    w3c-keyname "^2.2.4"
-
 "@emotion/babel-plugin@^11.12.0":
   version "11.12.0"
   resolved "https://registry.yarnpkg.com/@emotion/babel-plugin/-/babel-plugin-11.12.0.tgz#7b43debb250c313101b3f885eba634f1d723fcc2"