Select all column / row cells on menu trigger click

Signed-off-by: Émile Ré <emile@getprobo.com>

packages/ui/src/RichEditor/TableColumnMenu.tsx

@@ -423,6 +423,24 @@ export function TableColumnMenu({ editor }: TableColumnMenuProps) {
           // table may have changed
         }
       } else {
+        try {
+          const table = editor.state.doc.nodeAt(tableStart - 1);
+          if (table) {
+            const map = TableMap.get(table);
+            const anchorPos
+              = map.positionAt(0, fromCol, table) + tableStart;
+            const headPos
+              = map.positionAt(map.height - 1, fromCol, table) + tableStart;
+            const sel = CellSelection.create(
+              view.state.doc,
+              anchorPos,
+              headPos,
+            );
+            view.dispatch(view.state.tr.setSelection(sel));
+          }
+        } catch {
+          // table may have changed
+        }
         setMenuOpen(prev => !prev);
       }
 

packages/ui/src/RichEditor/TableRowMenu.tsx

@@ -419,6 +419,24 @@ export function TableRowMenu({ editor }: TableRowMenuProps) {
           // table may have changed
         }
       } else {
+        try {
+          const table = editor.state.doc.nodeAt(tableStart - 1);
+          if (table) {
+            const map = TableMap.get(table);
+            const anchorPos
+              = map.positionAt(fromRow, 0, table) + tableStart;
+            const headPos
+              = map.positionAt(fromRow, map.width - 1, table) + tableStart;
+            const sel = CellSelection.create(
+              view.state.doc,
+              anchorPos,
+              headPos,
+            );
+            view.dispatch(view.state.tr.setSelection(sel));
+          }
+        } catch {
+          // table may have changed
+        }
         setMenuOpen(prev => !prev);
       }