From 6f3b46ffc23ddcd56a09808d3e4533ff02c46d20 Mon Sep 17 00:00:00 2001
From: "J.C. Zhong" <exzjc@hotmail.com>
Date: Fri, 22 Nov 2024 01:30:33 +0000
Subject: [PATCH] fix: codemirror editor  tab behavior

---
 .../extensions/useKeyMapExtension.ts          |  2 -
 .../extensions/useOptionsExtension.ts         | 46 +++++++++++++++++--
 2 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/querybook/webapp/hooks/queryEditor/extensions/useKeyMapExtension.ts b/querybook/webapp/hooks/queryEditor/extensions/useKeyMapExtension.ts
index 26b24b3fa..fc9630e34 100644
--- a/querybook/webapp/hooks/queryEditor/extensions/useKeyMapExtension.ts
+++ b/querybook/webapp/hooks/queryEditor/extensions/useKeyMapExtension.ts
@@ -2,7 +2,6 @@ import { KeyBinding, keymap, Prec } from '@uiw/react-codemirror';
 import { useCallback, useMemo } from 'react';
 
 import { CodeMirrorKeyMap } from 'lib/codemirror';
-import { indentLess, insertTab } from '@codemirror/commands';
 
 export const useKeyMapExtension = ({
     keyMap = {},
@@ -41,7 +40,6 @@ export const useKeyMapExtension = ({
                     })),
                 ])
             ),
-            keymap.of([{ key: 'Tab', run: insertTab, shift: indentLess }]),
         ],
         [keyBindings, keyMap, transformKey]
     );
diff --git a/querybook/webapp/hooks/queryEditor/extensions/useOptionsExtension.ts b/querybook/webapp/hooks/queryEditor/extensions/useOptionsExtension.ts
index 0aa30fc62..d09772ddc 100644
--- a/querybook/webapp/hooks/queryEditor/extensions/useOptionsExtension.ts
+++ b/querybook/webapp/hooks/queryEditor/extensions/useOptionsExtension.ts
@@ -1,7 +1,29 @@
+import { indentLess, indentMore, insertTab } from '@codemirror/commands';
 import { indentUnit } from '@codemirror/language';
-import { EditorView } from '@uiw/react-codemirror';
+import { EditorSelection, EditorState, Extension } from '@codemirror/state';
+import { EditorView, keymap } from '@uiw/react-codemirror';
 import { useMemo } from 'react';
-import { EditorState, Extension } from '@codemirror/state';
+
+// Insert spaces instead of a real tab to reach the next tab stop
+const insertSpacesOnTab =
+    (tabSize: number) =>
+    ({ state, dispatch }) => {
+        const line = state.doc.lineAt(state.selection.main.from);
+        const column = state.selection.main.from - line.from;
+        const spacesToInsert = tabSize - (column % tabSize);
+
+        let changes = state.changeByRange((range) => ({
+            changes: {
+                from: range.from,
+                to: range.to,
+                insert: ' '.repeat(spacesToInsert),
+            },
+            range: EditorSelection.cursor(range.from + spacesToInsert),
+        }));
+
+        dispatch(state.update(changes, { userEvent: 'input' }));
+        return true;
+    };
 
 export const useOptionsExtension = ({
     lineWrapping = true,
@@ -15,9 +37,27 @@ export const useOptionsExtension = ({
 
         if (options.indentWithTabs) {
             extensions.push(indentUnit.of('\t'));
-            extensions.push(EditorState.tabSize.of(options.tabSize));
+            extensions.push(
+                keymap.of([
+                    {
+                        key: 'Tab',
+                        run: insertTab,
+                        shift: indentLess,
+                    },
+                ])
+            );
         } else {
+            extensions.push(EditorState.tabSize.of(options.tabSize));
             extensions.push(indentUnit.of(' '.repeat(options.indentUnit)));
+            extensions.push(
+                keymap.of([
+                    {
+                        key: 'Tab',
+                        run: insertSpacesOnTab(options.tabSize),
+                        shift: indentLess,
+                    },
+                ])
+            );
         }
 
         if (lineWrapping) {