fix: improve row selection handling and clean up code, fixes: #1504

ui/studio/grid/DataGrid.tsx

@@ -927,7 +927,7 @@ export function DataGrid(props: DataGridProps) {
       });
       const shouldEnablePreview = Boolean(
         dragDropTarget.compatibleOverId &&
-        dragDropTarget.compatibleOverId !== activeId,
+          dragDropTarget.compatibleOverId !== activeId,
       );
       setIsColumnReorderPreviewEnabled((current) =>
         current === shouldEnablePreview ? current : shouldEnablePreview,
@@ -1470,8 +1470,8 @@ export function DataGrid(props: DataGridProps) {
       const text = hasRowSelection
         ? getSelectedRowClipboardText()
         : hasCellSelection
-          ? getSelectedClipboardText()
-          : getFocusedCellClipboardText();
+        ? getSelectedClipboardText()
+        : getFocusedCellClipboardText();
 
       if (!text) {
         return;
@@ -2064,12 +2064,26 @@ export function DataGrid(props: DataGridProps) {
         return;
       }
 
-      selectSingleRowMode({ rowId, rowIndex, drag: isPrimaryButton });
+      clearNativeTextSelection();
+      clearCellSelectionState();
+
+      const nextSelection = { ...rowSelectionState };
+
+      if (nextSelection[rowId]) {
+        delete nextSelection[rowId];
+      } else {
+        nextSelection[rowId] = true;
+      }
+
+      setRowSelection(nextSelection);
+
+      rowSelectionDragRef.current = isPrimaryButton;
+      rowSelectionAnchorRef.current = rowIndex;
     },
     [
       clearCellSelectionState,
+      clearNativeTextSelection,
       rowSelectionState,
-      selectSingleRowMode,
       setRowSelection,
     ],
   );

ui/studio/grid/test-utils.tsx

@@ -2,6 +2,7 @@ import type { ColumnDef } from "@tanstack/react-table";
 import { act } from "react";
 import { type Mock, vi } from "vitest";
 
+import { CheckboxTable } from "../../components/ui/checkbox-table";
 import { TableHead } from "../../components/ui/table";
 import { Cell, type CellProps } from "../cell/Cell";
 
@@ -58,14 +59,28 @@ export function createReadOnlyColumns(args?: {
     enableSorting: false,
     size: 35,
     minSize: 35,
-    header() {
+    header({ table }) {
       return (props: Omit<CellProps, "children" | "ref">) => (
-        <TableHead {...props} aria-label="Row selection spacer" />
+        <TableHead {...props} aria-label="Row selection spacer">
+          <div className="flex items-center justify-center h-full w-full">
+            <CheckboxTable
+              checked={table.getIsAllRowsSelected()}
+              className="pointer-events-none h-4 w-4"
+            />
+          </div>
+        </TableHead>
       );
     },
-    cell() {
+    cell({ row }) {
       return (props: Omit<CellProps, "children" | "ref">) => (
-        <Cell data-select="true" {...props} />
+        <Cell data-select="true" {...props}>
+          <div className="flex items-center justify-center h-full w-full">
+            <CheckboxTable
+              checked={row.getIsSelected()}
+              className="pointer-events-none h-4 w-4"
+            />
+          </div>
+        </Cell>
       );
     },
   };

ui/studio/views/sql/SqlView.tsx

@@ -24,6 +24,7 @@ import { createSqlEditorSchemaFromIntrospection } from "../../../../data/sql-edi
 import { getTopLevelSqlStatementAtCursor } from "../../../../data/sql-statements";
 import { Button } from "../../../components/ui/button";
 import { Input } from "../../../components/ui/input";
+import { CheckboxTable } from "../../../components/ui/checkbox-table";
 import { TableHead, TableRow } from "../../../components/ui/table";
 import { useColumnPinning } from "../../../hooks/use-column-pinning";
 import { useIntrospection } from "../../../hooks/use-introspection";
@@ -103,15 +104,31 @@ const SQL_ROW_SELECTION_COLUMN_DEF = {
   size: 35,
   minSize: 35,
   header({ table }) {
-    void table;
     return (props: Omit<CellProps, "children" | "ref">) => {
-      return <TableHead {...props} aria-label="Row selection spacer" />;
+      return (
+        <TableHead {...props} aria-label="Row selection spacer">
+          <div className="flex items-center justify-center h-full w-full">
+            <CheckboxTable
+              checked={table.getIsAllRowsSelected()}
+              className="pointer-events-none h-4 w-4"
+            />
+          </div>
+        </TableHead>
+      );
     };
   },
   cell({ row }) {
-    void row;
     return (props: Omit<CellProps, "children" | "ref">) => {
-      return <Cell data-select="true" {...props} />;
+      return (
+        <Cell data-select="true" {...props}>
+          <div className="flex items-center justify-center h-full w-full">
+            <CheckboxTable
+              checked={row.getIsSelected()}
+              className="pointer-events-none h-4 w-4"
+            />
+          </div>
+        </Cell>
+      );
     };
   },
 } satisfies ColumnDef<Record<string, unknown>>;

ui/studio/views/table/ActiveTableView.tsx

@@ -39,6 +39,7 @@ import {
   DropdownMenuSeparator,
   DropdownMenuTrigger,
 } from "../../../components/ui/dropdown-menu";
+import { CheckboxTable } from "../../../components/ui/checkbox-table";
 import { TableHead } from "../../../components/ui/table";
 import { useActiveTableInsert } from "../../../hooks/use-active-table-insert";
 import { useActiveTableQuery } from "../../../hooks/use-active-table-query";
@@ -1516,15 +1517,31 @@ export function ActiveTableView(_props: ViewProps) {
       size: 35,
       minSize: 35,
       header({ table }) {
-        void table;
         return (props: Omit<CellProps, "children" | "ref">) => {
-          return <TableHead {...props} aria-label="Row selection spacer" />;
+          return (
+            <TableHead {...props} aria-label="Row selection spacer">
+              <div className="flex items-center justify-center h-full w-full">
+                <CheckboxTable
+                  checked={table.getIsAllRowsSelected()}
+                  className="pointer-events-none h-4 w-4"
+                />
+              </div>
+            </TableHead>
+          );
         };
       },
       cell({ row }) {
-        void row;
         return (props: Omit<CellProps, "children" | "ref">) => {
-          return <Cell data-select="true" {...props} />;
+          return (
+            <Cell data-select="true" {...props}>
+              <div className="flex items-center justify-center h-full w-full">
+                <CheckboxTable
+                  checked={row.getIsSelected()}
+                  className="pointer-events-none h-4 w-4"
+                />
+              </div>
+            </Cell>
+          );
         };
       },
     },